#2448: store last query response for db client

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

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:

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