oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merged PR 348: #2462 - Refactor of DatabaseClient and DatabaseServer

Browse Source 
## Summary
Refactor of `DatabaseClient` and `DatabaseService` to update how connection IDs are generated. These are now provided by DatabaseService when establishing a connection.
Creation of `DatabaseClientConnection` class. This is used by `DatabaseClient` to hold a dictionary of active db connections.

## Test process
Tests have been updated to reflect the changes and all pass

## Checklist
- [X] PR is linked to a **work item**
- [X] **acceptance criteria** of linked ticket are met
- [X] performed **self-review** of the code
- [X] written **tests** for any new functionality added with this PR
- [X] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [X] updated the **change log**
- [X] ran **pre-commit** checks for code style
- [X] attended to any **TO-DOs** left in the code

Related work items: #2462
This commit is contained in:
Charlie Crane
2024-04-26 14:52:21 +00:00
parent e1ac6255ad
commit 5ee23dcb17
21 changed files with 502 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
tests/integration_tests/network/test_capture_nmne.py
View File

@@ -2,7 +2,7 @@ from primaite.game.agent.observations.nic_observations import NICObservation
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.nmne import set_nmne_config
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
def test_capture_nmne(uc2_network):
@@ -15,7 +15,7 @@ def test_capture_nmne(uc2_network):
"""
web_server: Server = uc2_network.get_node_by_hostname("web_server") # noqa
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] # noqa
db_client.connect()
db_client_connection: DatabaseClientConnection = db_client.get_new_connection()
db_server: Server = uc2_network.get_node_by_hostname("database_server") # noqa
@@ -39,42 +39,42 @@ def test_capture_nmne(uc2_network):
assert db_server_nic.nmne == {}
# Perform a "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that it does not trigger an MNE capture.
assert web_server_nic.nmne == {}
assert db_server_nic.nmne == {}
# Perform a "DELETE" query
db_client.query("DELETE")
db_client_connection.query(sql="DELETE")
# Check that the web server's outbound interface and the database server's inbound interface register the MNE
assert web_server_nic.nmne == {"direction": {"outbound": {"keywords": {"*": 1}}}}
assert db_server_nic.nmne == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that no additional MNEs are captured
assert web_server_nic.nmne == {"direction": {"outbound": {"keywords": {"*": 1}}}}
assert db_server_nic.nmne == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "DELETE" query
db_client.query("DELETE")
db_client_connection.query(sql="DELETE")
# Check that the web server and database server interfaces register an additional MNE
assert web_server_nic.nmne == {"direction": {"outbound": {"keywords": {"*": 2}}}}
assert db_server_nic.nmne == {"direction": {"inbound": {"keywords": {"*": 2}}}}
# Perform an "ENCRYPT" query
db_client.query("ENCRYPT")
db_client_connection.query(sql="ENCRYPT")
# Check that the web server and database server interfaces register an additional MNE
assert web_server_nic.nmne == {"direction": {"outbound": {"keywords": {"*": 3}}}}
assert db_server_nic.nmne == {"direction": {"inbound": {"keywords": {"*": 3}}}}
# Perform another "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that no additional MNEs are captured
assert web_server_nic.nmne == {"direction": {"outbound": {"keywords": {"*": 3}}}}
@@ -92,7 +92,7 @@ def test_describe_state_nmne(uc2_network):
"""
web_server: Server = uc2_network.get_node_by_hostname("web_server") # noqa
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"] # noqa
db_client.connect()
db_client_connection: DatabaseClientConnection = db_client.get_new_connection()
db_server: Server = uc2_network.get_node_by_hostname("database_server") # noqa
@@ -119,7 +119,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {}
# Perform a "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that it does not trigger an MNE capture.
web_server_nic_state = web_server_nic.describe_state()
@@ -129,7 +129,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {}
# Perform a "DELETE" query
db_client.query("DELETE")
db_client_connection.query(sql="DELETE")
# Check that the web server's outbound interface and the database server's inbound interface register the MNE
web_server_nic_state = web_server_nic.describe_state()
@@ -139,7 +139,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that no additional MNEs are captured
web_server_nic_state = web_server_nic.describe_state()
@@ -149,7 +149,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {"direction": {"inbound": {"keywords": {"*": 1}}}}
# Perform another "DELETE" query
db_client.query("DELETE")
db_client_connection.query(sql="DELETE")
# Check that the web server and database server interfaces register an additional MNE
web_server_nic_state = web_server_nic.describe_state()
@@ -159,7 +159,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {"direction": {"inbound": {"keywords": {"*": 2}}}}
# Perform a "ENCRYPT" query
db_client.query("ENCRYPT")
db_client_connection.query(sql="ENCRYPT")
# Check that the web server's outbound interface and the database server's inbound interface register the MNE
web_server_nic_state = web_server_nic.describe_state()
@@ -169,7 +169,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {"direction": {"inbound": {"keywords": {"*": 3}}}}
# Perform another "SELECT" query
db_client.query("SELECT")
db_client_connection.query(sql="SELECT")
# Check that no additional MNEs are captured
web_server_nic_state = web_server_nic.describe_state()
@@ -179,7 +179,7 @@ def test_describe_state_nmne(uc2_network):
assert db_server_nic_state["nmne"] == {"direction": {"inbound": {"keywords": {"*": 3}}}}
# Perform another "ENCRYPT"
db_client.query("ENCRYPT")
db_client_connection.query(sql="ENCRYPT")
# Check that the web server and database server interfaces register an additional MNE
web_server_nic_state = web_server_nic.describe_state()
@@ -206,7 +206,7 @@ def test_capture_nmne_observations(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect()
db_client_connection: DatabaseClientConnection = db_client.get_new_connection()
# Set the NMNE configuration to capture DELETE/ENCRYPT queries as MNEs
nmne_config = {
@@ -228,7 +228,7 @@ def test_capture_nmne_observations(uc2_network):
for i in range(0, 20):
# Perform a "DELETE" query each iteration
for j in range(i):
db_client.query("DELETE")
db_client_connection.query(sql="DELETE")
# Observe the current state of NMNEs from the NICs of both the database and web servers
state = sim.describe_state()
@@ -253,7 +253,7 @@ def test_capture_nmne_observations(uc2_network):
for i in range(0, 20):
# Perform a "ENCRYPT" query each iteration
for j in range(i):
db_client.query("ENCRYPT")
db_client_connection.query(sql="ENCRYPT")
# Observe the current state of NMNEs from the NICs of both the database and web servers
state = sim.describe_state()

8
tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py
View File

@@ -10,7 +10,7 @@ from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.simulator.system.applications.red_applications.data_manipulation_bot import (
DataManipulationAttackStage,
DataManipulationBot,
@@ -141,8 +141,10 @@ def test_data_manipulation_disrupts_green_agent_connection(data_manipulation_db_
server: Server = network.get_node_by_hostname("server_1")
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
green_db_connection: DatabaseClientConnection = green_db_client.get_new_connection()
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD
assert green_db_client.query("SELECT")
assert green_db_connection.query("SELECT")
assert green_db_client.last_query_response.get("status_code") == 200
data_manipulation_bot.port_scan_p_of_success = 1
@@ -151,5 +153,5 @@ def test_data_manipulation_disrupts_green_agent_connection(data_manipulation_db_
data_manipulation_bot.attack()
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED
assert green_db_client.query("SELECT") is False
assert green_db_connection.query("SELECT") is False
assert green_db_client.last_query_response.get("status_code") != 200

7
tests/integration_tests/system/red_applications/test_ransomware_script.py
View File

@@ -10,7 +10,7 @@ from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.simulator.system.applications.red_applications.ransomware_script import (
RansomwareAttackStage,
RansomwareScript,
@@ -144,12 +144,13 @@ def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_
client_2: Computer = network.get_node_by_hostname("client_2")
green_db_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient")
green_db_client_connection: DatabaseClientConnection = green_db_client.get_new_connection()
server: Server = network.get_node_by_hostname("server_1")
db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService")
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD
assert green_db_client.query("SELECT")
assert green_db_client_connection.query("SELECT")
assert green_db_client.last_query_response.get("status_code") == 200
ransomware_script_application.target_scan_p_of_success = 1
@@ -159,5 +160,5 @@ def test_ransomware_disrupts_green_agent_connection(ransomware_script_db_server_
ransomware_script_application.attack()
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
assert green_db_client.query("SELECT") is True
assert green_db_client_connection.query("SELECT") is True
assert green_db_client.last_query_response.get("status_code") == 200

158
tests/integration_tests/system/test_database_on_node.py
View File

@@ -8,7 +8,8 @@ from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.system.applications.database_client import DatabaseClient, DatabaseClientConnection
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
@@ -56,11 +57,12 @@ def test_database_client_server_connection(peer_to_peer):
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"]
db_client.connect()
assert len(db_client.connections) == 1
assert len(db_client.client_connections) == 1
assert len(db_service.connections) == 1
db_client.disconnect()
assert len(db_client.connections) == 0
assert len(db_client.client_connections) == 0
assert len(db_service.connections) == 0
@@ -73,7 +75,7 @@ def test_database_client_server_correct_password(peer_to_peer_secure_db):
db_client.configure(server_ip_address=IPv4Address("192.168.0.11"), server_password="12345")
db_client.connect()
assert len(db_client.connections) == 1
assert len(db_client.client_connections) == 1
assert len(db_service.connections) == 1
@@ -95,14 +97,24 @@ def test_database_client_server_incorrect_password(peer_to_peer_secure_db):
assert len(db_service.connections) == 0
def test_database_client_query(uc2_network):
def test_database_client_native_connection_query(uc2_network):
"""Tests DB query across the network returns HTTP status 200 and date."""
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_client.connect()
assert db_client.query("SELECT")
assert db_client.query("INSERT")
assert db_client.query(sql="SELECT")
assert db_client.query(sql="INSERT")
def test_database_client_connection_query(uc2_network):
"""Tests DB query across the network returns HTTP status 200 and date."""
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
db_connection: DatabaseClientConnection = db_client.get_new_connection()
assert db_connection.query(sql="SELECT")
assert db_connection.query(sql="INSERT")
def test_create_database_backup(uc2_network):
@@ -172,7 +184,6 @@ def test_restore_backup_after_deleting_file_without_updating_scan(uc2_network):
db_server: Server = uc2_network.get_node_by_hostname("database_server")
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
# create a back up
assert db_service.backup_database() is True
db_service.db_file.corrupt() # corrupt the db
@@ -211,10 +222,13 @@ def test_database_client_cannot_query_offline_database_server(uc2_network):
web_server: Server = uc2_network.get_node_by_hostname("web_server")
db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
assert len(db_client.connections)
assert len(db_client.client_connections)
assert db_client.query("SELECT") is True
assert db_client.query("INSERT") is True
# Establish a new connection to the DatabaseService
db_connection: DatabaseClientConnection = db_client.get_new_connection()
assert db_connection.query("SELECT") is True
assert db_connection.query("INSERT") is True
db_server.power_off()
for i in range(db_server.shut_down_duration + 1):
@@ -223,5 +237,121 @@ def test_database_client_cannot_query_offline_database_server(uc2_network):
assert db_server.operating_state is NodeOperatingState.OFF
assert db_service.operating_state is ServiceOperatingState.STOPPED
assert db_client.query("SELECT") is False
assert db_client.query("INSERT") is False
assert db_connection.query("SELECT") is False
assert db_connection.query("INSERT") is False
def test_database_client_uninstall_terminates_connections(peer_to_peer):
node_a, node_b = peer_to_peer
db_client: DatabaseClient = node_a.software_manager.software["DatabaseClient"]
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"] # noqa
db_connection: DatabaseClientConnection = db_client.get_new_connection()
# Check that all connection counters are correct and that the client connection can query the database
assert len(db_service.connections) == 1
assert len(db_client.client_connections) == 1
assert db_connection.is_active
assert db_connection.query("SELECT")
# Perform the DatabaseClient uninstall
node_a.software_manager.uninstall("DatabaseClient")
# Check that all connection counters are updated accordingly and client connection can no longer query the database
assert len(db_service.connections) == 0
assert len(db_client.client_connections) == 0
assert not db_connection.query("SELECT")
assert not db_connection.is_active
def test_database_service_can_terminate_connection(peer_to_peer):
node_a, node_b = peer_to_peer
db_client: DatabaseClient = node_a.software_manager.software["DatabaseClient"]
db_service: DatabaseService = node_b.software_manager.software["DatabaseService"] # noqa
db_connection: DatabaseClientConnection = db_client.get_new_connection()
# Check that all connection counters are correct and that the client connection can query the database
assert len(db_service.connections) == 1
assert len(db_client.client_connections) == 1
assert db_connection.is_active
assert db_connection.query("SELECT")
# Perform the server-led connection termination
connection_id = next(iter(db_service.connections.keys()))
db_service.terminate_connection(connection_id)
# Check that all connection counters are updated accordingly and client connection can no longer query the database
assert len(db_service.connections) == 0
assert len(db_client.client_connections) == 0
assert not db_connection.query("SELECT")
assert not db_connection.is_active
def test_client_connection_terminate_does_not_terminate_another_clients_connection():
network = Network()
db_server = Server(
hostname="db_client", ip_address="192.168.0.11", subnet_mask="255.255.255.0", start_up_duration=0
)
db_server.power_on()
db_server.software_manager.install(DatabaseService)
db_service: DatabaseService = db_server.software_manager.software["DatabaseService"] # noqa
db_service.start()
client_a = Computer(
hostname="client_a", ip_address="192.168.0.12", subnet_mask="255.255.255.0", start_up_duration=0
)
client_a.power_on()
client_a.software_manager.install(DatabaseClient)
client_a.software_manager.software["DatabaseClient"].configure(server_ip_address=IPv4Address("192.168.0.11"))
client_a.software_manager.software["DatabaseClient"].run()
client_b = Computer(
hostname="client_b", ip_address="192.168.0.13", subnet_mask="255.255.255.0", start_up_duration=0
)
client_b.power_on()
client_b.software_manager.install(DatabaseClient)
client_b.software_manager.software["DatabaseClient"].configure(server_ip_address=IPv4Address("192.168.0.11"))
client_b.software_manager.software["DatabaseClient"].run()
switch = Switch(hostname="switch", start_up_duration=0, num_ports=3)
switch.power_on()
network.connect(endpoint_a=switch.network_interface[1], endpoint_b=db_server.network_interface[1])
network.connect(endpoint_a=switch.network_interface[2], endpoint_b=client_a.network_interface[1])
network.connect(endpoint_a=switch.network_interface[3], endpoint_b=client_b.network_interface[1])
db_client_a: DatabaseClient = client_a.software_manager.software["DatabaseClient"] # noqa
db_connection_a = db_client_a.get_new_connection()
assert db_connection_a.query("SELECT")
assert len(db_service.connections) == 1
db_client_b: DatabaseClient = client_b.software_manager.software["DatabaseClient"] # noqa
db_connection_b = db_client_b.get_new_connection()
assert db_connection_b.query("SELECT")
assert len(db_service.connections) == 2
db_connection_a.disconnect()
assert db_connection_b.query("SELECT")
assert len(db_service.connections) == 1

5
tests/integration_tests/test_simulation/test_request_response.py
View File

@@ -7,6 +7,7 @@
import pytest
from primaite.interface.request import RequestResponse
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
@@ -97,7 +98,7 @@ def test_request_fails_if_node_off(example_network, node_request):
class TestDataManipulationGreenRequests:
def test_node_off(self, uc2_network):
"""Test that green requests succeed when the node is on and fail if the node is off."""
net = uc2_network
net: Network = uc2_network
client_1_browser_execute = net.apply_request(["node", "client_1", "application", "WebBrowser", "execute"])
client_1_db_client_execute = net.apply_request(["node", "client_1", "application", "DatabaseClient", "execute"])
@@ -131,7 +132,7 @@ class TestDataManipulationGreenRequests:
def test_acl_block(self, uc2_network):
"""Test that green requests succeed when not blocked by ACLs but fail when blocked."""
net = uc2_network
net: Network = uc2_network
router: Router = net.get_node_by_hostname("router_1")
client_1: HostNode = net.get_node_by_hostname("client_1")