#1816 - Updated the DataManipulationBot to subclass DatabaseClient. Extended logging. Dropped the Link loading logging as it was clogging up the terminal output.

src/primaite/simulator/system/services/database_service.py

@@ -81,20 +81,21 @@ class DatabaseService(Service):
             status_code = 404  # service not found
         return {"status_code": status_code, "type": "connect_response", "response": status_code == 200}
 
-    def _process_sql(self, query: str) -> Dict[str, Union[int, List[Any]]]:
+    def _process_sql(self, query: str, query_id: str) -> Dict[str, Union[int, List[Any]]]:
         """
         Executes the given SQL query and returns the result.
 
         :param query: The SQL query to be executed.
         :return: Dictionary containing status code and data fetched.
         """
+        self.sys_log.info(f"{self.name}: Running {query}")
         try:
             self._cursor.execute(query)
-
             self._conn.commit()
         except OperationalError:
             # Handle the case where the table does not exist.
-            return {"status_code": 404, "data": []}
+            self.sys_log.error(f"{self.name}: Error, query failed")
+            return {"status_code": 404, "data": {}}
         data = []
         description = self._cursor.description
         if description:
@@ -104,7 +105,7 @@ class DatabaseService(Service):
             data = self._cursor.fetchall()
             if data and headers:
                 data = {row[0]: {header: value for header, value in zip(headers, row)} for row in data}
-        return {"status_code": 200, "type": "sql", "data": data}
+        return {"status_code": 200, "type": "sql", "data": data, "uuid": query_id}
 
     def describe_state(self) -> Dict:
         """
@@ -134,7 +135,7 @@ class DatabaseService(Service):
                     self.connections.pop(session_id)
             elif payload["type"] == "sql":
                 if session_id in self.connections:
-                    result = self._process_sql(payload.get("sql"))
+                    result = self._process_sql(query=payload["sql"], query_id=payload["uuid"])
                 else:
                     result = {"status_code": 401, "type": "sql"}
         self.send(payload=result, session_id=session_id)

src/primaite/simulator/system/services/red_services/data_manipulation_bot.py

@@ -0,0 +1,49 @@
+from ipaddress import IPv4Address
+from typing import Optional
+
+from primaite.simulator.system.applications.database_client import DatabaseClient
+
+
+class DataManipulationBot(DatabaseClient):
+    """
+    Red Agent Data Integration Service.
+
+    The Service represents a bot that causes files/folders in the File System to
+    become corrupted.
+    """
+
+    server_ip_address: Optional[IPv4Address] = None
+    payload: Optional[str] = None
+    server_password: Optional[str] = None
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.name = "DataManipulationBot"
+
+    def configure(
+        self, server_ip_address: IPv4Address, server_password: Optional[str] = None, payload: Optional[str] = None
+    ):
+        """
+        Configure the DataManipulatorBot to communicate with a DatabaseService.
+
+        :param server_ip_address: The IP address of the Node the DatabaseService is on.
+        :param server_password: The password on the DatabaseService.
+        :param payload: The data manipulation query payload.
+        """
+        self.server_ip_address = server_ip_address
+        self.payload = payload
+        self.server_password = server_password
+        self.sys_log.info(f"Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}.")
+
+    def run(self):
+        """Run the DataManipulationBot."""
+        if self.server_ip_address and self.payload:
+            self.sys_log.info(f"Attempting to start the {self.name}")
+            super().run()
+            if not self.connected:
+                self.connect()
+            if self.connected:
+                self.query(self.payload)
+                self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
+        else:
+            self.sys_log.error(f"Failed to start the {self.name} as it requires both a target_io_address and payload.")

src/primaite/simulator/system/services/red_services/data_manipulator_service.py

@@ -1,34 +0,0 @@
-from ipaddress import IPv4Address
-from typing import Any, Optional
-
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
-from primaite.simulator.system.services.service import Service
-
-
-class DataManipulatorService(Service):
-    """
-    Red Agent Data Integration Service.
-
-    The Service represents a bot that causes files/folders in the File System to
-    become corrupted.
-    """
-
-    def __init__(self, **kwargs):
-        kwargs["name"] = "DataManipulatorBot"
-        kwargs["port"] = Port.POSTGRES_SERVER
-        kwargs["protocol"] = IPProtocol.TCP
-        super().__init__(**kwargs)
-
-    def start(self, target_ip_address: IPv4Address, payload: Optional[Any] = "DELETE TABLE users", **kwargs):
-        """
-        Run the DataManipulatorService actions.
-
-        :param: target_ip_address: The IP address of the target machine to attack
-        :param: payload: The payload to send to the target machine
-        """
-        super().start()
-
-        self.software_manager.send_payload_to_session_manager(
-            payload=payload, dest_ip_address=target_ip_address, dest_port=self.port
-        )