#2689 Final docustring updates before PR.

src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py

@@ -87,7 +87,10 @@ class AbstractC2(Application, identifier="AbstractC2"):
         Creates and returns a Masquerade Packet using the parameters given.
 
         The packet uses the current c2 configuration and parameters given
-        to construct a C2 Packet.
+        to construct the base networking information such as the masquerade
+        protocol/port. Additionally all C2 Traffic packets pass the currently
+        in use C2 configuration. This ensures that the all C2 applications
+        can keep their configuration in sync.
 
         :param c2_payload: The type of C2 Traffic ot be sent
         :type c2_payload: C2Payload
@@ -208,6 +211,8 @@ class AbstractC2(Application, identifier="AbstractC2"):
         :type payload: C2Packet
         :param session_id: The transport session_id that the payload is originating from.
         :type session_id: str
+        :return: Returns a bool if the traffic was received correctly (See _handle_c2_payload.)
+        :rtype: bool
         """
         return self._handle_c2_payload(payload, session_id)
 
@@ -306,7 +311,13 @@ class AbstractC2(Application, identifier="AbstractC2"):
         return True
 
     def _reset_c2_connection(self) -> None:
-        """Resets all currently established C2 communications to their default setting."""
+        """
+        Resets all currently established C2 communications to their default setting.
+
+        This method is called once a C2 application considers their remote connection
+        severed and reverts back to default settings. Worth noting that that this will
+        revert any non-default configuration that a user/agent may have set.
+        """
         self.c2_connection_active = False
         self.c2_session = None
         self.keep_alive_inactivity = 0
@@ -359,7 +370,9 @@ class AbstractC2(Application, identifier="AbstractC2"):
         Validation method: Checks that the C2 application is capable of sending C2 Command input/output.
 
         Performs a series of connection validation to ensure that the C2 application is capable of
-        sending and responding to the remote c2 connection.
+        sending and responding to the remote c2 connection. This method is used to confirm connection
+        before carrying out Agent Commands hence why this method also returns a tuple
+        containing both a success boolean as well as RequestResponse.
 
         :return: A tuple containing a boolean True/False and a corresponding Request Response
         :rtype: tuple[bool, RequestResponse]

src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py

@@ -156,6 +156,11 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
         masquerade_protocol  | What protocol should the C2 traffic masquerade as? (HTTP, FTP or DNS)
         masquerade_port      | What port should the C2 traffic use? (TCP or UDP)
 
+        These configuration options are used to reassign the fields in the inherited inner class
+        ``c2_config``.
+
+        If a connection is already in progress then this method also sends a keep alive to the C2
+        Server in order for the C2 Server to sync with the new configuration settings.
 
         :param c2_server_ip_address: The IP Address of the C2 Server. Used to establish connection.
         :type c2_server_ip_address: IPv4Address
@@ -165,6 +170,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
         :type masquerade_protocol: Enum (IPProtocol)
         :param masquerade_port: The Port that the C2 Traffic will masquerade as. Defaults to FTP.
         :type masquerade_port: Enum (Port)
+        :return: Returns True if the configuration was successful, False otherwise.
         """
         self.c2_remote_connection = IPv4Address(c2_server_ip_address)
         self.c2_config.keep_alive_frequency = keep_alive_frequency
@@ -183,7 +189,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
         # Send a keep alive to the C2 Server if we already have a keep alive.
         if self.c2_connection_active is True:
             self.sys_log.info(f"{self.name}: Updating C2 Server with updated C2 configuration.")
-            self._send_keep_alive(self.c2_session.uuid if not None else None)
+            return self._send_keep_alive(self.c2_session.uuid if not None else None)
         return True
 
     def establish(self) -> bool:
@@ -193,14 +199,14 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
             return False
         self.run()
         self.num_executions += 1
+        # Creates a new session if using the establish method.
         return self._send_keep_alive(session_id=None)
 
     def _handle_command_input(self, payload: C2Packet, session_id: Optional[str]) -> bool:
         """
         Handles the parsing of C2 Commands from C2 Traffic (Masquerade Packets).
 
-        Dependant the C2 Command contained within the payload.
+        Dependant the C2 Command parsed from the payload, the following methods are called and returned:
-        The following methods are called and returned.
 
         C2 Command           | Internal Method
         ---------------------|------------------------

src/primaite/simulator/system/applications/red_applications/c2/c2_server.py

@@ -103,15 +103,15 @@ class C2Server(AbstractC2, identifier="C2Server"):
         """
         Handles the parsing of C2 Command Output from C2 Traffic (Masquerade Packets).
 
-        Parses the Request Response within C2Packet's payload attribute (Inherited from Data packet).
+        Parses the Request Response from the given C2Packet's payload attribute (Inherited from Data packet).
-        The class attribute self.current_command_output is then set to this Request Response.
+        This RequestResponse is then stored in the C2 Server class attribute self.current_command_output.
 
         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: C2Packet
-        :return: Returns True if the self.current_command_output is currently updated, false otherwise.
+        :return: Returns True if the self.current_command_output was updated, false otherwise.
         :rtype Bool:
         """
         self.sys_log.info(f"{self.name}: Received command response from C2 Beacon: {payload}.")
@@ -130,20 +130,27 @@ class C2Server(AbstractC2, identifier="C2Server"):
         """
         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
+        Abstract method inherited from abstract C2.
-        is set to true and the keep alive inactivity is reset after receiving one keep alive.
+
+        In the C2 Server implementation of this method the following logic is performed:
+
+        1. The ``self.c2_connection_active`` is set to True. (Indicates that we're received a connection)
+        2. The received keep alive (Payload parameter) is then resolved by _resolve_keep_alive.
+        3. After the keep alive is resolved, a keep alive is sent back to confirm connection.
 
         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)
+        the C2 Beacon rather than the other way around.
+
+        The C2 Beacon/Server communication is akin to that of a real-world reverse shells.
 
         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.
 
         :param payload: The Keep Alive payload received.
         :type payload: C2Packet
-        :param session_id: The transport session_id that the payload is originating from.
+        :param session_id: The transport session_id that the payload originates from.
         :type session_id: str
-        :return: True if successfully handled, false otherwise.
+        :return: True if the keep alive was successfully handled, false otherwise.
         :rtype: Bool
         """
         self.sys_log.info(f"{self.name}: Keep Alive Received. Attempting to resolve the remote connection details.")
@@ -155,16 +162,22 @@ class C2Server(AbstractC2, identifier="C2Server"):
             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}.")
 
+        # If this method returns true then we have sent successfully sent a keep alive response back.
         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.
+        Sends a C2 command to the C2 Beacon using the given parameters.
+
+        C2 Command           | Command Synopsis
+        ---------------------|------------------------
+        RANSOMWARE_CONFIGURE | Configures an installed ransomware script based on the passed parameters.
+        RANSOMWARE_LAUNCH    | Launches the installed ransomware script.
+        TERMINAL             | Executes a command via the terminal installed on the C2 Beacons Host.
 
         Currently, these commands leverage the pre-existing capability of other applications.
         However, the commands are sent via the network rather than the game layer which
@@ -174,12 +187,6 @@ class C2Server(AbstractC2, identifier="C2Server"):
         more complex red agent behaviour such as file extraction, establishing further fall back channels
         or introduce red applications that are only installable via C2 Servers. (T1105)
 
-        C2 Command           | Meaning
-        ---------------------|------------------------
-        RANSOMWARE_CONFIGURE | Configures an installed ransomware script based on the passed parameters.
-        RANSOMWARE_LAUNCH    | Launches the installed ransomware script.
-        TERMINAL             | Executes a command via the terminal installed on the C2 Beacons Host.
-
         For more information on the impact of these commands please refer to the terminal
         and the ransomware applications.
 
@@ -225,6 +232,46 @@ class C2Server(AbstractC2, identifier="C2Server"):
             )
         return self.current_command_output
 
+    def _confirm_remote_connection(self, timestep: int) -> bool:
+        """Checks the suitability of the current C2 Beacon connection.
+
+        Inherited Abstract Method.
+
+        If a C2 Server has not received a keep alive within the current set
+        keep alive frequency (self._keep_alive_frequency) then the C2 beacons
+        connection is considered dead and any commands will be rejected.
+
+        This method is called on each timestep (Called by .apply_timestep)
+
+        :param timestep: The current timestep of the simulation.
+        :type timestep: Int
+        :return: Returns False if the C2 beacon is considered dead. Otherwise True.
+        :rtype bool:
+        """
+        if self.keep_alive_inactivity > self.c2_config.keep_alive_frequency:
+            self.sys_log.info(f"{self.name}: C2 Beacon connection considered dead due to inactivity.")
+            self.sys_log.debug(
+                f"{self.name}: Did not receive expected keep alive connection from {self.c2_remote_connection}"
+                f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.c2_config.keep_alive_frequency}"
+                f"{self.name}: Last Keep Alive received at {(timestep - self.keep_alive_inactivity)}"
+            )
+            self._reset_c2_connection()
+            return False
+        return True
+
+    # Abstract method inherited from abstract C2.
+    # C2 Servers do not currently receive any input commands from the C2 beacon.
+    def _handle_command_input(self, payload: C2Packet) -> None:
+        """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 C2Packet
+        :type payload: C2Packet.
+        """
+        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.
@@ -261,41 +308,3 @@ class C2Server(AbstractC2, identifier="C2Server"):
             ]
         )
         print(table)
-
-    # Abstract method inherited from abstract C2.
-    # C2 Servers do not currently receive any input commands from the C2 beacon.
-    def _handle_command_input(self, payload: C2Packet) -> None:
-        """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 C2Packet
-        :type payload: C2Packet.
-        """
-        self.sys_log.warning(f"{self.name}: C2 Server received an unexpected INPUT payload: {payload}")
-        pass
-
-    def _confirm_remote_connection(self, timestep: int) -> bool:
-        """Checks the suitability of the current C2 Beacon connection.
-
-        If a C2 Server has not received a keep alive within the current set
-        keep alive frequency (self._keep_alive_frequency) then the C2 beacons
-        connection is considered dead and any commands will be rejected.
-
-        This method is used to
-
-        :param timestep: The current timestep of the simulation.
-        :type timestep: Int
-        :return: Returns False if the C2 beacon is considered dead. Otherwise True.
-        :rtype bool:
-        """
-        if self.keep_alive_inactivity > self.c2_config.keep_alive_frequency:
-            self.sys_log.info(f"{self.name}: C2 Beacon connection considered dead due to inactivity.")
-            self.sys_log.debug(
-                f"{self.name}: Did not receive expected keep alive connection from {self.c2_remote_connection}"
-                f"{self.name}: Expected at timestep: {timestep} due to frequency: {self.c2_config.keep_alive_frequency}"
-                f"{self.name}: Last Keep Alive received at {(timestep - self.keep_alive_inactivity)}"
-            )
-            self._reset_c2_connection()
-            return False
-        return True

tests/unit_tests/_primaite/_simulator/_system/_applications/_red_applications/test_c2_suite.py

@@ -219,13 +219,3 @@ def test_c2_handles_1_timestep_keep_alive(basic_c2_network):
 
     assert c2_beacon.c2_connection_active is True
     assert c2_server.c2_connection_active is True
-
-
-def test_c2_server_runs_on_default(basic_c2_network):
-    """Tests that the C2 Server begins running by default."""
-    network: Network = basic_c2_network
-
-    computer_a: Computer = network.get_node_by_hostname("computer_a")
-    c2_server: C2Server = computer_a.software_manager.software.get("C2Server")
-
-    assert c2_server.operating_state == ApplicationOperatingState.RUNNING