Merged PR 520: Backport core changes from internal

src/primaite/game/agent/rewards.py

@@ -190,8 +190,12 @@ class WebServer404Penalty(AbstractReward):
     def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
         """Calculate the reward for the current state.
 
-        :param state: The current state of the simulation.
+        :param state: Current simulation state
         :type state: Dict
+        :param last_action_response: Current agent history state
+        :type last_action_response: AgentHistoryItem state
+        :return: Reward value
+        :rtype: float
         """
         web_service_state = access_from_nested_dict(state, self.location_in_state)
 
@@ -263,6 +267,12 @@ class WebpageUnavailablePenalty(AbstractReward):
         When the green agent requests to execute the browser application, and that request fails, this reward
         component will keep track of that information. In that case, it doesn't matter whether the last webpage
         had a 200 status code, because there has been an unsuccessful request since.
+        :param state: Current simulation state
+        :type state: Dict
+        :param last_action_response: Current agent history state
+        :type last_action_response: AgentHistoryItem state
+        :return: Reward value
+        :rtype: float
         """
         web_browser_state = access_from_nested_dict(state, self.location_in_state)
 
@@ -519,6 +529,7 @@ class RewardFunction:
             weight = comp_and_weight[1]
             total += weight * comp.calculate(state=state, last_action_response=last_action_response)
         self.current_reward = total
+
         return self.current_reward
 
     @classmethod

src/primaite/simulator/network/airspace.py

@@ -60,13 +60,13 @@ class AirSpaceFrequency(Enum):
     @property
     def maximum_data_rate_bps(self) -> float:
         """
-        Retrieves the maximum data transmission rate in bits per second (bps) for the frequency.
+        Retrieves the maximum data transmission rate in bits per second (bps).
 
-        The maximum rates are predefined for known frequencies:
-            - For WIFI_2_4, it returns 100,000,000 bps (100 Mbps).
-            - For WIFI_5, it returns 500,000,000 bps (500 Mbps).
+        The maximum rates are predefined for frequencies.:
+            - WIFI 2.4 supports 100,000,000 bps
+            - WIFI 5 supports 500,000,000 bps
 
-        :return: The maximum data rate in bits per second. If the frequency is not recognized, returns 0.0.
+        :return: The maximum data rate in bits per second.
         """
         if self == AirSpaceFrequency.WIFI_2_4:
             return 100_000_000.0  # 100 Megabits per second

src/primaite/simulator/network/hardware/base.py

@@ -1149,7 +1149,7 @@ class UserSessionManager(Service):
     local_session_timeout_steps: int = 30
     """The number of steps before a local session times out due to inactivity."""
 
-    remote_session_timeout_steps: int = 5
+    remote_session_timeout_steps: int = 30
     """The number of steps before a remote session times out due to inactivity."""
 
     max_remote_sessions: int = 3
@@ -1179,15 +1179,14 @@ class UserSessionManager(Service):
         """
         rm = super()._init_request_manager()
 
-        # todo add doc about request schemas
-        rm.add_request(
-            "remote_login",
-            RequestType(
-                func=lambda request, context: RequestResponse.from_bool(
-                    self.remote_login(username=request[0], password=request[1], remote_ip_address=request[2])
-                )
-            ),
-        )
+        def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse:
+            """Request should take the form [username, password, remote_ip_address]."""
+            username, password, remote_ip_address = request
+            response = RequestResponse.from_bool(self.remote_login(username, password, remote_ip_address))
+            response.data = {"remote_hostname": self.parent.hostname, "username": username}
+            return response
+
+        rm.add_request("remote_login", RequestType(func=_remote_login))
 
         rm.add_request(
             "remote_logout",
@@ -1422,6 +1421,7 @@ class UserSessionManager(Service):
             self.local_session = None
 
         if not local and remote_session_id:
+            self.parent.terminal._disconnect(remote_session_id)
             session = self.remote_sessions.pop(remote_session_id)
         if session:
             self.historic_sessions.append(session)

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

@@ -446,6 +446,7 @@ class AbstractC2(Application, identifier="AbstractC2"):
         if (
             self.operating_state is ApplicationOperatingState.RUNNING
             and self.health_state_actual is SoftwareHealthState.GOOD
+            and self.c2_connection_active is True
         ):
             self.keep_alive_inactivity += 1
             self._confirm_remote_connection(timestep)

src/primaite/simulator/system/services/terminal/terminal.py

@@ -171,7 +171,8 @@ class Terminal(Service):
                 return RequestResponse(
                     status="success",
                     data={
-                        "ip_address": login.ip_address,
+                        "ip_address": str(login.ip_address),
+                        "username": request[0],
                     },
                 )
             else:
@@ -189,15 +190,9 @@ class Terminal(Service):
             if remote_connection:
                 outcome = self._disconnect(remote_connection.connection_uuid)
                 if outcome:
-                    return RequestResponse(
-                        status="success",
-                        data={},
-                    )
-            else:
-                return RequestResponse(
-                    status="failure",
-                    data={"reason": "No remote connection held."},
-                )
+                    return RequestResponse(status="success", data={})
+
+            return RequestResponse(status="failure", data={})
 
         rm.add_request("remote_logoff", request_type=RequestType(func=_remote_logoff))
 
@@ -464,6 +459,10 @@ class Terminal(Service):
                 command = payload.ssh_command
                 valid_connection = self._check_client_connection(payload.connection_uuid)
                 if valid_connection:
+                    remote_session = self.software_manager.node.user_session_manager.remote_sessions.get(
+                        payload.connection_uuid
+                    )
+                    remote_session.last_active_step = self.software_manager.node.user_session_manager.current_timestep
                     self.execute(command)
                     return True
                 else:
@@ -484,7 +483,7 @@ class Terminal(Service):
 
             if payload["type"] == "user_timeout":
                 connection_id = payload["connection_id"]
-                valid_id = self._check_client_connection(connection_id)
+                valid_id = connection_id in self._connections
                 if valid_id:
                     connection = self._connections.pop(connection_id)
                     connection.is_active = False
@@ -500,11 +499,14 @@ class Terminal(Service):
         :param connection_uuid: Connection ID that we want to disconnect.
         :return True if successful, False otherwise.
         """
+        # TODO: Handle the possibility of attempting to disconnect
         if not self._connections:
             self.sys_log.warning(f"{self.name}: No remote connection present")
             return False
 
-        connection = self._connections.pop(connection_uuid)
+        connection = self._connections.pop(connection_uuid, None)
+        if not connection:
+            return False
         connection.is_active = False
 
         if isinstance(connection, RemoteTerminalConnection):