Merged PR 175: Red Service Data Manipulator Bot

## Summary
Implementation of the Data Manipulation Bot. This service sends a SQL query payload to the database server (or a given machine IP and port)

## Test process
Added a test
https://dev.azure.com/ma-dev-uk/PrimAITE/_git/PrimAITE/pullrequest/175?_a=files&path=/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulator_service.py

## Checklist
- [x] This PR is linked to a **work item**
- [x] I have performed **self-review** of the code
- [x] I have written **tests** for any new functionality added with this PR
- [ ] I have updated the **documentation** if this PR changes or adds functionality
- [ ] I have written/updated **design docs** if this PR implements new functionality
- [X] I have update the **change log**
- [x] I have run **pre-commit** checks for code style

Related work items: #1814

CHANGELOG.md

@@ -27,6 +27,8 @@ SessionManager.
 - File System - ability to emulate a node's file system during a simulation
 - Example notebooks - There is currently 1 jupyter notebook which walks through using PrimAITE
   1. Creating a simulation - this notebook explains how to build up a simulation using the Python package. (WIP)
+- Red Agent Services:
+  - Data Manipulator Bot - A red agent service which sends a payload to a target machine. (By default this payload is a SQL query that breaks a database)
 
 ## [2.0.0] - 2023-07-26
 

src/primaite/simulator/core.py

@@ -1,6 +1,6 @@
 """Core of the PrimAITE Simulator."""
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional
 from uuid import uuid4
 
 from pydantic import BaseModel, ConfigDict
@@ -140,7 +140,7 @@ class SimComponent(BaseModel):
             kwargs["uuid"] = str(uuid4())
         super().__init__(**kwargs)
         self._action_manager: ActionManager = self._init_action_manager()
-        self._parent: Optional["SimComponent"] = None
+        self.parent: Optional["SimComponent"] = None
 
     def _init_action_manager(self) -> ActionManager:
         """
@@ -213,24 +213,3 @@ class SimComponent(BaseModel):
         Override this method with anything that needs to happen within the component for it to be reset.
         """
         pass
-
-    @property
-    def parent(self) -> "SimComponent":
-        """Reference to the parent object which manages this object.
-
-        :return: Parent object.
-        :rtype: SimComponent
-        """
-        return self._parent
-
-    @parent.setter
-    def parent(self, new_parent: Union["SimComponent", None]) -> None:
-        if self._parent and new_parent:
-            msg = f"Overwriting parent of {self.uuid}. Old parent: {self._parent.uuid}, New parent: {new_parent.uuid}"
-            _LOGGER.warn(msg)
-            raise RuntimeWarning(msg)
-        self._parent = new_parent
-
-    @parent.deleter
-    def parent(self) -> None:
-        self._parent = None

src/primaite/simulator/network/transmission/transport_layer.py

@@ -59,6 +59,8 @@ class Port(Enum):
     "Alternative port for HTTP (HTTP_ALT) - Often used as an alternative HTTP port for web applications."
     HTTPS_ALT = 8443
     "Alternative port for HTTPS (HTTPS_ALT) - Used in some configurations for secure web traffic."
+    POSTGRES_SERVER = 5432
+    "Postgres SQL Server."
 
 
 class UDPHeader(BaseModel):

src/primaite/simulator/system/core/session_manager.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from ipaddress import IPv4Address
-from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
+
+from prettytable import MARKDOWN, PrettyTable
 
 from primaite.simulator.core import SimComponent
-from primaite.simulator.network.transmission.data_link_layer import Frame
+from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame
-from primaite.simulator.network.transmission.network_layer import IPProtocol
+from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.network.transmission.transport_layer import Port, TCPHeader
 
 if TYPE_CHECKING:
     from primaite.simulator.network.hardware.base import ARPCache
@@ -135,7 +137,14 @@ class SessionManager:
             dst_port = None
         return protocol, src_ip_address, dst_ip_address, src_port, dst_port
 
-    def receive_payload_from_software_manager(self, payload: Any, session_id: Optional[int] = None):
+    def receive_payload_from_software_manager(
+        self,
+        payload: Any,
+        dest_ip_address: Optional[IPv4Address] = None,
+        dest_port: Optional[Port] = None,
+        session_id: Optional[str] = None,
+        is_reattempt: bool = False,
+    ) -> Union[Any, None]:
         """
         Receive a payload from the SoftwareManager.
 
@@ -144,9 +153,50 @@ class SessionManager:
         :param payload: The payload to be sent.
         :param session_id: The Session ID the payload is to originate from. Optional. If None, one will be created.
         """
-        # TODO: Implement session creation and
+        if session_id:
+            dest_ip_address = self.sessions_by_uuid[session_id].dst_ip_address
+            dest_port = self.sessions_by_uuid[session_id].dst_port
 
-        self.send_payload_to_nic(payload, session_id)
+        dst_mac_address = self.arp_cache.get_arp_cache_mac_address(dest_ip_address)
+
+        if dst_mac_address:
+            outbound_nic = self.arp_cache.get_arp_cache_nic(dest_ip_address)
+        else:
+            if not is_reattempt:
+                self.arp_cache.send_arp_request(dest_ip_address)
+                return self.receive_payload_from_software_manager(
+                    payload=payload,
+                    dest_ip_address=dest_ip_address,
+                    dest_port=dest_port,
+                    session_id=session_id,
+                    is_reattempt=True,
+                )
+            else:
+                return
+
+        frame = Frame(
+            ethernet=EthernetHeader(src_mac_addr=outbound_nic.mac_address, dst_mac_addr=dst_mac_address),
+            ip=IPPacket(
+                src_ip_address=outbound_nic.ip_address,
+                dst_ip_address=dest_ip_address,
+            ),
+            tcp=TCPHeader(
+                src_port=dest_port,
+                dst_port=dest_port,
+            ),
+            payload=payload,
+        )
+
+        if not session_id:
+            session_key = self._get_session_key(frame, from_source=True)
+            session = self.sessions_by_key.get(session_key)
+            if not session:
+                # Create new session
+                session = Session.from_session_key(session_key)
+                self.sessions_by_key[session_key] = session
+                self.sessions_by_uuid[session.uuid] = session
+
+        outbound_nic.send_frame(frame)
 
     def send_payload_to_software_manager(self, payload: Any, session_id: int):
         """
@@ -157,18 +207,6 @@ class SessionManager:
         """
         self.software_manager.receive_payload_from_session_manger()
 
-    def send_payload_to_nic(self, payload: Any, session_id: int):
-        """
-        Send a payload across the Network.
-
-        Takes a payload and a session_id. Builds a Frame and sends it across the network via a NIC.
-
-        :param payload: The payload to be sent.
-        :param session_id: The Session ID the payload originates from
-        """
-        # TODO: Implement frame construction and sent to NIC.
-        pass
-
     def receive_payload_from_nic(self, frame: Frame):
         """
         Receive a Frame from the NIC.
@@ -187,3 +225,22 @@ class SessionManager:
             self.sessions_by_uuid[session.uuid] = session
         self.software_manager.receive_payload_from_session_manger(payload=frame, session=session)
         # TODO: Implement the frame deconstruction and send to SoftwareManager.
+
+    def show(self, markdown: bool = False):
+        """
+        Print tables describing the SessionManager.
+
+        Generate and print PrettyTable instances that show details about
+        session's destination IP Address, destination Ports and the protocol to use.
+        Output can be in Markdown format.
+
+        :param markdown: Use Markdown style in table output. Defaults to False.
+        """
+        table = PrettyTable(["Destination IP", "Port", "Protocol"])
+        if markdown:
+            table.set_style(MARKDOWN)
+        table.align = "l"
+        table.title = f"{self.sys_log.hostname} Session Manager"
+        for session in self.sessions_by_key.values():
+            table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
+        print(table)

src/primaite/simulator/system/core/software_manager.py

@@ -1,5 +1,8 @@
+from ipaddress import IPv4Address
 from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING, Union
 
+from prettytable import MARKDOWN, PrettyTable
+
 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
@@ -12,6 +15,10 @@ if TYPE_CHECKING:
     from primaite.simulator.system.core.session_manager import SessionManager
     from primaite.simulator.system.core.sys_log import SysLog
 
+from typing import Type, TypeVar
+
+ServiceClass = TypeVar("ServiceClass", bound=Service)
+
 
 class SoftwareManager:
     """A class that manages all running Services and Applications on a Node and facilitates their communication."""
@@ -28,18 +35,17 @@ class SoftwareManager:
         self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {}
         self.sys_log: SysLog = sys_log
 
-    def add_service(self, name: str, service: Service, port: Port, protocol: IPProtocol):
+    def add_service(self, service_class: Type[ServiceClass]):
         """
         Add a Service to the manager.
 
-        :param name: The name of the service.
+        :param: service_class: The class of the service to add
-        :param service: The service instance.
-        :param port: The port used by the service.
-        :param protocol: The network protocol used by the service.
         """
+        service = service_class(software_manager=self, sys_log=self.sys_log)
+
         service.software_manager = self
-        self.services[name] = service
+        self.services[service.name] = service
-        self.port_protocol_mapping[(port, protocol)] = service
+        self.port_protocol_mapping[(service.port, service.protocol)] = service
 
     def add_application(self, name: str, application: Application, port: Port, protocol: IPProtocol):
         """
@@ -75,14 +81,24 @@ class SoftwareManager:
         else:
             raise ValueError(f"No {target_software_type.name.lower()} found with the name {target_software}")
 
-    def send_payload_to_session_manger(self, payload: Any, session_id: Optional[int] = None):
+    def send_payload_to_session_manager(
+        self,
+        payload: Any,
+        dest_ip_address: Optional[IPv4Address] = None,
+        dest_port: Optional[Port] = None,
+        session_id: Optional[int] = None,
+    ):
         """
         Send a payload to the SessionManager.
 
         :param payload: The payload to be sent.
+        :param dest_ip_address: The ip address of the payload destination.
+        :param dest_port: The port of the payload destination.
         :param session_id: The Session ID the payload is to originate from. Optional.
         """
-        self.session_manager.receive_payload_from_software_manager(payload, session_id)
+        self.session_manager.receive_payload_from_software_manager(
+            payload=payload, dest_ip_address=dest_ip_address, dest_port=dest_port, session_id=session_id
+        )
 
     def receive_payload_from_session_manger(self, payload: Any, session: Session):
         """
@@ -97,3 +113,20 @@ class SoftwareManager:
         # else:
         #     raise ValueError(f"No service or application found for port {port} and protocol {protocol}")
         pass
+
+    def show(self, markdown: bool = False):
+        """
+        Prints a table of the SwitchPorts on the Switch.
+
+        :param markdown: If True, outputs the table in markdown format. Default is False.
+        """
+        table = PrettyTable(["Name", "Operating State", "Health State", "Port"])
+        if markdown:
+            table.set_style(MARKDOWN)
+        table.align = "l"
+        table.title = f"{self.sys_log.hostname} Software Manager"
+        for service in self.services.values():
+            table.add_row(
+                [service.name, service.operating_state.name, service.health_state_actual.name, service.port.value]
+            )
+        print(table)

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

@@ -0,0 +1,34 @@
+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
+        )

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

@@ -1,4 +1,3 @@
-from abc import abstractmethod
 from enum import Enum
 from typing import Any, Dict, Optional
 
@@ -33,7 +32,7 @@ class Service(IOSoftware):
     Services are programs that run in the background and may perform input/output operations.
     """
 
-    operating_state: ServiceOperatingState
+    operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED
     "The current operating state of the Service."
     restart_duration: int = 5
     "How many timesteps does it take to restart this service."
@@ -51,7 +50,6 @@ class Service(IOSoftware):
         am.add_action("enable", Action(func=lambda request, context: self.enable()))
         return am
 
-    @abstractmethod
     def describe_state(self) -> Dict:
         """
         Produce a dictionary describing the current state of this object.
@@ -102,49 +100,49 @@ class Service(IOSoftware):
         """Stop the service."""
         _LOGGER.debug(f"Stopping service {self.name}")
         if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
-            self.parent.sys_log.info(f"Stopping service {self.name}")
+            self.sys_log.info(f"Stopping service {self.name}")
             self.operating_state = ServiceOperatingState.STOPPED
 
-    def start(self) -> None:
+    def start(self, **kwargs) -> None:
         """Start the service."""
         _LOGGER.debug(f"Starting service {self.name}")
         if self.operating_state == ServiceOperatingState.STOPPED:
-            self.parent.sys_log.info(f"Starting service {self.name}")
+            self.sys_log.info(f"Starting service {self.name}")
             self.operating_state = ServiceOperatingState.RUNNING
 
     def pause(self) -> None:
         """Pause the service."""
         _LOGGER.debug(f"Pausing service {self.name}")
         if self.operating_state == ServiceOperatingState.RUNNING:
-            self.parent.sys_log.info(f"Pausing service {self.name}")
+            self.sys_log.info(f"Pausing service {self.name}")
             self.operating_state = ServiceOperatingState.PAUSED
 
     def resume(self) -> None:
         """Resume paused service."""
         _LOGGER.debug(f"Resuming service {self.name}")
         if self.operating_state == ServiceOperatingState.PAUSED:
-            self.parent.sys_log.info(f"Resuming service {self.name}")
+            self.sys_log.info(f"Resuming service {self.name}")
             self.operating_state = ServiceOperatingState.RUNNING
 
     def restart(self) -> None:
         """Restart running service."""
         _LOGGER.debug(f"Restarting service {self.name}")
         if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
-            self.parent.sys_log.info(f"Pausing service {self.name}")
+            self.sys_log.info(f"Pausing service {self.name}")
             self.operating_state = ServiceOperatingState.RESTARTING
             self.restart_countdown = self.restarting_duration
 
     def disable(self) -> None:
         """Disable the service."""
         _LOGGER.debug(f"Disabling service {self.name}")
-        self.parent.sys_log.info(f"Disabling Application {self.name}")
+        self.sys_log.info(f"Disabling Application {self.name}")
         self.operating_state = ServiceOperatingState.DISABLED
 
     def enable(self) -> None:
         """Enable the disabled service."""
         _LOGGER.debug(f"Enabling service {self.name}")
         if self.operating_state == ServiceOperatingState.DISABLED:
-            self.parent.sys_log.info(f"Enabling Application {self.name}")
+            self.sys_log.info(f"Enabling Application {self.name}")
             self.operating_state = ServiceOperatingState.STOPPED
 
     def apply_timestep(self, timestep: int) -> None:

src/primaite/simulator/system/software.py

@@ -1,9 +1,10 @@
 from abc import abstractmethod
 from enum import Enum
-from typing import Any, Dict, Set
+from typing import Any, Dict
 
 from primaite.simulator.core import Action, ActionManager, SimComponent
 from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.core.sys_log import SysLog
 
 
 class SoftwareType(Enum):
@@ -62,11 +63,11 @@ class Software(SimComponent):
 
     name: str
     "The name of the software."
-    health_state_actual: SoftwareHealthState
+    health_state_actual: SoftwareHealthState = SoftwareHealthState.GOOD
     "The actual health state of the software."
-    health_state_visible: SoftwareHealthState
+    health_state_visible: SoftwareHealthState = SoftwareHealthState.GOOD
     "The health state of the software visible to the red agent."
-    criticality: SoftwareCriticality
+    criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
     "The criticality level of the software."
     patching_count: int = 0
     "The count of patches applied to the software, defaults to 0."
@@ -74,6 +75,10 @@ class Software(SimComponent):
     "The count of times the software has been scanned, defaults to 0."
     revealed_to_red: bool = False
     "Indicates if the software has been revealed to red agent, defaults is False."
+    software_manager: Any = None
+    "An instance of Software Manager that is used by the parent node."
+    sys_log: SysLog = None
+    "An instance of SysLog that is used by the parent node."
 
     def _init_action_manager(self) -> ActionManager:
         am = super()._init_action_manager()
@@ -132,7 +137,6 @@ class Software(SimComponent):
         """
         self.health_state_actual = health_state
 
-    @abstractmethod
     def install(self) -> None:
         """
         Perform first-time setup of this service on a node.
@@ -175,8 +179,8 @@ class IOSoftware(Software):
     "Indicates if the software uses TCP protocol for communication. Default is True."
     udp: bool = True
     "Indicates if the software uses UDP protocol for communication. Default is True."
-    ports: Set[Port]
+    port: Port
-    "The set of ports to which the software is connected."
+    "The port to which the software is connected."
 
     @abstractmethod
     def describe_state(self) -> Dict:

tests/integration_tests/system/test_database_on_node.py

@@ -11,9 +11,7 @@ def test_installing_database():
         health_state_actual=SoftwareHealthState.GOOD,
         health_state_visible=SoftwareHealthState.GOOD,
         criticality=SoftwareCriticality.MEDIUM,
-        ports=[
+        port=Port.SQL_SERVER,
-            Port.SQL_SERVER,
-        ],
         operating_state=ServiceOperatingState.RUNNING,
     )
 
@@ -40,9 +38,7 @@ def test_uninstalling_database():
         health_state_actual=SoftwareHealthState.GOOD,
         health_state_visible=SoftwareHealthState.GOOD,
         criticality=SoftwareCriticality.MEDIUM,
-        ports=[
+        port=Port.SQL_SERVER,
-            Port.SQL_SERVER,
-        ],
         operating_state=ServiceOperatingState.RUNNING,
     )
 

tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulator_service.py

@@ -0,0 +1,32 @@
+from ipaddress import IPv4Address
+
+from primaite.simulator.network.hardware.base import Node
+from primaite.simulator.network.networks import arcd_uc2_network
+from primaite.simulator.network.transmission.network_layer import IPProtocol
+from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.simulator.system.services.red_services.data_manipulator_service import DataManipulatorService
+
+
+def test_creation():
+    network = arcd_uc2_network()
+
+    client_1: Node = network.get_node_by_hostname("client_1")
+
+    client_1.software_manager.add_service(service_class=DataManipulatorService)
+
+    data_manipulator_service: DataManipulatorService = client_1.software_manager.services["DataManipulatorBot"]
+
+    assert data_manipulator_service.name == "DataManipulatorBot"
+    assert data_manipulator_service.port == Port.POSTGRES_SERVER
+    assert data_manipulator_service.protocol == IPProtocol.TCP
+
+    # should have no session yet
+    assert len(client_1.session_manager.sessions_by_uuid) == 0
+
+    try:
+        data_manipulator_service.start(target_ip_address=IPv4Address("192.168.1.14"))
+    except Exception as e:
+        assert False, f"Test was not supposed to throw exception: {e}"
+
+    # there should be a session after the service is started
+    assert len(client_1.session_manager.sessions_by_uuid) == 1

tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py

@@ -10,8 +10,6 @@ def test_creation():
         health_state_actual=SoftwareHealthState.GOOD,
         health_state_visible=SoftwareHealthState.GOOD,
         criticality=SoftwareCriticality.MEDIUM,
-        ports=[
+        port=Port.SQL_SERVER,
-            Port.SQL_SERVER,
-        ],
         operating_state=ServiceOperatingState.RUNNING,
     )