diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 8cddbcda..5ef8c14c 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -36,8 +36,6 @@ from primaite.simulator.system.applications.red_applications.data_manipulation_b ) from primaite.simulator.system.applications.red_applications.dos_bot import DoSBot # noqa: F401 from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript # noqa: F401 -from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon -from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server from primaite.simulator.system.applications.web_browser import WebBrowser # noqa: F401 from primaite.simulator.system.services.database.database_service import DatabaseService from primaite.simulator.system.services.dns.dns_client import DNSClient diff --git a/src/primaite/simulator/network/protocols/masquerade.py b/src/primaite/simulator/network/protocols/masquerade.py index 93554f57..7ef17fc0 100644 --- a/src/primaite/simulator/network/protocols/masquerade.py +++ b/src/primaite/simulator/network/protocols/masquerade.py @@ -5,19 +5,6 @@ from typing import Optional from primaite.simulator.network.protocols.packet import DataPacket -class C2Payload(Enum): - """Represents the different types of command and control payloads.""" - - KEEP_ALIVE = "keep_alive" - """C2 Keep Alive payload. Used by the C2 beacon and C2 Server to confirm their connection.""" - - INPUT = "input_command" - """C2 Input Command payload. Used by the C2 Server to send a command to the c2 beacon.""" - - OUTPUT = "output_command" - """C2 Output Command. Used by the C2 Beacon to send the results of a Input command to the c2 server.""" - - class MasqueradePacket(DataPacket): """Represents an generic malicious packet that is masquerading as another protocol.""" @@ -25,6 +12,6 @@ class MasqueradePacket(DataPacket): masquerade_port: Enum # The 'Masquerade' port that is currently in use - payload_type: C2Payload # The type of C2 traffic (e.g keep alive, command or command out) + payload_type: Enum # The type of C2 traffic (e.g keep alive, command or command out) - command: Optional[str] # Used to pass the actual C2 Command in C2 INPUT + command: Optional[Enum] = None # Used to pass the actual C2 Command in C2 INPUT 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 9c840616..af701e8c 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 @@ -6,7 +6,7 @@ from typing import Dict, Optional from pydantic import validate_call -from primaite.simulator.network.protocols.masquerade import C2Payload, MasqueradePacket +from primaite.simulator.network.protocols.masquerade import MasqueradePacket 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 @@ -20,9 +20,7 @@ from primaite.simulator.system.core.session_manager import Session class C2Command(Enum): - """ - Enumerations representing the different commands the C2 suite currently supports. - """ + """Enumerations representing the different commands the C2 suite currently supports.""" RANSOMWARE_CONFIGURE = "Ransomware Configure" "Instructs the c2 beacon to configure the ransomware with the provided options." @@ -36,12 +34,25 @@ class C2Command(Enum): # The terminal command should also be able to pass a session which can be used for remote connections. +class C2Payload(Enum): + """Represents the different types of command and control payloads.""" + + KEEP_ALIVE = "keep_alive" + """C2 Keep Alive payload. Used by the C2 beacon and C2 Server to confirm their connection.""" + + INPUT = "input_command" + """C2 Input Command payload. Used by the C2 Server to send a command to the c2 beacon.""" + + OUTPUT = "output_command" + """C2 Output Command. Used by the C2 Beacon to send the results of a Input command to the c2 server.""" + + class AbstractC2(Application, identifier="AbstractC2"): """ An abstract command and control (c2) application. Extends the application class to provide base functionality for c2 suite applications - such as c2 beacons and c2 servers. + such as c2 beacons and c2 servers. Provides the base methods for handling ``Keep Alive`` connections, configuring masquerade ports and protocols as well as providing the abstract methods for sending, receiving and parsing commands. @@ -55,13 +66,6 @@ class AbstractC2(Application, identifier="AbstractC2"): c2_remote_connection: IPv4Address = None """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.""" - - # We should set the application to NOT_RUNNING if the inactivity count reaches a certain thresh hold. - keep_alive_inactivity: int = 0 - """Indicates how many timesteps since the last time the c2 application received a keep alive.""" - # These two attributes are set differently in the c2 server and c2 beacon. # The c2 server parses the keep alive and sets these accordingly. # The c2 beacon will set this attributes upon installation and configuration @@ -75,9 +79,6 @@ class AbstractC2(Application, identifier="AbstractC2"): 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: """ @@ -104,10 +105,10 @@ class AbstractC2(Application, identifier="AbstractC2"): :rtype: Dict """ return super().describe_state() - + def __init__(self, **kwargs): - kwargs["port"] = Port.HTTP # TODO: Update this post application/services requiring to listen to multiple ports - kwargs["protocol"] = IPProtocol.TCP # Update this as well + 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. @@ -145,7 +146,7 @@ class AbstractC2(Application, identifier="AbstractC2"): elif payload.payload_type == C2Payload.OUTPUT: self.sys_log.info(f"{self.name} received an OUTPUT COMMAND payload.") - return self._handle_command_input(payload, session_id) + return self._handle_command_output(payload) else: self.sys_log.warning( @@ -168,41 +169,7 @@ class AbstractC2(Application, identifier="AbstractC2"): pass def _handle_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: - """ - Handles receiving and sending keep alive payloads. This method is only called if we receive a keep alive. - - Returns False if a keep alive was unable to be sent. - Returns True if a keep alive was successfully sent or already has been sent this timestep. - - :return: True if successfully handled, false otherwise. - :rtype: Bool - """ - 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 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. - - self.sys_log.info(f"{self.name}: Connection successfully established with {self.c2_remote_connection}.") - - return self._send_keep_alive(session_id) - + """Abstract Method: The C2 Server and the C2 Beacon handle the KEEP ALIVEs differently.""" # from_network_interface=from_network_interface def receive(self, payload: MasqueradePacket, session_id: Optional[str] = None, **kwargs) -> bool: @@ -217,23 +184,23 @@ class AbstractC2(Application, identifier="AbstractC2"): """Sends a C2 keep alive payload to the self.remote_connection IPv4 Address.""" # Checking that the c2 application is capable of performing both actions and has an enabled NIC # (Using NOT to improve code readability) - if self.c2_remote_connection == None: - self.sys_log.error(f"{self.name}: Unable to Establish connection as the C2 Server's IP Address has not been given.") - + if self.c2_remote_connection is None: + self.sys_log.error( + f"{self.name}: Unable to Establish connection as the C2 Server's IP Address has not been given." + ) + if not self._can_perform_network_action(): self.sys_log.warning(f"{self.name}: Unable to perform network actions.") return False - # We also Pass masquerade protocol/port so that the c2 server can reply on the correct protocol/port. + # We also Pass masquerade proto`col/port so that the c2 server can reply on the correct protocol/port. # (This also lays the foundations for switching masquerade port/protocols mid episode.) keep_alive_packet = MasqueradePacket( masquerade_protocol=self.current_masquerade_protocol, masquerade_port=self.current_masquerade_port, payload_type=C2Payload.KEEP_ALIVE, - command=None + command=None, ) - # 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( payload=keep_alive_packet, @@ -242,16 +209,16 @@ class AbstractC2(Application, identifier="AbstractC2"): ip_protocol=self.current_masquerade_protocol, session_id=session_id, ): + self.keep_alive_sent = True 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}") return True else: self.sys_log.warning( - f"{self.name}: failed to send a Keep Alive. The node may be unable to access the network." + f"{self.name}: failed to send a Keep Alive. The node may be unable to access the ``network." ) return False - def _resolve_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: """ Parses the Masquerade Port/Protocol within the received Keep Alive packet. @@ -260,7 +227,7 @@ class AbstractC2(Application, identifier="AbstractC2"): Returns True on successfully extracting and configuring the masquerade port/protocols. Returns False otherwise. - + :param payload: The Keep Alive payload received. :type payload: MasqueradePacket :return: True on successful configuration, false otherwise. @@ -268,22 +235,24 @@ class AbstractC2(Application, identifier="AbstractC2"): """ # Validating that they are valid Enums. 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)}.") + self.sys_log.warning( + f"{self.name}: Received invalid Masquerade Values within Keep Alive." + f"Port: {payload.masquerade_port} Protocol: {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) - + # Setting the Ports self.current_masquerade_port = payload.masquerade_port self.current_masquerade_protocol = payload.masquerade_protocol # This statement is intended to catch on the C2 Application that is listening for connection. (C2 Beacon) - if self.c2_remote_connection == None: + if self.c2_remote_connection is 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 14c7af02..b00b7c57 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,34 +1,44 @@ # © 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 -from primaite.simulator.network.transmission.network_layer import IPProtocol -from ipaddress import IPv4Address -from typing import Dict,Optional -from primaite.simulator.network.transmission.transport_layer import Port from enum import Enum -from primaite.simulator.system.software import SoftwareHealthState +from ipaddress import IPv4Address +from typing import Dict, Optional + +# from primaite.simulator.system.services.terminal.terminal import Terminal +from prettytable import MARKDOWN, PrettyTable + +from primaite.interface.request import RequestFormat, RequestResponse +from primaite.simulator.core import RequestManager, RequestType +from primaite.simulator.network.protocols.masquerade import MasqueradePacket +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 ApplicationOperatingState +from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload +from primaite.simulator.system.software import SoftwareHealthState + class C2Beacon(AbstractC2, identifier="C2 Beacon"): """ C2 Beacon Application. - Represents a generic C2 beacon which can be used in conjunction with the C2 Server - to simulate malicious communications within primAITE. + Represents a vendor generic C2 beacon is used in conjunction with the C2 Server + to simulate malicious communications and infrastructure within primAITE. + + Must be configured with the C2 Server's IP Address upon installation. - Must be configured with the C2 Server's Ip Address upon installation. - Extends the Abstract C2 application to include the following: 1. Receiving commands from the C2 Server (Command input) 2. Leveraging the terminal application to execute requests (dependant on the command given) 3. Sending the RequestResponse back to the C2 Server (Command output) """ - + + keep_alive_attempted: bool = False + """Indicates if a keep alive has been attempted to be sent this timestep. Used to prevent packet storms.""" + + # We should set the application to NOT_RUNNING if the inactivity count reaches a certain thresh hold. + keep_alive_inactivity: int = 0 + """Indicates how many timesteps since the last time the c2 application received a keep alive.""" + keep_alive_frequency: int = 5 "The frequency at which ``Keep Alive`` packets are sent to the C2 Server from the C2 Beacon." @@ -38,8 +48,8 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): # Implement a command output method that sends the RequestResponse to the C2 server. # Uncomment the terminal Import and the terminal property after terminal PR - #@property - #def _host_terminal(self) -> Terminal: + # @property + # def _host_terminal(self) -> Terminal: # """Return the Terminal that is installed on the same machine as the C2 Beacon.""" # host_terminal: Terminal = self.software_manager.software.get("Terminal") # if host_terminal: is None: @@ -69,27 +79,33 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): :return: RequestResponse object with a success code reflecting whether the configuration could be applied. :rtype: RequestResponse """ - server_ip = request[-1].get("c2_server_ip_address") - if server_ip == None: + c2_remote_ip = request[-1].get("c2_server_ip_address") + if c2_remote_ip is None: self.sys_log.error(f"{self.name}: Did not receive C2 Server IP in configuration parameters.") - RequestResponse(status="failure", data={"No C2 Server IP given to C2 beacon. Unable to configure C2 Beacon"}) + RequestResponse( + status="failure", data={"No C2 Server IP given to C2 beacon. Unable to configure C2 Beacon"} + ) c2_remote_ip = IPv4Address(c2_remote_ip) frequency = request[-1].get("keep_alive_frequency") - protocol= request[-1].get("masquerade_protocol") + protocol = request[-1].get("masquerade_protocol") port = request[-1].get("masquerade_port") - return RequestResponse.from_bool(self.configure(c2_server_ip_address=server_ip, - keep_alive_frequency=frequency, - masquerade_protocol=protocol, - masquerade_port=port)) + return RequestResponse.from_bool( + self.configure( + c2_server_ip_address=c2_remote_ip, + keep_alive_frequency=frequency, + masquerade_protocol=protocol, + masquerade_port=port, + ) + ) rm.add_request("configure", request_type=RequestType(func=_configure)) return rm - + def __init__(self, **kwargs): kwargs["name"] = "C2Beacon" super().__init__(**kwargs) - + def configure( self, c2_server_ip_address: IPv4Address = None, @@ -100,6 +116,7 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): """ Configures the C2 beacon to communicate with the C2 server with following additional parameters. + # TODO: Expand docustring. :param c2_server_ip_address: The IP Address of the C2 Server. Used to establish connection. :type c2_server_ip_address: IPv4Address @@ -117,43 +134,70 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): self.sys_log.info( f"{self.name}: Configured {self.name} with remote C2 server connection: {c2_server_ip_address=}." ) - self.sys_log.debug(f"{self.name}: configured with the following settings:" - f"Remote C2 Server: {c2_server_ip_address}" - f"Keep Alive Frequency {keep_alive_frequency}" - f"Masquerade Protocol: {masquerade_protocol}" - f"Masquerade Port: {masquerade_port}") + self.sys_log.debug( + f"{self.name}: configured with the following settings:" + f"Remote C2 Server: {c2_server_ip_address}" + f"Keep Alive Frequency {keep_alive_frequency}" + f"Masquerade Protocol: {masquerade_protocol}" + f"Masquerade Port: {masquerade_port}" + ) return True - # I THINK that once the application is running it can respond to incoming traffic but I'll need to test this later. def establish(self) -> bool: - """Establishes connection to the C2 server via a send alive. Must be called after the C2 Beacon is configured.""" + """Establishes connection to the C2 server via a send alive. The C2 Beacon must already be configured.""" + if self.c2_remote_connection is None: + self.sys_log.info(f"{self.name}: Failed to establish connection. C2 Beacon has not been configured.") + return False self.run() self.num_executions += 1 return self._send_keep_alive(session_id=None) 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. - - :param payload: The INPUT C2 Payload + Handles the parsing of C2 Commands from C2 Traffic (Masquerade Packets). + + Dependant the C2 Command contained within the payload. + The following methods are called and returned. + + C2 Command | Internal Method + ---------------------|------------------------ + RANSOMWARE_CONFIGURE | self._command_ransomware_config() + RANSOMWARE_LAUNCH | self._command_ransomware_launch() + Terminal | self._command_terminal() + + Please see each method individually for further information regarding + the implementation of these commands. + + :param payload: The INPUT C2 Payload :type payload: MasqueradePacket :return: The Request Response provided by the terminal execute method. :rtype Request Response: """ - command = payload.payload_type - if command != C2Payload: + # TODO: Probably could refactor this to be a more clean. + # The elif's are a bit ugly when they are all calling the same method. + command = payload.command + if not isinstance(command, C2Command): self.sys_log.warning(f"{self.name}: Received unexpected C2 command. Unable to resolve command") - return self._return_command_output(RequestResponse(status="failure", data={"Received unexpected C2Command. Unable to resolve command."})) + return self._return_command_output( + command_output=RequestResponse( + status="failure", + data={"Reason": "C2 Beacon received unexpected C2Command. Unable to resolve command."}, + ), + session_id=session_id, + ) if command == C2Command.RANSOMWARE_CONFIGURE: self.sys_log.info(f"{self.name}: Received a ransomware configuration C2 command.") - return self._return_command_output(command_output=self._command_ransomware_config(payload), session_id=session_id) + 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(command_output=self._command_ransomware_launch(payload), session_id=session_id) + 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.") @@ -161,22 +205,30 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): 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."})) + return self._return_command_output( + RequestResponse(status="failure", data={"Reason": "Unexpected Behaviour. Unable to resolve command."}) + ) + def _return_command_output(self, command_output: RequestResponse, session_id: Optional[str] = None) -> bool: + """Responsible for responding to the C2 Server with the output of the given command. - 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.""" + :param command_output: The RequestResponse returned by the terminal application's execute method. + :type command_output: Request Response + :param session_id: The current session established with the C2 Server. + :type session_id: Str + """ output_packet = MasqueradePacket( masquerade_protocol=self.current_masquerade_protocol, masquerade_port=self.current_masquerade_port, payload_type=C2Payload.OUTPUT, - payload=command_output + payload=command_output, ) if self.send( payload=output_packet, dest_ip_address=self.c2_remote_connection, dest_port=self.current_masquerade_port, ip_protocol=self.current_masquerade_protocol, + session_id=session_id, ): 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}") @@ -189,68 +241,121 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): def _command_ransomware_config(self, payload: MasqueradePacket) -> RequestResponse: """ - C2 Command: Ransomware Configuration + C2 Command: Ransomware Configuration. Creates a request that configures the ransomware based off the configuration options given. This request is then sent to the terminal service in order to be executed. + :payload MasqueradePacket: The incoming INPUT command. + :type Masquerade Packet: MasqueradePacket. :return: Returns the Request Response returned by the Terminal execute method. :rtype: Request Response """ - pass - #return self._host_terminal.execute(command) + # TODO: replace and use terminal + return RequestResponse(status="success", data={"Reason": "Placeholder."}) def _command_ransomware_launch(self, payload: MasqueradePacket) -> RequestResponse: """ - C2 Command: Ransomware Execute + C2 Command: Ransomware Launch. Creates a request that executes the ransomware script. This request is then sent to the terminal service in order to be executed. + + :payload MasqueradePacket: The incoming INPUT command. + :type Masquerade Packet: MasqueradePacket. :return: Returns the Request Response returned by the Terminal execute method. :rtype: Request Response - - Creates a Request that launches the ransomware. """ - pass - #return self._host_terminal.execute(command) + # TODO: replace and use terminal + return RequestResponse(status="success", data={"Reason": "Placeholder."}) def _command_terminal(self, payload: MasqueradePacket) -> RequestResponse: """ - C2 Command: Ransomware Execute + C2 Command: Terminal. - Creates a request that executes the ransomware script. + Creates a request that executes a terminal command. This request is then sent to the terminal service in order to be executed. + :payload MasqueradePacket: The incoming INPUT command. + :type Masquerade Packet: MasqueradePacket. :return: Returns the Request Response returned by the Terminal execute method. :rtype: Request Response - - Creates a Request that launches the ransomware. """ - pass - #return self._host_terminal.execute(command) + # TODO: uncomment and replace (uses terminal) + return RequestResponse(status="success", data={"Reason": "Placeholder."}) + # return self._host_terminal.execute(command) + def _handle_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: + """ + Handles receiving and sending keep alive payloads. This method is only called if we receive a keep alive. + + In the C2 Beacon implementation of this method the c2 connection active boolean + is set to true and the keep alive inactivity is reset only after sending a keep alive + as wel as receiving a response back from the C2 Server. + + This is because the C2 Server is the listener and thus will only ever receive packets from + the C2 Beacon rather than the other way around. (The C2 Beacon is akin to a reverse shell) + + Therefore, we need a response back from the listener (C2 Server) + before the C2 beacon is able to confirm it's connection. + + Returns False if a keep alive was unable to be sent. + Returns True if a keep alive was successfully sent or already has been sent this timestep. + + :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. + # This guard clause triggers on the c2 suite that establishes connection. + if self.keep_alive_attempted is True: + 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] + + # We set keep alive_attempted here to show that we've achieved connection. + self.keep_alive_attempted = False + self.sys_log.warning(f"{self.name}: Connection successfully Established with C2 Server.") + 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 self._resolve_keep_alive(payload, session_id) is False: + self.sys_log.warning(f"{self.name}: Keep Alive Could not be resolved correctly. Refusing Keep Alive.") + return False + + self.keep_alive_attempted = True + # If this method returns true then we have sent successfully sent a keep alive. + return self._send_keep_alive(session_id) - # Not entirely sure if this actually works. def apply_timestep(self, timestep: int) -> None: - """ - Apply a timestep to the c2_beacon. - Used to keep track of when the c2 beacon should send another keep alive. + """Apply a timestep to the c2_beacon. + Used to keep track of when the c2 beacon should send another keep alive. The following logic is applied: 1. Each timestep the keep_alive_inactivity is increased. 2. If the keep alive inactivity eclipses that of the keep alive frequency then another keep alive is sent. - 3. If the c2 beacon receives a keep alive response packet then the ``keep_alive_inactivity`` attribute is set to 0 - - Therefore, if ``keep_alive_inactivity`` attribute is not 0, then the connection is considered severed and c2 beacon will shut down. - + 3. If a keep alive response packet is received then the ``keep_alive_inactivity`` attribute is reset. + + Therefore, if ``keep_alive_inactivity`` attribute is not 0 after a keep alive is sent + then the connection is considered severed and c2 beacon will shut down. + :param timestep: The current timestep of the simulation. + :type timestep: Int + :return bool: Returns false if connection was lost. Returns True if connection is active or re-established. + :rtype bool: """ super().apply_timestep(timestep=timestep) - if self.operating_state is ApplicationOperatingState.RUNNING and self.health_state_actual is SoftwareHealthState.GOOD: + self.keep_alive_attempted = False # Resetting keep alive sent. + if ( + self.operating_state is ApplicationOperatingState.RUNNING + and self.health_state_actual is SoftwareHealthState.GOOD + ): self.keep_alive_inactivity += 1 if not self._check_c2_connection(timestep): self.sys_log.error(f"{self.name}: Connection Severed - Application Closing.") @@ -259,34 +364,67 @@ class C2Beacon(AbstractC2, identifier="C2 Beacon"): self.close() return + def _check_c2_connection(self, timestep: int) -> bool: + """Checks the suitability of the current C2 Server connection. - 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 a connection cannot be confirmed then this method will return false otherwise true. + + :param timestep: The current timestep of the simulation. + :type timestep: Int + :return: Returns False if connection was lost. Returns True if connection is active or re-established. + :rtype bool: + """ 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.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.") + self.sys_log.warning( + f"{self.name}: Did not receive keep alive from c2 Server. Connection considered severed." + ) return False return True - # 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.""" + def _handle_command_output(self, payload: MasqueradePacket): + """C2 Beacons currently does 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}.") pass - + def show(self, markdown: bool = False): """ - Prints a table of the current C2 attributes on a C2 Beacon. + Prints a table of the current status of the C2 Beacon. + + Displays the current values of the following C2 attributes: + + ``C2 Connection Active``: + If the C2 Beacon is currently connected to the C2 Server + + ``C2 Remote Connection``: + The IP of the C2 Server. (Configured by upon installation) + + ``Keep Alive Inactivity``: + How many timesteps have occurred since the last keep alive. + + ``Keep Alive Frequency``: + How often should the C2 Beacon attempt a keep alive? :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"]) + 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]) + 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 8fe8c00c..9d02224e 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 @@ -1,15 +1,34 @@ # © 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.network.protocols.masquerade import C2Payload, MasqueradePacket -from primaite.simulator.core import RequestManager, RequestType -from primaite.interface.request import RequestFormat, RequestResponse +from typing import Dict, Optional + from prettytable import MARKDOWN, PrettyTable -from typing import Dict,Optional +from pydantic import validate_call + +from primaite.interface.request import RequestFormat, RequestResponse +from primaite.simulator.core import RequestManager, RequestType +from primaite.simulator.network.protocols.masquerade import MasqueradePacket +from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload + class C2Server(AbstractC2, identifier="C2 Server"): - # TODO: - # Implement the request manager and agent actions. - # Implement the output handling methods. (These need to interface with the actions) + """ + C2 Server Application. + + Represents a vendor generic C2 Server is used in conjunction with the C2 beacon + to simulate malicious communications and infrastructure within primAITE. + + The C2 Server must be installed and be in a running state before it's able to receive + red agent actions and send commands to the C2 beacon. + + Extends the Abstract C2 application to include the following: + + 1. Sending commands to the C2 Beacon. (Command input) + 2. Parsing terminal RequestResponses back to the Agent. + """ + + current_command_output: RequestResponse = None + """The Request Response by the last command send. This attribute is updated by the method _handle_command_output.""" + def _init_request_manager(self) -> RequestManager: """ Initialise the request manager. @@ -34,7 +53,7 @@ class C2Server(AbstractC2, identifier="C2 Server"): def _launch_ransomware_action(request: RequestFormat, context: Dict) -> RequestResponse: """Agent Action - Sends a RANSOMWARE_LAUNCH C2Command to the C2 Beacon with the given parameters. - + :param request: Request with one element containing a dict of parameters for the configure method. :type request: RequestFormat :param context: additional context for resolving this action, currently unused @@ -48,7 +67,7 @@ class C2Server(AbstractC2, identifier="C2 Server"): def _remote_terminal_action(request: RequestFormat, context: Dict) -> RequestResponse: """Agent Action - Sends a TERMINAL C2Command to the C2 Beacon with the given parameters. - + :param request: Request with one element containing a dict of parameters for the configure method. :type request: RequestFormat :param context: additional context for resolving this action, currently unused @@ -78,28 +97,69 @@ class C2Server(AbstractC2, identifier="C2 Server"): kwargs["name"] = "C2Server" super().__init__(**kwargs) - def _handle_command_output(self, payload: MasqueradePacket) -> RequestResponse: + def _handle_command_output(self, payload: MasqueradePacket) -> bool: """ - Handles the parsing of C2 Command Output from C2 Traffic (Masquerade Packets) - as well as then calling the relevant method dependant on the C2 Command. - - :param payload: The OUTPUT C2 Payload + Handles the parsing of C2 Command Output from C2 Traffic (Masquerade Packets). + + Parses the Request Response within MasqueradePacket's payload attribute (Inherited from Data packet). + The class attribute self.current_command_output is then set to this Request Response. + + If the payload attribute does not contain a RequestResponse, then an error will be raised in syslog and + the self.current_command_output is updated to reflect the error. + + :param payload: The OUTPUT C2 Payload :type payload: MasqueradePacket - :return: Returns the Request Response of the C2 Beacon's host terminal service execute method. - :rtype Request Response: + :return: Returns True if the self.current_command_output is currently updated, false otherwise. + :rtype Bool: """ self.sys_log.info(f"{self.name}: Received command response from C2 Beacon: {payload}.") command_output = payload.payload - if command_output != MasqueradePacket: - self.sys_log.warning(f"{self.name}: Received invalid command response: {command_output}.") - return RequestResponse(status="failure", data={"Received unexpected C2 Response."}) - return command_output - + if not isinstance(command_output, RequestResponse): + self.sys_log.warning(f"{self.name}: C2 Server received invalid command response: {command_output}.") + self.current_command_output = RequestResponse( + status="failure", data={"Reason": "Received unexpected C2 Response."} + ) + return False + self.current_command_output = command_output + return True + + def _handle_keep_alive(self, payload: MasqueradePacket, session_id: Optional[str]) -> bool: + """ + Handles receiving and sending keep alive payloads. This method is only called if we receive a keep alive. + + In the C2 Server implementation of this method the c2 connection active boolean + is set to true and the keep alive inactivity is reset after receiving one keep alive. + + This is because the C2 Server is the listener and thus will only ever receive packets from + the C2 Beacon rather than the other way around. (The C2 Beacon is akin to a reverse shell) + + Returns False if a keep alive was unable to be sent. + Returns True if a keep alive was successfully sent or already has been sent this timestep. + + :return: True if successfully handled, false otherwise. + :rtype: Bool + """ + self.sys_log.info(f"{self.name}: Keep Alive Received. Attempting to resolve the remote connection details.") + + self.c2_connection_active = True # Sets the connection to active + self.current_c2_session = self.software_manager.session_manager.sessions_by_uuid[session_id] + + 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. + self.sys_log.info(f"{self.name}: Remote connection successfully established: {self.c2_remote_connection}.") + self.sys_log.debug(f"{self.name}: Attempting to send Keep Alive response back to {self.c2_remote_connection}.") + + return self._send_keep_alive(session_id) + + @validate_call def _send_command(self, given_command: C2Command, command_options: Dict) -> RequestResponse: """ Sends a command to the C2 Beacon. - + # TODO: Expand this docustring. :param given_command: The C2 command to be sent to the C2 Beacon. @@ -109,31 +169,41 @@ class C2Server(AbstractC2, identifier="C2 Server"): :return: Returns the Request Response of the C2 Beacon's host terminal service execute method. :rtype: RequestResponse """ - if given_command != C2Payload: + if not isinstance(given_command, C2Command): self.sys_log.warning(f"{self.name}: Received unexpected C2 command. Unable to send command.") - return RequestResponse(status="failure", data={"Received unexpected C2Command. Unable to send command."}) + return RequestResponse( + status="failure", data={"Reason": "Received unexpected C2Command. Unable to send command."} + ) + + if self._can_perform_network_action == False: + self.sys_log.warning(f"{self.name}: Unable to make leverage networking resources. Rejecting Command.") + return RequestResponse( + status="failure", data={"Reason": "Unable to access networking resources. Unable to send command."} + ) self.sys_log.info(f"{self.name}: Attempting to send command {given_command}.") command_packet = self._craft_packet(given_command=given_command, command_options=command_options) # Need to investigate if this is correct. - if self.send(payload=command_packet, + 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, + session_id=self.current_c2_session.uuid, + dest_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) + # If the command output was handled currently, the self.current_command_output will contain the RequestResponse. + return self.current_command_output # TODO: Perhaps make a new pydantic base model for command_options? # TODO: Perhaps make the return optional? Returns False is the packet was unable to be crafted. def _craft_packet(self, given_command: C2Command, command_options: Dict) -> MasqueradePacket: """ Creates a Masquerade Packet based off the command parameter and the arguments given. - + :param given_command: The C2 command to be sent to the C2 Beacon. :type given_command: C2Command. :param command_options: The relevant C2 Beacon parameters.F @@ -147,27 +217,33 @@ class C2Server(AbstractC2, identifier="C2 Server"): masquerade_port=self.current_masquerade_port, payload_type=C2Payload.INPUT, command=given_command, - payload=command_options + payload=command_options, ) return constructed_packet - + + # TODO: I think I can just overload the methods rather than setting it as abstract_method? # Defining this abstract method - def _handle_command_input(self, payload): - """C2 Servers currently do not receive input commands coming from the C2 Beacons.""" + def _handle_command_input(self, payload: MasqueradePacket): + """Defining this method (Abstract method inherited from abstract C2) in order to instantiate the class. + + C2 Servers currently do not receive input commands coming from the C2 Beacons. + + :param payload: The incoming MasqueradePacket + :type payload: MasqueradePacket. + """ 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"]) + table = PrettyTable(["C2 Connection Active", "C2 Remote Connection"]) 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]) + table.add_row([self.c2_connection_active, self.c2_remote_connection]) 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 3014cd19..7f869e92 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 @@ -12,14 +12,15 @@ from primaite.simulator.network.hardware.nodes.network.router import ACLAction, from primaite.simulator.network.hardware.nodes.network.switch import Switch from primaite.simulator.network.transmission.network_layer import IPProtocol from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon +from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server from primaite.simulator.system.services.dns.dns_server import DNSServer from primaite.simulator.system.services.service import ServiceOperatingState from primaite.simulator.system.services.web_server.web_server import WebServer -from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon -from primaite.simulator.system.applications.red_applications.c2.c2_server import C2Server # TODO: Update these tests. + @pytest.fixture(scope="function") def c2_server_on_computer() -> Tuple[C2Beacon, Computer]: computer: Computer = Computer( @@ -30,6 +31,7 @@ def c2_server_on_computer() -> Tuple[C2Beacon, Computer]: return [c2_beacon, computer] + @pytest.fixture(scope="function") def c2_server_on_computer() -> Tuple[C2Server, Computer]: computer: Computer = Computer( @@ -41,7 +43,6 @@ def c2_server_on_computer() -> Tuple[C2Server, Computer]: return [c2_server, computer] - @pytest.fixture(scope="function") def basic_network() -> Network: network = Network() @@ -57,6 +58,7 @@ def basic_network() -> Network: return network + def test_c2_suite_setup_receive(basic_network): """Test that C2 Beacon can successfully establish connection with the c2 Server.""" network: Network = basic_network @@ -68,5 +70,5 @@ def test_c2_suite_setup_receive(basic_network): 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 + + c2_beacon.sys_log.show()