From a1553fb1b45410e215030a5cb7b05275ac3a798f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 28 Aug 2024 10:20:32 +0100 Subject: [PATCH] Backport core changes from internal --- src/primaite/game/agent/rewards.py | 13 +++++++++- src/primaite/simulator/network/airspace.py | 10 +++---- .../simulator/network/hardware/base.py | 20 +++++++------- .../red_applications/c2/abstract_c2.py | 1 + .../system/services/terminal/terminal.py | 26 ++++++++++--------- 5 files changed, 42 insertions(+), 28 deletions(-) diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index b97b7c5a..1de34b40 100644 --- a/src/primaite/game/agent/rewards.py +++ b/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 diff --git a/src/primaite/simulator/network/airspace.py b/src/primaite/simulator/network/airspace.py index 9c736383..cdb01514 100644 --- a/src/primaite/simulator/network/airspace.py +++ b/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 diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 4f73ad7b..ef2d47c3 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/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) 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 7316dd63..5d4cc8e0 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 @@ -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) diff --git a/src/primaite/simulator/system/services/terminal/terminal.py b/src/primaite/simulator/system/services/terminal/terminal.py index 406facd1..e98e8555 100644 --- a/src/primaite/simulator/system/services/terminal/terminal.py +++ b/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):