Merge remote-tracking branch 'origin/dev' into dev-v3.0.0b6

src/primaite/simulator/system/applications/database_client.py

@@ -5,7 +5,7 @@ from uuid import uuid4
 from primaite import getLogger
 from primaite.simulator.network.transmission.network_layer import IPProtocol
 from primaite.simulator.network.transmission.transport_layer import Port
-from primaite.simulator.system.applications.application import Application, ApplicationOperatingState
+from primaite.simulator.system.applications.application import Application
 from primaite.simulator.system.core.software_manager import SoftwareManager
 
 _LOGGER = getLogger(__name__)
@@ -23,7 +23,6 @@ class DatabaseClient(Application):
 
     server_ip_address: Optional[IPv4Address] = None
     server_password: Optional[str] = None
-    connected: bool = False
     _query_success_tracker: Dict[str, bool] = {}
 
     def __init__(self, **kwargs):
@@ -66,18 +65,24 @@ class DatabaseClient(Application):
         self.server_password = server_password
         self.sys_log.info(f"{self.name}: Configured the {self.name} with {server_ip_address=}, {server_password=}.")
 
-    def connect(self) -> bool:
+    def connect(self, connection_id: Optional[str] = None) -> bool:
         """Connect to a Database Service."""
         if not self._can_perform_action():
             return False
 
-        if not self.connected:
-            return self._connect(self.server_ip_address, self.server_password)
-        # already connected
-        return True
+        if not connection_id:
+            connection_id = str(uuid4())
+
+        return self._connect(
+            server_ip_address=self.server_ip_address, password=self.server_password, connection_id=connection_id
+        )
 
     def _connect(
-        self, server_ip_address: IPv4Address, password: Optional[str] = None, is_reattempt: bool = False
+        self,
+        server_ip_address: IPv4Address,
+        connection_id: Optional[str] = None,
+        password: Optional[str] = None,
+        is_reattempt: bool = False,
     ) -> bool:
         """
         Connects the DatabaseClient to the DatabaseServer.
@@ -92,33 +97,58 @@ class DatabaseClient(Application):
         :type: is_reattempt: Optional[bool]
         """
         if is_reattempt:
-            if self.connected:
-                self.sys_log.info(f"{self.name}: DatabaseClient connection to {server_ip_address} authorised")
+            if self.connections.get(connection_id):
+                self.sys_log.info(
+                    f"{self.name} {connection_id=}: DatabaseClient connection to {server_ip_address} authorised"
+                )
                 self.server_ip_address = server_ip_address
-                return self.connected
+                return True
             else:
-                self.sys_log.info(f"{self.name}: DatabaseClient connection to {server_ip_address} declined")
+                self.sys_log.info(
+                    f"{self.name} {connection_id=}: DatabaseClient connection to {server_ip_address} declined"
+                )
                 return False
-        payload = {"type": "connect_request", "password": password}
+        payload = {
+            "type": "connect_request",
+            "password": password,
+            "connection_id": connection_id,
+        }
         software_manager: SoftwareManager = self.software_manager
         software_manager.send_payload_to_session_manager(
             payload=payload, dest_ip_address=server_ip_address, dest_port=self.port
         )
-        return self._connect(server_ip_address, password, True)
+        return self._connect(
+            server_ip_address=server_ip_address, password=password, connection_id=connection_id, is_reattempt=True
+        )
 
-    def disconnect(self):
+    def disconnect(self, connection_id: Optional[str] = None) -> bool:
         """Disconnect from the Database Service."""
-        if self.connected and self.operating_state is ApplicationOperatingState.RUNNING:
-            software_manager: SoftwareManager = self.software_manager
-            software_manager.send_payload_to_session_manager(
-                payload={"type": "disconnect"}, dest_ip_address=self.server_ip_address, dest_port=self.port
-            )
+        if not self._can_perform_action():
+            self.sys_log.error(f"Unable to disconnect - {self.name} is {self.operating_state.name}")
+            return False
 
-            self.sys_log.info(f"{self.name}: DatabaseClient disconnected from {self.server_ip_address}")
-            self.server_ip_address = None
-            self.connected = False
+        # if there are no connections - nothing to disconnect
+        if not len(self.connections):
+            self.sys_log.error(f"Unable to disconnect - {self.name} has no active connections.")
+            return False
 
-    def _query(self, sql: str, query_id: str, is_reattempt: bool = False) -> bool:
+        # if no connection provided, disconnect the first connection
+        if not connection_id:
+            connection_id = list(self.connections.keys())[0]
+
+        software_manager: SoftwareManager = self.software_manager
+        software_manager.send_payload_to_session_manager(
+            payload={"type": "disconnect", "connection_id": connection_id},
+            dest_ip_address=self.server_ip_address,
+            dest_port=self.port,
+        )
+        self.remove_connection(connection_id=connection_id)
+
+        self.sys_log.info(
+            f"{self.name}: DatabaseClient disconnected connection {connection_id} from {self.server_ip_address}"
+        )
+
+    def _query(self, sql: str, query_id: str, connection_id: str, is_reattempt: bool = False) -> bool:
         """
         Send a query to the connected database server.
 
@@ -141,19 +171,17 @@ class DatabaseClient(Application):
         else:
             software_manager: SoftwareManager = self.software_manager
             software_manager.send_payload_to_session_manager(
-                payload={"type": "sql", "sql": sql, "uuid": query_id},
+                payload={"type": "sql", "sql": sql, "uuid": query_id, "connection_id": connection_id},
                 dest_ip_address=self.server_ip_address,
                 dest_port=self.port,
             )
-            return self._query(sql=sql, query_id=query_id, is_reattempt=True)
+            return self._query(sql=sql, query_id=query_id, connection_id=connection_id, is_reattempt=True)
 
     def run(self) -> None:
         """Run the DatabaseClient."""
         super().run()
-        if self.operating_state == ApplicationOperatingState.RUNNING:
-            self.connect()
 
-    def query(self, sql: str, is_reattempt: bool = False) -> bool:
+    def query(self, sql: str, connection_id: Optional[str] = None) -> bool:
         """
         Send a query to the Database Service.
 
@@ -164,20 +192,17 @@ class DatabaseClient(Application):
         if not self._can_perform_action():
             return False
 
-        if self.connected:
-            query_id = str(uuid4())
+        if connection_id is None:
+            connection_id = str(uuid4())
+
+        if not self.connections.get(connection_id):
+            if not self.connect(connection_id=connection_id):
+                return False
 
             # Initialise the tracker of this ID to False
-            self._query_success_tracker[query_id] = False
-            return self._query(sql=sql, query_id=query_id)
-        else:
-            if is_reattempt:
-                return False
-
-            if not self.connect():
-                return False
-
-            self.query(sql=sql, is_reattempt=True)
+        uuid = str(uuid4())
+        self._query_success_tracker[uuid] = False
+        return self._query(sql=sql, query_id=uuid, connection_id=connection_id)
 
     def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
         """
@@ -192,13 +217,13 @@ class DatabaseClient(Application):
 
         if isinstance(payload, dict) and payload.get("type"):
             if payload["type"] == "connect_response":
-                self.connected = payload["response"] == True
+                if payload["response"] is True:
+                    # add connection
+                    self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id)
             elif payload["type"] == "sql":
                 query_id = payload.get("uuid")
                 status_code = payload.get("status_code")
                 self._query_success_tracker[query_id] = status_code == 200
                 if self._query_success_tracker[query_id]:
                     _LOGGER.debug(f"Received payload {payload}")
-                else:
-                    self.connected = False
         return True

src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py

@@ -0,0 +1,207 @@
+from enum import IntEnum
+from ipaddress import IPv4Address
+from typing import Optional
+
+from primaite import getLogger
+from primaite.game.science import simulate_trial
+from primaite.simulator.core import RequestManager, RequestType
+from primaite.simulator.system.applications.database_client import DatabaseClient
+
+_LOGGER = getLogger(__name__)
+
+
+class DataManipulationAttackStage(IntEnum):
+    """
+    Enumeration representing different stages of a data manipulation attack.
+
+    This enumeration defines the various stages a data manipulation attack can be in during its lifecycle in the
+    simulation. Each stage represents a specific phase in the attack process.
+    """
+
+    NOT_STARTED = 0
+    "Indicates that the attack has not started yet."
+    LOGON = 1
+    "The stage where logon procedures are simulated."
+    PORT_SCAN = 2
+    "Represents the stage of performing a horizontal port scan on the target."
+    ATTACKING = 3
+    "Stage of actively attacking the target."
+    SUCCEEDED = 4
+    "Indicates the attack has been successfully completed."
+    FAILED = 5
+    "Signifies that the attack has failed."
+
+
+class DataManipulationBot(DatabaseClient):
+    """A bot that simulates a script which performs a SQL injection attack."""
+
+    server_ip_address: Optional[IPv4Address] = None
+    payload: Optional[str] = None
+    server_password: Optional[str] = None
+    port_scan_p_of_success: float = 0.1
+    data_manipulation_p_of_success: float = 0.1
+
+    attack_stage: DataManipulationAttackStage = DataManipulationAttackStage.NOT_STARTED
+    repeat: bool = False
+    "Whether to repeat attacking once finished."
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.name = "DataManipulationBot"
+
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting DataManipulationBot original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {
+            "server_ip_address",
+            "payload",
+            "server_password",
+            "port_scan_p_of_success",
+            "data_manipulation_p_of_success",
+            "attack_stage",
+            "repeat",
+        }
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting DataManipulationBot state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
+    def _init_request_manager(self) -> RequestManager:
+        rm = super()._init_request_manager()
+
+        rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.attack()))
+
+        return rm
+
+    def configure(
+        self,
+        server_ip_address: IPv4Address,
+        server_password: Optional[str] = None,
+        payload: Optional[str] = None,
+        port_scan_p_of_success: float = 0.1,
+        data_manipulation_p_of_success: float = 0.1,
+        repeat: bool = True,
+    ):
+        """
+        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.
+        :param port_scan_p_of_success: The probability of success for the port scan stage.
+        :param data_manipulation_p_of_success: The probability of success for the data manipulation stage.
+        :param repeat: Whether to repeat attacking once finished.
+        """
+        self.server_ip_address = server_ip_address
+        self.payload = payload
+        self.server_password = server_password
+        self.port_scan_p_of_success = port_scan_p_of_success
+        self.data_manipulation_p_of_success = data_manipulation_p_of_success
+        self.repeat = repeat
+        self.sys_log.info(
+            f"{self.name}: Configured the {self.name} with {server_ip_address=}, {payload=}, {server_password=}, "
+            f"{repeat=}."
+        )
+
+    def _logon(self):
+        """
+        Simulate the logon process as the initial stage of the attack.
+
+        Advances the attack stage to `LOGON` if successful.
+        """
+        if self.attack_stage == DataManipulationAttackStage.NOT_STARTED:
+            # Bypass this stage as we're not dealing with logon for now
+            self.sys_log.info(f"{self.name}: ")
+            self.attack_stage = DataManipulationAttackStage.LOGON
+
+    def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
+        """
+        Perform a simulated port scan to check for open SQL ports.
+
+        Advances the attack stage to `PORT_SCAN` if successful.
+
+        :param p_of_success: Probability of successful port scan, by default 0.1.
+        """
+        if self.attack_stage == DataManipulationAttackStage.LOGON:
+            # perform a port scan to identify that the SQL port is open on the server
+            if simulate_trial(p_of_success):
+                self.sys_log.info(f"{self.name}: Performing port scan")
+                # perform the port scan
+                port_is_open = True  # Temporary; later we can implement NMAP port scan.
+                if port_is_open:
+                    self.sys_log.info(f"{self.name}: ")
+                    self.attack_stage = DataManipulationAttackStage.PORT_SCAN
+
+    def _perform_data_manipulation(self, p_of_success: Optional[float] = 0.1):
+        """
+        Execute the data manipulation attack on the target.
+
+        Advances the attack stage to `COMPLETE` if successful, or 'FAILED' if unsuccessful.
+
+        :param p_of_success: Probability of successfully performing data manipulation, by default 0.1.
+        """
+        if self.attack_stage == DataManipulationAttackStage.PORT_SCAN:
+            # perform the actual data manipulation attack
+            if simulate_trial(p_of_success):
+                self.sys_log.info(f"{self.name}: Performing data manipulation")
+                # perform the attack
+                if not len(self.connections):
+                    self.connect()
+                if len(self.connections):
+                    self.query(self.payload)
+                    self.sys_log.info(f"{self.name} payload delivered: {self.payload}")
+                    attack_successful = True
+                    if attack_successful:
+                        self.sys_log.info(f"{self.name}: Data manipulation successful")
+                        self.attack_stage = DataManipulationAttackStage.SUCCEEDED
+                    else:
+                        self.sys_log.info(f"{self.name}: Data manipulation failed")
+                        self.attack_stage = DataManipulationAttackStage.FAILED
+
+    def run(self):
+        """
+        Run the Data Manipulation Bot.
+
+        Calls the parent classes execute method before starting the application loop.
+        """
+        super().run()
+
+    def attack(self):
+        """Perform the attack steps after opening the application."""
+        if not self._can_perform_action():
+            _LOGGER.debug("Data manipulation application attempted to execute but it cannot perform actions right now.")
+            self.run()
+        self._application_loop()
+
+    def _application_loop(self):
+        """
+        The main application loop of the bot, handling the attack process.
+
+        This is the core loop where the bot sequentially goes through the stages of the attack.
+        """
+        if not self._can_perform_action():
+            return
+        if self.server_ip_address and self.payload:
+            self.sys_log.info(f"{self.name}: Running")
+            self._logon()
+            self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
+            self._perform_data_manipulation(p_of_success=self.data_manipulation_p_of_success)
+
+            if self.repeat and self.attack_stage in (
+                DataManipulationAttackStage.SUCCEEDED,
+                DataManipulationAttackStage.FAILED,
+            ):
+                self.attack_stage = DataManipulationAttackStage.NOT_STARTED
+        else:
+            self.sys_log.error(f"{self.name}: Failed to start as it requires both a target_ip_address and payload.")
+
+    def apply_timestep(self, timestep: int) -> None:
+        """
+        Apply a timestep to the bot, triggering the application loop.
+
+        :param timestep: The timestep value to update the bot's state.
+        """
+        pass

src/primaite/simulator/system/applications/red_applications/dos_bot.py

@@ -0,0 +1,192 @@
+from enum import IntEnum
+from ipaddress import IPv4Address
+from typing import Optional
+
+from primaite import getLogger
+from primaite.game.science import simulate_trial
+from primaite.simulator.core import RequestManager, RequestType
+from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.applications.application import Application
+from primaite.simulator.system.applications.database_client import DatabaseClient
+
+_LOGGER = getLogger(__name__)
+
+
+class DoSAttackStage(IntEnum):
+    """Enum representing the different stages of a Denial of Service attack."""
+
+    NOT_STARTED = 0
+    "Attack not yet started."
+
+    PORT_SCAN = 1
+    "Attack is in discovery stage - checking if provided ip and port are open."
+
+    ATTACKING = 2
+    "Denial of Service attack is in progress."
+
+    COMPLETED = 3
+    "Attack is completed."
+
+
+class DoSBot(DatabaseClient, Application):
+    """A bot that simulates a Denial of Service attack."""
+
+    target_ip_address: Optional[IPv4Address] = None
+    """IP address of the target service."""
+
+    target_port: Optional[Port] = None
+    """Port of the target service."""
+
+    payload: Optional[str] = None
+    """Payload to deliver to the target service as part of the denial of service attack."""
+
+    repeat: bool = False
+    """If true, the Denial of Service bot will keep performing the attack."""
+
+    attack_stage: DoSAttackStage = DoSAttackStage.NOT_STARTED
+    """Current stage of the DoS kill chain."""
+
+    port_scan_p_of_success: float = 0.1
+    """Probability of port scanning being sucessful."""
+
+    dos_intensity: float = 1.0
+    """How much of the max sessions will be used by the DoS when attacking."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.name = "DoSBot"
+        self.max_sessions = 1000  # override normal max sessions
+
+    def set_original_state(self):
+        """Set the original state of the Denial of Service Bot."""
+        _LOGGER.debug(f"Setting {self.name} original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {
+            "target_ip_address",
+            "target_port",
+            "payload",
+            "repeat",
+            "attack_stage",
+            "max_sessions",
+            "port_scan_p_of_success",
+            "dos_intensity",
+        }
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting {self.name} state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
+    def _init_request_manager(self) -> RequestManager:
+        rm = super()._init_request_manager()
+
+        rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run()))
+
+        return rm
+
+    def configure(
+        self,
+        target_ip_address: IPv4Address,
+        target_port: Optional[Port] = Port.POSTGRES_SERVER,
+        payload: Optional[str] = None,
+        repeat: bool = False,
+        port_scan_p_of_success: float = 0.1,
+        dos_intensity: float = 1.0,
+        max_sessions: int = 1000,
+    ):
+        """
+        Configure the Denial of Service bot.
+
+        :param: target_ip_address: The IP address of the Node containing the target service.
+        :param: target_port: The port of the target service. Optional - Default is `Port.HTTP`
+        :param: payload: The payload the DoS Bot will throw at the target service. Optional - Default is `None`
+        :param: repeat: If True, the bot will maintain the attack. Optional - Default is `True`
+        :param: port_scan_p_of_success: The chance of the port scan being sucessful. Optional - Default is 0.1 (10%)
+        :param: dos_intensity: The intensity of the DoS attack.
+            Multiplied with the application's max session - Default is 1.0
+        :param: max_sessions: The maximum number of sessions the DoS bot will attack with. Optional - Default is 1000
+        """
+        self.target_ip_address = target_ip_address
+        self.target_port = target_port
+        self.payload = payload
+        self.repeat = repeat
+        self.port_scan_p_of_success = port_scan_p_of_success
+        self.dos_intensity = dos_intensity
+        self.max_sessions = max_sessions
+        self.sys_log.info(
+            f"{self.name}: Configured the {self.name} with {target_ip_address=}, {target_port=}, {payload=}, "
+            f"{repeat=}, {port_scan_p_of_success=}, {dos_intensity=}, {max_sessions=}."
+        )
+
+    def run(self):
+        """Run the Denial of Service Bot."""
+        super().run()
+        self._application_loop()
+
+    def _application_loop(self):
+        """
+        The main application loop for the Denial of Service bot.
+
+        The loop goes through the stages of a DoS attack.
+        """
+        if not self._can_perform_action():
+            return
+
+        # DoS bot cannot do anything without a target
+        if not self.target_ip_address or not self.target_port:
+            self.sys_log.error(
+                f"{self.name} is not properly configured. {self.target_ip_address=}, {self.target_port=}"
+            )
+            return
+
+        self.clear_connections()
+        self._perform_port_scan(p_of_success=self.port_scan_p_of_success)
+        self._perform_dos()
+
+        if self.repeat and self.attack_stage is DoSAttackStage.ATTACKING:
+            self.attack_stage = DoSAttackStage.NOT_STARTED
+        else:
+            self.attack_stage = DoSAttackStage.COMPLETED
+
+    def _perform_port_scan(self, p_of_success: Optional[float] = 0.1):
+        """
+        Perform a simulated port scan to check for open SQL ports.
+
+        Advances the attack stage to `PORT_SCAN` if successful.
+
+        :param p_of_success: Probability of successful port scan, by default 0.1.
+        """
+        if self.attack_stage == DoSAttackStage.NOT_STARTED:
+            # perform a port scan to identify that the SQL port is open on the server
+            if simulate_trial(p_of_success):
+                self.sys_log.info(f"{self.name}: Performing port scan")
+                # perform the port scan
+                port_is_open = True  # Temporary; later we can implement NMAP port scan.
+                if port_is_open:
+                    self.sys_log.info(f"{self.name}: ")
+                    self.attack_stage = DoSAttackStage.PORT_SCAN
+
+    def _perform_dos(self):
+        """
+        Perform the Denial of Service attack.
+
+        DoSBot does this by clogging up the available connections to a service.
+        """
+        if not self.attack_stage == DoSAttackStage.PORT_SCAN:
+            return
+        self.attack_stage = DoSAttackStage.ATTACKING
+        self.server_ip_address = self.target_ip_address
+        self.port = self.target_port
+
+        dos_sessions = int(float(self.max_sessions) * self.dos_intensity)
+        for i in range(dos_sessions):
+            self.connect()
+
+    def apply_timestep(self, timestep: int) -> None:
+        """
+        Apply a timestep to the bot, iterate through the application loop.
+
+        :param timestep: The timestep value to update the bot's state.
+        """
+        self._application_loop()