From 383cf051df249b7fd3fcfece4ca38e0784ef7e72 Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Thu, 4 Apr 2024 14:17:34 +0100 Subject: [PATCH] #2448: store last query response for db client --- .../system/applications/database_client.py | 6 + .../services/database/database_service.py | 6 +- .../test_data_manipulation_bot_and_server.py | 155 ++++++++++++++++++ 3 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index d3afef59..1de75dc5 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -29,6 +29,8 @@ class DatabaseClient(Application): _query_success_tracker: Dict[str, bool] = {} _last_connection_successful: Optional[bool] = None """Keep track of connections that were established or verified during this step. Used for rewards.""" + last_query_response: Optional[Dict] = None + """Keep track of the latest query response. Used to determine rewards.""" def __init__(self, **kwargs): kwargs["name"] = "DatabaseClient" @@ -219,6 +221,9 @@ class DatabaseClient(Application): if not self._can_perform_action(): return False + # reset last query response + self.last_query_response = None + if connection_id is None: if self.connections: connection_id = list(self.connections.keys())[-1] @@ -252,6 +257,7 @@ class DatabaseClient(Application): # add connection self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id) elif payload["type"] == "sql": + self.last_query_response = payload query_id = payload.get("uuid") status_code = payload.get("status_code") self._query_success_tracker[query_id] = status_code == 200 diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index ede2a54f..321d9088 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -204,7 +204,7 @@ class DatabaseService(Service): if not self.db_file: self.sys_log.info(f"{self.name}: Failed to run {query} because the database file is missing.") - return {"status_code": 404, "data": False} + return {"status_code": 404, "type": "sql", "data": False} if query == "SELECT": if self.db_file.health_status == FileSystemItemHealthStatus.GOOD: @@ -216,7 +216,7 @@ class DatabaseService(Service): "connection_id": connection_id, } else: - return {"status_code": 404, "data": False} + return {"status_code": 404, "type": "sql", "data": False} elif query == "DELETE": self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED return { @@ -236,7 +236,7 @@ class DatabaseService(Service): "connection_id": connection_id, } else: - return {"status_code": 404, "data": False} + return {"status_code": 404, "type": "sql", "data": False} elif query == "SELECT * FROM pg_stat_activity": # Check if the connection is active. if self.health_state_actual == SoftwareHealthState.GOOD: diff --git a/tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py b/tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py new file mode 100644 index 00000000..1106d6ca --- /dev/null +++ b/tests/integration_tests/system/red_applications/test_data_manipulation_bot_and_server.py @@ -0,0 +1,155 @@ +from ipaddress import IPv4Address +from typing import Tuple + +import pytest + +from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus +from primaite.simulator.network.container import Network +from primaite.simulator.network.hardware.nodes.host.computer import Computer +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.red_applications.data_manipulation_bot import ( + DataManipulationAttackStage, + DataManipulationBot, +) +from primaite.simulator.system.applications.red_applications.dos_bot import DoSAttackStage, DoSBot +from primaite.simulator.system.services.database.database_service import DatabaseService +from primaite.simulator.system.software import SoftwareHealthState + + +@pytest.fixture(scope="function") +def data_manipulation_bot_and_db_server(client_server) -> Tuple[DataManipulationBot, Computer, DatabaseService, Server]: + computer, server = client_server + + # install db client on computer + computer.software_manager.install(DatabaseClient) + db_client: DatabaseClient = computer.software_manager.software.get("DatabaseClient") + db_client.run() + + # Install DoSBot on computer + computer.software_manager.install(DataManipulationBot) + + data_manipulation_bot: DataManipulationBot = computer.software_manager.software.get("DataManipulationBot") + data_manipulation_bot.configure( + server_ip_address=IPv4Address(server.network_interface[1].ip_address), payload="DELETE" + ) + + # Install DB Server service on server + server.software_manager.install(DatabaseService) + db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService") + db_server_service.start() + + return data_manipulation_bot, computer, db_server_service, server + + +@pytest.fixture(scope="function") +def data_manipulation_db_server_green_client(example_network) -> Network: + network: Network = example_network + + router_1: Router = example_network.get_node_by_hostname("router_1") + router_1.acl.add_rule( + action=ACLAction.PERMIT, src_port=Port.POSTGRES_SERVER, dst_port=Port.POSTGRES_SERVER, position=0 + ) + + client_1: Computer = network.get_node_by_hostname("client_1") + client_2: Computer = network.get_node_by_hostname("client_2") + server: Server = network.get_node_by_hostname("server_1") + + # install db client on client 1 + client_1.software_manager.install(DatabaseClient) + db_client: DatabaseClient = client_1.software_manager.software.get("DatabaseClient") + db_client.run() + + # install Data Manipulation bot on client 1 + client_1.software_manager.install(DataManipulationBot) + + data_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot") + data_manipulation_bot.configure( + server_ip_address=IPv4Address(server.network_interface[1].ip_address), payload="DELETE" + ) + + # install db server service on server + server.software_manager.install(DatabaseService) + db_server_service: DatabaseService = server.software_manager.software.get("DatabaseService") + db_server_service.start() + + # Install DB client (green) on client 2 + client_2.software_manager.install(DatabaseClient) + + database_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient") + database_client.configure(server_ip_address=IPv4Address(server.network_interface[1].ip_address)) + database_client.run() + + return network + + +def test_repeating_data_manipulation_attack(data_manipulation_bot_and_db_server): + """Test a repeating data manipulation attack.""" + data_manipulation_bot, computer, db_server_service, server = data_manipulation_bot_and_db_server + + assert db_server_service.health_state_actual is SoftwareHealthState.GOOD + + data_manipulation_bot.port_scan_p_of_success = 1 + data_manipulation_bot.data_manipulation_p_of_success = 1 + data_manipulation_bot.repeat = True + data_manipulation_bot.attack() + + assert data_manipulation_bot.attack_stage == DataManipulationAttackStage.NOT_STARTED + assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED + + computer.apply_timestep(timestep=1) + server.apply_timestep(timestep=1) + + assert data_manipulation_bot.attack_stage == DataManipulationAttackStage.NOT_STARTED + assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED + + +def test_non_repeating_data_manipulation_attack(data_manipulation_bot_and_db_server): + """Test a non repeating data manipulation attack.""" + data_manipulation_bot, computer, db_server_service, server = data_manipulation_bot_and_db_server + + assert db_server_service.health_state_actual is SoftwareHealthState.GOOD + + data_manipulation_bot.port_scan_p_of_success = 1 + data_manipulation_bot.data_manipulation_p_of_success = 1 + data_manipulation_bot.repeat = False + data_manipulation_bot.attack() + + assert data_manipulation_bot.attack_stage == DataManipulationAttackStage.SUCCEEDED + assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED + + computer.apply_timestep(timestep=1) + server.apply_timestep(timestep=1) + + assert data_manipulation_bot.attack_stage == DataManipulationAttackStage.SUCCEEDED + assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.COMPROMISED + + +def test_data_manipulation_disrupts_green_agent_connection(data_manipulation_db_server_green_client): + """Test to see that the data manipulation bot affects a green agent query.""" + network: Network = data_manipulation_db_server_green_client + + client_1: Computer = network.get_node_by_hostname("client_1") + data_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot") + + client_2: Computer = network.get_node_by_hostname("client_2") + green_db_client: DatabaseClient = client_2.software_manager.software.get("DatabaseClient") + + 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.last_query_response.get("status_code") == 200 + + data_manipulation_bot.port_scan_p_of_success = 1 + data_manipulation_bot.data_manipulation_p_of_success = 1 + data_manipulation_bot.repeat = False + 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_client.last_query_response.get("status_code") != 200