diff --git a/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py b/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py index af5c37b9..9c840616 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py @@ -10,6 +10,7 @@ from primaite.simulator.network.protocols.masquerade import C2Payload, Masquerad 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 +from primaite.simulator.system.core.session_manager import Session # TODO: # Complete C2 Server and C2 Beacon TODOs @@ -52,7 +53,7 @@ class AbstractC2(Application, identifier="AbstractC2"): """Indicates if the c2 server and c2 beacon are currently connected.""" c2_remote_connection: IPv4Address = None - """The IPv4 Address of the remote c2 connection. (Either the IP of the beacon or the server)""" + """The IPv4 Address of the remote c2 connection. (Either the IP of the beacon or the server).""" keep_alive_sent: bool = False """Indicates if a keep alive has been sent this timestep. Used to prevent packet storms.""" @@ -65,12 +66,18 @@ class AbstractC2(Application, identifier="AbstractC2"): # The c2 server parses the keep alive and sets these accordingly. # The c2 beacon will set this attributes upon installation and configuration - current_masquerade_protocol: Enum = IPProtocol.TCP + current_masquerade_protocol: IPProtocol = IPProtocol.TCP """The currently chosen protocol that the C2 traffic is masquerading as. Defaults as TCP.""" - current_masquerade_port: Enum = Port.HTTP + current_masquerade_port: Port = Port.HTTP """The currently chosen port that the C2 traffic is masquerading as. Defaults at HTTP.""" + current_c2_session: Session = None + """The currently active session that the C2 Traffic is using. Set after establishing connection.""" + + # TODO: Create a attribute call 'LISTENER' which indicates whenever the c2 application should listen for incoming connections or establish connections. + # This in order to simulate a blind shell (the current implementation is more akin to a reverse shell) + # TODO: Move this duplicate method from NMAP class into 'Application' to adhere to DRY principle. def _can_perform_network_action(self) -> bool: """ @@ -99,8 +106,8 @@ class AbstractC2(Application, identifier="AbstractC2"): return super().describe_state() def __init__(self, **kwargs): - kwargs["port"] = Port.NONE - kwargs["protocol"] = IPProtocol.NONE + kwargs["port"] = Port.HTTP # TODO: Update this post application/services requiring to listen to multiple ports + kwargs["protocol"] = IPProtocol.TCP # Update this as well super().__init__(**kwargs) # Validate call ensures we are only handling Masquerade Packets. @@ -147,7 +154,7 @@ class AbstractC2(Application, identifier="AbstractC2"): return False # Abstract method - # Used in C2 server to prase and receive the output of commands sent to the c2 beacon. + # Used in C2 server to parse and receive the output of commands sent to the c2 beacon. @abstractmethod def _handle_command_output(payload): """Abstract Method: Used in C2 server to prase and receive the output of commands sent to the c2 beacon.""" @@ -170,32 +177,35 @@ class AbstractC2(Application, identifier="AbstractC2"): :return: True if successfully handled, false otherwise. :rtype: Bool """ - self.sys_log.info(f"{self.name}: Keep Alive Received from {self.c2_remote_connection}") - # Using this guard clause to prevent packet storms and recognise that we've achieved a connection. - if self.keep_alive_sent: - self.sys_log.info(f"{self.name}: Connection successfully established with {self.c2_remote_connection}") - self.c2_connection_active = True # Sets the connection to active - self.keep_alive_inactivity = 0 # Sets the keep alive inactivity to zero + self.sys_log.info(f"{self.name}: Keep Alive Received from {self.c2_remote_connection}.") + + self.c2_connection_active = True # Sets the connection to active + self.keep_alive_inactivity = 0 # Sets the keep alive inactivity to zero + self.current_c2_session = self.software_manager.session_manager.sessions_by_uuid[session_id] + + # Using this guard clause to prevent packet storms and recognise that we've achieved a connection. + # This guard clause triggers on the c2 suite that establishes connection. + if self.keep_alive_sent == True: # Return early without sending another keep alive and then setting keep alive_sent false for next timestep. self.keep_alive_sent = False return True # If we've reached this part of the method then we've received a keep alive but haven't sent a reply. # Therefore we also need to configure the masquerade attributes based off the keep alive sent. - if not self._resolve_keep_alive(self, payload): + if self._resolve_keep_alive(payload, session_id) == False: + self.sys_log.warning(f"{self.name}: Keep Alive Could not be resolved correctly. Refusing Keep Alive.") return False # If this method returns true then we have sent successfully sent a keep alive. - if self._send_keep_alive(self, session_id): - self.keep_alive_sent = True # Setting the guard clause to true (prevents packet storms.) - return True - # Return false if we're unable to send handle the keep alive correctly. - else: - return False + self.sys_log.info(f"{self.name}: Connection successfully established with {self.c2_remote_connection}.") + + return self._send_keep_alive(session_id) - def receive(self, payload: MasqueradePacket, session_id: Optional[str] = None) -> bool: + + # from_network_interface=from_network_interface + def receive(self, payload: MasqueradePacket, session_id: Optional[str] = None, **kwargs) -> bool: """Receives masquerade packets. Used by both c2 server and c2 client. :param payload: The Masquerade Packet to be received. @@ -220,20 +230,20 @@ class AbstractC2(Application, identifier="AbstractC2"): masquerade_protocol=self.current_masquerade_protocol, masquerade_port=self.current_masquerade_port, payload_type=C2Payload.KEEP_ALIVE, + command=None ) - - # C2 Server will need to c2_remote_connection after it receives it's first keep alive. + # We need to set this guard clause to true before sending the keep alive (prevents packet storms.) + self.keep_alive_sent = True + # C2 Server will need to configure c2_remote_connection after it receives it's first keep alive. if self.send( - self, payload=keep_alive_packet, dest_ip_address=self.c2_remote_connection, - port=self.current_masquerade_port, - protocol=self.current_masquerade_protocol, + dest_port=self.current_masquerade_port, + ip_protocol=self.current_masquerade_protocol, session_id=session_id, ): self.sys_log.info(f"{self.name}: Keep Alive sent to {self.c2_remote_connection}") self.sys_log.debug(f"{self.name}: on {self.current_masquerade_port} via {self.current_masquerade_protocol}") - self.receive(payload=keep_alive_packet) return True else: self.sys_log.warning( @@ -242,7 +252,7 @@ class AbstractC2(Application, identifier="AbstractC2"): return False - def _resolve_keep_alive(self, payload: MasqueradePacket) -> bool: + def _resolve_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: """ Parses the Masquerade Port/Protocol within the received Keep Alive packet. @@ -257,8 +267,8 @@ class AbstractC2(Application, identifier="AbstractC2"): :rtype: bool """ # Validating that they are valid Enums. - if payload.masquerade_port or payload.masquerade_protocol != Enum: - self.sys_log.warning(f"{self.name}: Received invalid Masquerade type. Port: {type(payload.masquerade_port)} Protocol: {type(payload.masquerade_protocol)}") + if not isinstance(payload.masquerade_port, Port) or not isinstance(payload.masquerade_protocol, IPProtocol): + self.sys_log.warning(f"{self.name}: Received invalid Masquerade type. Port: {type(payload.masquerade_port)} Protocol: {type(payload.masquerade_protocol)}.") return False # TODO: Validation on Ports (E.g only allow HTTP, FTP etc) # Potentially compare to IPProtocol & Port children (Same way that abstract TAP does it with kill chains) @@ -266,4 +276,14 @@ class AbstractC2(Application, identifier="AbstractC2"): # Setting the Ports self.current_masquerade_port = payload.masquerade_port self.current_masquerade_protocol = payload.masquerade_protocol - return True \ No newline at end of file + + # This statement is intended to catch on the C2 Application that is listening for connection. (C2 Beacon) + if self.c2_remote_connection == None: + self.sys_log.debug(f"{self.name}: Attempting to configure remote C2 connection based off received output.") + self.c2_remote_connection = self.current_c2_session.with_ip_address + + self.c2_connection_active = True # Sets the connection to active + self.keep_alive_inactivity = 0 # Sets the keep alive inactivity to zeroW + + return True + \ No newline at end of file diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py index 1d61e3b1..14c7af02 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py @@ -1,6 +1,7 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command #from primaite.simulator.system.services.terminal.terminal import Terminal +from prettytable import MARKDOWN, PrettyTable from primaite.simulator.core import RequestManager, RequestType from primaite.interface.request import RequestFormat, RequestResponse from primaite.simulator.network.protocols.masquerade import C2Payload, MasqueradePacket @@ -12,7 +13,7 @@ from enum import Enum from primaite.simulator.system.software import SoftwareHealthState from primaite.simulator.system.applications.application import ApplicationOperatingState -class C2Beacon(AbstractC2, identifier="C2Beacon"): +class C2Beacon(AbstractC2, identifier="C2 Beacon"): """ C2 Beacon Application. @@ -128,11 +129,10 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): def establish(self) -> bool: """Establishes connection to the C2 server via a send alive. Must be called after the C2 Beacon is configured.""" self.run() - self._send_keep_alive() self.num_executions += 1 - + return self._send_keep_alive(session_id=None) - def _handle_command_input(self, payload: MasqueradePacket) -> bool: + def _handle_command_input(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: """ Handles the parsing of C2 Commands from C2 Traffic (Masquerade Packets) as well as then calling the relevant method dependant on the C2 Command. @@ -149,22 +149,22 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): if command == C2Command.RANSOMWARE_CONFIGURE: self.sys_log.info(f"{self.name}: Received a ransomware configuration C2 command.") - return self._return_command_output(self._command_ransomware_config(payload)) + return self._return_command_output(command_output=self._command_ransomware_config(payload), session_id=session_id) elif command == C2Command.RANSOMWARE_LAUNCH: self.sys_log.info(f"{self.name}: Received a ransomware launch C2 command.") - return self._return_command_output(self._command_ransomware_launch(payload)) + return self._return_command_output(command_output=self._command_ransomware_launch(payload), session_id=session_id) elif command == C2Command.TERMINAL: self.sys_log.info(f"{self.name}: Received a terminal C2 command.") - return self._return_command_output(self._command_terminal(payload)) + return self._return_command_output(command_output=self._command_terminal(payload), session_id=session_id) else: self.sys_log.error(f"{self.name}: Received an C2 command: {command} but was unable to resolve command.") return self._return_command_output(RequestResponse(status="failure", data={"Unexpected Behaviour. Unable to resolve command."})) - def _return_command_output(self, command_output: RequestResponse) -> bool: + def _return_command_output(self, command_output: RequestResponse, session_id) -> bool: """Responsible for responding to the C2 Server with the output of the given command.""" output_packet = MasqueradePacket( masquerade_protocol=self.current_masquerade_protocol, @@ -173,11 +173,10 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): payload=command_output ) if self.send( - self, payload=output_packet, dest_ip_address=self.c2_remote_connection, - port=self.current_masquerade_port, - protocol=self.current_masquerade_protocol, + dest_port=self.current_masquerade_port, + ip_protocol=self.current_masquerade_protocol, ): self.sys_log.info(f"{self.name}: Command output sent to {self.c2_remote_connection}") self.sys_log.debug(f"{self.name}: on {self.current_masquerade_port} via {self.current_masquerade_protocol}") @@ -256,15 +255,16 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): if not self._check_c2_connection(timestep): self.sys_log.error(f"{self.name}: Connection Severed - Application Closing.") self.clear_connections() + # TODO: Shouldn't this close() method also set the health state to 'UNUSED'? self.close() return def _check_c2_connection(self, timestep) -> bool: """Checks the C2 Server connection. If a connection cannot be confirmed then this method will return false otherwise true.""" - if self.keep_alive_inactivity > self.keep_alive_frequency: - self.sys_log.info(f"{self.name}: Keep Alive sent to {self.c2_remote_connection} at timestep {timestep}.") - self._send_keep_alive() + if self.keep_alive_inactivity == self.keep_alive_frequency: + self.sys_log.info(f"{self.name}: Attempting to Send Keep Alive to {self.c2_remote_connection} at timestep {timestep}.") + self._send_keep_alive(session_id=self.current_c2_session.uuid) if self.keep_alive_inactivity != 0: self.sys_log.warning(f"{self.name}: Did not receive keep alive from c2 Server. Connection considered severed.") return False @@ -274,6 +274,19 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"): # Defining this abstract method from Abstract C2 def _handle_command_output(self, payload): """C2 Beacons currently do not need to handle output commands coming from the C2 Servers.""" - self.sys_log.warning(f"{self.name}: C2 Beacon received an unexpected OUTPUT payload: {payload}") + self.sys_log.warning(f"{self.name}: C2 Beacon received an unexpected OUTPUT payload: {payload}.") pass - \ No newline at end of file + + def show(self, markdown: bool = False): + """ + Prints a table of the current C2 attributes on a C2 Beacon. + + :param markdown: If True, outputs the table in markdown format. Default is False. + """ + table = PrettyTable(["C2 Connection Active", "C2 Remote Connection", "Keep Alive Inactivity", "Keep Alive Frequency"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"{self.name} Running Status" + table.add_row([self.c2_connection_active, self.c2_remote_connection, self.keep_alive_inactivity, self.keep_alive_frequency]) + print(table) diff --git a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py index 6cff1972..8fe8c00c 100644 --- a/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py +++ b/src/primaite/simulator/system/applications/red_applications/c2/c2_server.py @@ -3,6 +3,7 @@ from primaite.simulator.system.applications.red_applications.c2.abstract_c2 impo from primaite.simulator.network.protocols.masquerade import C2Payload, MasqueradePacket from primaite.simulator.core import RequestManager, RequestType from primaite.interface.request import RequestFormat, RequestResponse +from prettytable import MARKDOWN, PrettyTable from typing import Dict,Optional class C2Server(AbstractC2, identifier="C2 Server"): @@ -94,6 +95,7 @@ class C2Server(AbstractC2, identifier="C2 Server"): return RequestResponse(status="failure", data={"Received unexpected C2 Response."}) return command_output + def _send_command(self, given_command: C2Command, command_options: Dict) -> RequestResponse: """ Sends a command to the C2 Beacon. @@ -115,8 +117,12 @@ class C2Server(AbstractC2, identifier="C2 Server"): command_packet = self._craft_packet(given_command=given_command, command_options=command_options) # Need to investigate if this is correct. - if self.send(self, payload=command_packet,dest_ip_address=self.c2_remote_connection, - port=self.current_masquerade_port, protocol=self.current_masquerade_protocol,session_id=None): + if self.send(payload=command_packet, + dest_ip_address=self.c2_remote_connection, + src_port=self.current_masquerade_port, + dst_port=self.current_masquerade_port, + ip_protocol=self.current_masquerade_protocol, + session_id=None): self.sys_log.info(f"{self.name}: Successfully sent {given_command}.") self.sys_log.info(f"{self.name}: Awaiting command response {given_command}.") return self._handle_command_output(command_packet) @@ -150,3 +156,18 @@ class C2Server(AbstractC2, identifier="C2 Server"): """C2 Servers currently do not receive input commands coming from the C2 Beacons.""" self.sys_log.warning(f"{self.name}: C2 Server received an unexpected INPUT payload: {payload}") pass + + + def show(self, markdown: bool = False): + """ + Prints a table of the current C2 attributes on a C2 Server. + + :param markdown: If True, outputs the table in markdown format. Default is False. + """ + table = PrettyTable(["C2 Connection Active", "C2 Remote Connection", "Keep Alive Inactivity"]) + if markdown: + table.set_style(MARKDOWN) + table.align = "l" + table.title = f"{self.name} Running Status" + table.add_row([self.c2_connection_active, self.c2_remote_connection, self.keep_alive_inactivity]) + print(table) diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py index 20da4140..3014cd19 100644 --- a/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py +++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py @@ -48,19 +48,25 @@ def basic_network() -> Network: node_a = Computer(hostname="node_a", ip_address="192.168.0.10", subnet_mask="255.255.255.0", start_up_duration=0) node_a.power_on() node_a.software_manager.get_open_ports() - + node_a.software_manager.install(software_class=C2Server) node_b = Computer(hostname="node_b", ip_address="192.168.0.11", subnet_mask="255.255.255.0", start_up_duration=0) + node_b.software_manager.install(software_class=C2Beacon) node_b.power_on() network.connect(node_a.network_interface[1], node_b.network_interface[1]) return network def test_c2_suite_setup_receive(basic_network): - """Test that C2 Beacon can successfully establish connection with the c2 Server""" + """Test that C2 Beacon can successfully establish connection with the c2 Server.""" network: Network = basic_network computer_a: Computer = network.get_node_by_hostname("node_a") c2_server: C2Server = computer_a.software_manager.software.get("C2Server") computer_b: Computer = network.get_node_by_hostname("node_b") - c2_beacon: C2Server = computer_a.software_manager.software.get("C2Beacon") + c2_beacon: C2Beacon = computer_b.software_manager.software.get("C2Beacon") + + c2_beacon.configure(c2_server_ip_address="192.168.0.10") + c2_beacon.establish() + + c2_beacon.sys_log.show() \ No newline at end of file