#2689 Remote connections now successfully establishing however current issues with keep alive inactivity causing the c2 beacon to close even when it does have connection to the c2 server.

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["port"] = Port.HTTP # TODO: Update this post application/services requiring to listen to multiple ports
-        kwargs["protocol"] = IPProtocol.NONE
+        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}")
+        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.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.
+        self.sys_log.info(f"{self.name}: Connection successfully established with {self.c2_remote_connection}.")
-        else:
-            return False
                 
-    def receive(self, payload: MasqueradePacket, session_id: Optional[str] = None) -> bool:
+        return self._send_keep_alive(session_id)
+
+
+    # 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
         )
-
+        # We need to set this guard clause to true before sending the keep alive (prevents packet storms.)
-        # C2 Server will need to c2_remote_connection after it receives it's first keep alive.
+        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,
+            dest_port=self.current_masquerade_port,
-            protocol=self.current_masquerade_protocol,
+            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:
+        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 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
+
+        # 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
+    

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
@@ -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, session_id: Optional[str]) -> bool:
-    def _handle_command_input(self, payload: MasqueradePacket) -> 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,
+            dest_port=self.current_masquerade_port,
-            protocol=self.current_masquerade_protocol,
+            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:
+        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.sys_log.info(f"{self.name}: Attempting to Send Keep Alive to {self.c2_remote_connection} at timestep {timestep}.")
-            self._send_keep_alive()
+            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
     
+    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)

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,
+        if self.send(payload=command_packet,            
-            port=self.current_masquerade_port, protocol=self.current_masquerade_protocol,session_id=None):
+            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)

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()