diff --git a/src/primaite/simulator/network/transmission/transport_layer.py b/src/primaite/simulator/network/transmission/transport_layer.py index b95b4a74..d4318baf 100644 --- a/src/primaite/simulator/network/transmission/transport_layer.py +++ b/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): diff --git a/src/primaite/simulator/system/core/session_manager.py b/src/primaite/simulator/system/core/session_manager.py index 7f3d22c5..be20a28d 100644 --- a/src/primaite/simulator/system/core/session_manager.py +++ b/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.network_layer import IPProtocol -from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.network.transmission.data_link_layer import EthernetHeader, Frame +from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol +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) diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 411fb6e9..312f6d84 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/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: The service instance. - :param port: The port used by the service. - :param protocol: The network protocol used by the service. + :param service_class: The class of the service to add """ + service = service_class(software_manager=self, sys_log=self.sys_log) + service.software_manager = self - self.services[name] = service - self.port_protocol_mapping[(port, protocol)] = service + self.services[service.name] = 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) diff --git a/src/primaite/simulator/system/services/red_services/data_manipulator_service.py b/src/primaite/simulator/system/services/red_services/data_manipulator_service.py new file mode 100644 index 00000000..29f0d3f8 --- /dev/null +++ b/src/primaite/simulator/system/services/red_services/data_manipulator_service.py @@ -0,0 +1,28 @@ +from ipaddress import IPv4Address + +from primaite.simulator.network.transmission.network_layer import IPProtocol +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.core.software_manager import SoftwareManager +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 run(self): + """Run the DataManipulatorService actions.""" + software_manager: SoftwareManager = self.software_manager + software_manager.send_payload_to_session_manager( + payload="SELECT * FROM users", dest_ip_address=IPv4Address("192.168.1.14"), dest_port=self.port + ) diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index f9cc784d..6a8c9abf 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/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. diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 605a062b..7f206311 100644 --- a/src/primaite/simulator/system/software.py +++ b/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] - "The set of ports to which the software is connected." + port: Port + "The port to which the software is connected." @abstractmethod def describe_state(self) -> Dict: diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index 73d19339..058bb590 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/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.SQL_SERVER, - ], + port=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.SQL_SERVER, - ], + port=Port.SQL_SERVER, operating_state=ServiceOperatingState.RUNNING, ) diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py index ea5c1b83..ebc5536f 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py +++ b/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.SQL_SERVER, - ], + port=Port.SQL_SERVER, operating_state=ServiceOperatingState.RUNNING, )