From 3df55a708d31f192a8a414673dce3e23e9126486 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Sun, 11 Aug 2024 23:24:29 +0100 Subject: [PATCH] #2769 - add actions and tests for terminal --- src/primaite/game/agent/actions.py | 28 ++- .../system/services/terminal/terminal.py | 120 +++++++------ tests/conftest.py | 7 +- .../actions/test_terminal_actions.py | 165 ++++++++++++++++++ .../test_remote_user_account_actions.py | 49 ------ .../test_user_account_change_password.py | 23 --- 6 files changed, 253 insertions(+), 139 deletions(-) create mode 100644 tests/integration_tests/game_layer/actions/test_terminal_actions.py delete mode 100644 tests/integration_tests/game_layer/actions/user_account_actions/test_remote_user_account_actions.py delete mode 100644 tests/integration_tests/game_layer/actions/user_account_actions/test_user_account_change_password.py diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 2ddeff3d..f421cb0b 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1101,15 +1101,14 @@ class NodeSessionsRemoteLoginAction(AbstractAction): def form_request(self, node_id: str, username: str, password: str, remote_ip: str) -> RequestFormat: """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - # TODO: change this so it creates a remote connection using terminal rather than a local remote login node_name = self.manager.get_node_name_by_idx(node_id) return [ "network", "node", node_name, "service", - "UserSessionManager", - "remote_login", + "Terminal", + "ssh_to_remote", username, password, remote_ip, @@ -1122,11 +1121,21 @@ class NodeSessionsRemoteLogoutAction(AbstractAction): def __init__(self, manager: "ActionManager", **kwargs) -> None: super().__init__(manager=manager) - def form_request(self, node_id: str, remote_session_id: str) -> RequestFormat: + def form_request(self, node_id: str, remote_ip: str) -> RequestFormat: """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" - # TODO: change this so it destroys a remote connection using terminal rather than a local remote login node_name = self.manager.get_node_name_by_idx(node_id) - return ["network", "node", node_name, "service", "UserSessionManager", "remote_logout", remote_session_id] + return ["network", "node", node_name, "service", "Terminal", "remote_logoff", remote_ip] + + +class NodeSendRemoteCommandAction(AbstractAction): + """Action which sends a terminal command to a remote node via SSH.""" + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, remote_ip: str, command: RequestFormat) -> RequestFormat: + node_name = self.manager.get_node_name_by_idx(node_id) + return ["network", "node", node_name, "service", "Terminal", "send_remote_command", remote_ip, command] class ActionManager: @@ -1180,9 +1189,10 @@ class ActionManager: "CONFIGURE_DATABASE_CLIENT": ConfigureDatabaseClientAction, "CONFIGURE_RANSOMWARE_SCRIPT": ConfigureRansomwareScriptAction, "CONFIGURE_DOSBOT": ConfigureDoSBotAction, - "NODE_ACCOUNTS_CHANGEPASSWORD": NodeAccountsChangePasswordAction, - "NODE_SESSIONS_REMOTE_LOGIN": NodeSessionsRemoteLoginAction, - "NODE_SESSIONS_REMOTE_LOGOUT": NodeSessionsRemoteLogoutAction, + "NODE_ACCOUNTS_CHANGE_PASSWORD": NodeAccountsChangePasswordAction, + "SSH_TO_REMOTE": NodeSessionsRemoteLoginAction, + "SSH_LOGOUT_LOGOUT": NodeSessionsRemoteLogoutAction, + "NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction, } """Dictionary which maps action type strings to the corresponding action class.""" diff --git a/src/primaite/simulator/system/services/terminal/terminal.py b/src/primaite/simulator/system/services/terminal/terminal.py index 876b1694..ead5c66a 100644 --- a/src/primaite/simulator/system/services/terminal/terminal.py +++ b/src/primaite/simulator/system/services/terminal/terminal.py @@ -92,7 +92,7 @@ class LocalTerminalConnection(TerminalClientConnection): if not self.is_active: self.parent_terminal.sys_log.warning("Connection inactive, cannot execute") return None - return self.parent_terminal.execute(command, connection_id=self.connection_uuid) + return self.parent_terminal.execute(command) class RemoteTerminalConnection(TerminalClientConnection): @@ -162,22 +162,36 @@ class Terminal(Service): def _init_request_manager(self) -> RequestManager: """Initialise Request manager.""" rm = super()._init_request_manager() - rm.add_request( - "send", - request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())), - ) + # rm.add_request( + # "send", + # request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.send())), + # ) - def _login(request: RequestFormat, context: Dict) -> RequestResponse: - login = self._process_local_login(username=request[0], password=request[1]) - if login: - return RequestResponse( - status="success", - data={ - "ip_address": login.ip_address, - }, - ) - else: - return RequestResponse(status="failure", data={"reason": "Invalid login credentials"}) + # def _login(request: RequestFormat, context: Dict) -> RequestResponse: + # login = self._process_local_login(username=request[0], password=request[1]) + # if login: + # return RequestResponse( + # status="success", + # data={ + # "ip_address": login.ip_address, + # }, + # ) + # else: + # return RequestResponse(status="failure", data={"reason": "Invalid login credentials"}) + # + # rm.add_request( + # "Login", + # request_type=RequestType(func=_login), + # ) + + # def _logoff(request: RequestFormat, context: Dict) -> RequestResponse: + # """Logoff from connection.""" + # connection_uuid = request[0] + # self.parent.user_session_manager.local_logout(connection_uuid) + # self._disconnect(connection_uuid) + # return RequestResponse(status="success", data={}) + # + # rm.add_request("Logoff", request_type=RequestType(func=_logoff)) def _remote_login(request: RequestFormat, context: Dict) -> RequestResponse: login = self._send_remote_login(username=request[0], password=request[1], ip_address=request[2]) @@ -191,10 +205,34 @@ class Terminal(Service): else: return RequestResponse(status="failure", data={}) + rm.add_request( + "ssh_to_remote", + request_type=RequestType(func=_remote_login), + ) + + def _remote_logoff(request: RequestFormat, context: Dict) -> RequestResponse: + """Logoff from remote connection.""" + ip_address = IPv4Address(request[0]) + remote_connection = self._get_connection_from_ip(ip_address=ip_address) + 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."}, + ) + + rm.add_request("remote_logoff", request_type=RequestType(func=_remote_logoff)) + def remote_execute_request(request: RequestFormat, context: Dict) -> RequestResponse: """Execute an instruction.""" - command: str = request[0] - ip_address: IPv4Address = IPv4Address(request[1]) + ip_address: IPv4Address = IPv4Address(request[0]) + command: str = request[1] remote_connection = self._get_connection_from_ip(ip_address=ip_address) if remote_connection: outcome = remote_connection.execute(command) @@ -209,30 +247,11 @@ class Terminal(Service): data={}, ) - def _logoff(request: RequestFormat, context: Dict) -> RequestResponse: - """Logoff from connection.""" - connection_uuid = request[0] - self.parent.user_session_manager.local_logout(connection_uuid) - self._disconnect(connection_uuid) - return RequestResponse(status="success", data={}) - rm.add_request( - "Login", - request_type=RequestType(func=_login), - ) - - rm.add_request( - "Remote Login", - request_type=RequestType(func=_remote_login), - ) - - rm.add_request( - "Execute", + "send_remote_command", request_type=RequestType(func=remote_execute_request), ) - rm.add_request("Logoff", request_type=RequestType(func=_logoff)) - return rm def execute(self, command: List[Any]) -> Optional[RequestResponse]: @@ -280,13 +299,9 @@ class Terminal(Service): if self.operating_state != ServiceOperatingState.RUNNING: self.sys_log.warning(f"{self.name}: Cannot login as service is not running.") return None - connection_request_id = str(uuid4()) - self._client_connection_requests[connection_request_id] = None if ip_address: # Assuming that if IP is passed we are connecting to remote - return self._send_remote_login( - username=username, password=password, ip_address=ip_address, connection_request_id=connection_request_id - ) + return self._send_remote_login(username=username, password=password, ip_address=ip_address) else: return self._process_local_login(username=username, password=password) @@ -320,32 +335,24 @@ class Terminal(Service): username: str, password: str, ip_address: IPv4Address, - connection_request_id: str, + connection_request_id: Optional[str] = None, is_reattempt: bool = False, ) -> Optional[RemoteTerminalConnection]: """Send a remote login attempt and connect to Node. :param: username: Username used to connect to the remote node. :type: username: str - :param: password: Password used to connect to the remote node :type: password: str - :param: ip_address: Target Node IP address for login attempt. :type: ip_address: IPv4Address - - :param: connection_request_id: Connection Request ID - :type: connection_request_id: str - + :param: connection_request_id: Connection Request ID, if not provided, a new one is generated + :type: connection_request_id: Optional[str] :param: is_reattempt: True if the request has been reattempted. Default False. :type: is_reattempt: Optional[bool] - :return: RemoteTerminalConnection: Connection Object for sending further commands if successful, else False. - """ - self.sys_log.info( - f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}" - ) + connection_request_id = connection_request_id or str(uuid4()) if is_reattempt: valid_connection_request = self._validate_client_connection_request(connection_id=connection_request_id) if valid_connection_request: @@ -360,6 +367,9 @@ class Terminal(Service): self.sys_log.warning(f"{self.name}: Remote connection to {ip_address} declined.") return None + self.sys_log.info( + f"{self.name}: Sending Remote login attempt to {ip_address}. Connection_id is {connection_request_id}" + ) transport_message: SSHTransportMessage = SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST connection_message: SSHConnectionMessage = SSHConnectionMessage.SSH_MSG_CHANNEL_DATA user_details: SSHUserCredentials = SSHUserCredentials(username=username, password=password) diff --git a/tests/conftest.py b/tests/conftest.py index d2f9bb2f..2ae6299d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -458,9 +458,10 @@ def game_and_agent(): {"type": "HOST_NIC_DISABLE"}, {"type": "NETWORK_PORT_ENABLE"}, {"type": "NETWORK_PORT_DISABLE"}, - {"type": "NODE_ACCOUNTS_CHANGEPASSWORD"}, - {"type": "NODE_SESSIONS_REMOTE_LOGIN"}, - {"type": "NODE_SESSIONS_REMOTE_LOGOUT"}, + {"type": "NODE_ACCOUNTS_CHANGE_PASSWORD"}, + {"type": "SSH_TO_REMOTE"}, + {"type": "SSH_LOGOUT_LOGOUT"}, + {"type": "NODE_SEND_REMOTE_COMMAND"}, ] action_space = ActionManager( diff --git a/tests/integration_tests/game_layer/actions/test_terminal_actions.py b/tests/integration_tests/game_layer/actions/test_terminal_actions.py new file mode 100644 index 00000000..ce0810eb --- /dev/null +++ b/tests/integration_tests/game_layer/actions/test_terminal_actions.py @@ -0,0 +1,165 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +from typing import Tuple + +import pytest + +from primaite.game.agent.interface import ProxyAgent +from primaite.game.game import PrimaiteGame +from primaite.simulator.network.hardware.base import UserManager +from primaite.simulator.network.hardware.nodes.host.computer import Computer +from primaite.simulator.network.hardware.nodes.host.server import Server +from primaite.simulator.system.services.service import ServiceOperatingState +from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection + + +@pytest.fixture +def game_and_agent_fixture(game_and_agent): + """Create a game with a simple agent that can be controlled by the tests.""" + game, agent = game_and_agent + + client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") + client_1.start_up_duration = 3 + + return (game, agent) + + +def test_remote_login(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): + game, agent = game_and_agent_fixture + + server_1: Server = game.simulation.network.get_node_by_hostname("server_1") + client_1 = game.simulation.network.get_node_by_hostname("client_1") + + # create a new user account on server_1 that will be logged into remotely + server_1_usm: UserManager = server_1.software_manager.software["UserManager"] + server_1_usm.add_user("user123", "password", is_admin=True) + + action = ( + "SSH_TO_REMOTE", + { + "node_id": 0, + "username": "user123", + "password": "password", + "remote_ip": str(server_1.network_interface[1].ip_address), + }, + ) + agent.store_action(action) + game.step() + assert agent.history[-1].response.status == "success" + + connection_established = False + for conn_str, conn_obj in client_1.terminal.connections.items(): + conn_obj: RemoteTerminalConnection + if conn_obj.ip_address == server_1.network_interface[1].ip_address: + connection_established = True + if not connection_established: + pytest.fail("Remote SSH connection could not be established") + + +def test_remote_login_wrong_password(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): + game, agent = game_and_agent_fixture + + server_1: Server = game.simulation.network.get_node_by_hostname("server_1") + client_1 = game.simulation.network.get_node_by_hostname("client_1") + + # create a new user account on server_1 that will be logged into remotely + server_1_usm: UserManager = server_1.software_manager.software["UserManager"] + server_1_usm.add_user("user123", "password", is_admin=True) + + action = ( + "SSH_TO_REMOTE", + { + "node_id": 0, + "username": "user123", + "password": "wrong_password", + "remote_ip": str(server_1.network_interface[1].ip_address), + }, + ) + agent.store_action(action) + game.step() + assert agent.history[-1].response.status == "failure" + + connection_established = False + for conn_str, conn_obj in client_1.terminal.connections.items(): + conn_obj: RemoteTerminalConnection + if conn_obj.ip_address == server_1.network_interface[1].ip_address: + connection_established = True + if connection_established: + pytest.fail("Remote SSH connection was established despite wrong password") + + +def test_remote_login_change_password(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): + game, agent = game_and_agent_fixture + + server_1: Server = game.simulation.network.get_node_by_hostname("server_1") + client_1 = game.simulation.network.get_node_by_hostname("client_1") + + # create a new user account on server_1 that will be logged into remotely + server_1_um: UserManager = server_1.software_manager.software["UserManager"] + server_1_um.add_user("user123", "password", is_admin=True) + + action = ( + "NODE_ACCOUNTS_CHANGE_PASSWORD", + { + "node_id": 1, # server_1 + "username": "user123", + "current_password": "password", + "new_password": "different_password", + }, + ) + agent.store_action(action) + game.step() + assert agent.history[-1].response.status == "success" + assert server_1_um.users["user123"].password == "different_password" + + +def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): + game, agent = game_and_agent_fixture + + server_1: Server = game.simulation.network.get_node_by_hostname("server_1") + client_1 = game.simulation.network.get_node_by_hostname("client_1") + + # create a new user account on server_1 that will be logged into remotely + server_1_usm: UserManager = server_1.software_manager.software["UserManager"] + server_1_usm.add_user("user123", "password", is_admin=True) + + # Log in remotely + action = ( + "SSH_TO_REMOTE", + { + "node_id": 0, + "username": "user123", + "password": "password", + "remote_ip": str(server_1.network_interface[1].ip_address), + }, + ) + agent.store_action(action) + game.step() + + # Change password + action = ( + "NODE_ACCOUNTS_CHANGE_PASSWORD", + { + "node_id": 1, # server_1 + "username": "user123", + "current_password": "password", + "new_password": "different_password", + }, + ) + agent.store_action(action) + game.step() + + # Assert that the user cannot execute an action + # TODO: should the db conn object get destroyed on both nodes? or is that not realistic? + action = ( + "NODE_SEND_REMOTE_COMMAND", + { + "node_id": 0, + "remote_ip": server_1.network_interface[1].ip_address, + "command": ["file_system", "create", "file", "folder123", "doggo.pdf", False], + }, + ) + agent.store_action(action) + game.step() + + assert server_1.file_system.get_folder("folder123") is None + assert server_1.file_system.get_file("folder123", "doggo.pdf") is None diff --git a/tests/integration_tests/game_layer/actions/user_account_actions/test_remote_user_account_actions.py b/tests/integration_tests/game_layer/actions/user_account_actions/test_remote_user_account_actions.py deleted file mode 100644 index 25079226..00000000 --- a/tests/integration_tests/game_layer/actions/user_account_actions/test_remote_user_account_actions.py +++ /dev/null @@ -1,49 +0,0 @@ -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK -from primaite.simulator.network.hardware.nodes.host.computer import Computer - - -def test_remote_logon(game_and_agent): - """Test that the remote session login action works.""" - game, agent = game_and_agent - - client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") - - client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True) - - action = ( - "NODE_SESSIONS_REMOTE_LOGIN", - {"node_id": 0, "username": "test_user", "password": "password", "remote_ip": "10.0.2.2"}, - ) - agent.store_action(action) - game.step() - - assert len(client_1.user_session_manager.remote_sessions) == 1 - - -def test_remote_logoff(game_and_agent): - """Test that the remote session logout action works.""" - game, agent = game_and_agent - - client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") - - client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True) - - action = ( - "NODE_SESSIONS_REMOTE_LOGIN", - {"node_id": 0, "username": "test_user", "password": "password"}, - ) - agent.store_action(action) - game.step() - - assert len(client_1.user_session_manager.remote_sessions) == 1 - - remote_session_id = client_1.user_session_manager.remote_sessions[0].uuid - - action = ( - "NODE_SESSIONS_REMOTE_LOGOUT", - {"node_id": 0, "remote_session_id": remote_session_id}, - ) - agent.store_action(action) - game.step() - - assert len(client_1.user_session_manager.remote_sessions) == 0 diff --git a/tests/integration_tests/game_layer/actions/user_account_actions/test_user_account_change_password.py b/tests/integration_tests/game_layer/actions/user_account_actions/test_user_account_change_password.py deleted file mode 100644 index 3e6f55f6..00000000 --- a/tests/integration_tests/game_layer/actions/user_account_actions/test_user_account_change_password.py +++ /dev/null @@ -1,23 +0,0 @@ -# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK -from primaite.simulator.network.hardware.nodes.host.computer import Computer - - -def test_remote_logon(game_and_agent): - """Test that the remote session login action works.""" - game, agent = game_and_agent - - client_1: Computer = game.simulation.network.get_node_by_hostname("client_1") - - client_1.user_manager.add_user(username="test_user", password="password", bypass_can_perform_action=True) - user = next((user for user in client_1.user_manager.users.values() if user.username == "test_user"), None) - - assert user.password == "password" - - action = ( - "NODE_ACCOUNTS_CHANGEPASSWORD", - {"node_id": 0, "username": user.username, "current_password": user.password, "new_password": "test_pass"}, - ) - agent.store_action(action) - game.step() - - assert user.password == "test_pass"