diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 2e6189c0..3dc1f514 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -1266,6 +1266,28 @@ class NodeSendRemoteCommandAction(AbstractAction): ] +class NodeSendLocalCommandAction(AbstractAction): + """Action which sends a terminal command using a local terminal session.""" + + def __init__(self, manager: "ActionManager", **kwargs) -> None: + super().__init__(manager=manager) + + def form_request(self, node_id: int, username: str, password: str, command: RequestFormat) -> RequestFormat: + """Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" + node_name = self.manager.get_node_name_by_idx(node_id) + return [ + "network", + "node", + node_name, + "service", + "Terminal", + "send_local_command", + username, + password, + {"command": command}, + ] + + class TerminalC2ServerAction(AbstractAction): """Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed.""" @@ -1372,6 +1394,7 @@ class ActionManager: "SSH_TO_REMOTE": NodeSessionsRemoteLoginAction, "SESSIONS_REMOTE_LOGOFF": NodeSessionsRemoteLogoutAction, "NODE_SEND_REMOTE_COMMAND": NodeSendRemoteCommandAction, + "NODE_SEND_LOCAL_COMMAND": NodeSendLocalCommandAction, } """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 e98e8555..9b88bbe8 100644 --- a/src/primaite/simulator/system/services/terminal/terminal.py +++ b/src/primaite/simulator/system/services/terminal/terminal.py @@ -208,7 +208,6 @@ class Terminal(Service): status="success", data={}, ) - else: return RequestResponse( status="failure", data={}, @@ -219,6 +218,27 @@ class Terminal(Service): request_type=RequestType(func=remote_execute_request), ) + def local_execute_request(request: RequestFormat, context: Dict) -> RequestResponse: + """Executes a command using a local terminal session.""" + command: str = request[2]["command"] + local_connection = self._process_local_login(username=request[0], password=request[1]) + if local_connection: + outcome = local_connection.execute(command) + if outcome: + return RequestResponse( + status="success", + data={"reason": outcome}, + ) + return RequestResponse( + status="success", + data={"reason": "Local Terminal failed to resolve command. Potentially invalid credentials?"}, + ) + + rm.add_request( + "send_local_command", + request_type=RequestType(func=local_execute_request), + ) + return rm def execute(self, command: List[Any]) -> Optional[RequestResponse]: diff --git a/tests/conftest.py b/tests/conftest.py index 1bbff8f2..8717abfa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -467,6 +467,7 @@ def game_and_agent(): {"type": "SSH_TO_REMOTE"}, {"type": "SESSIONS_REMOTE_LOGOFF"}, {"type": "NODE_SEND_REMOTE_COMMAND"}, + {"type": "NODE_SEND_LOCAL_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 index d011c1e8..d2ea7202 100644 --- a/tests/integration_tests/game_layer/actions/test_terminal_actions.py +++ b/tests/integration_tests/game_layer/actions/test_terminal_actions.py @@ -164,3 +164,54 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam assert server_1.file_system.get_folder("folder123") is None assert server_1.file_system.get_file("folder123", "doggo.pdf") is None + + +def test_local_terminal(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]): + game, agent = game_and_agent_fixture + + 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 + client_1_usm: UserManager = client_1.software_manager.software["UserManager"] + client_1_usm.add_user("user123", "password", is_admin=True) + + action = ( + "NODE_SEND_LOCAL_COMMAND", + { + "node_id": 0, + "username": "user123", + "password": "password", + "command": ["file_system", "create", "file", "folder123", "doggo.pdf", False], + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_folder("folder123") + assert client_1.file_system.get_file("folder123", "doggo.pdf") + + # Change password + action = ( + "NODE_ACCOUNTS_CHANGE_PASSWORD", + { + "node_id": 0, # server_1 + "username": "user123", + "current_password": "password", + "new_password": "different_password", + }, + ) + agent.store_action(action) + game.step() + + action = ( + "NODE_SEND_LOCAL_COMMAND", + { + "node_id": 0, + "username": "user123", + "password": "password", + "command": ["file_system", "create", "file", "folder123", "cat.pdf", False], + }, + ) + agent.store_action(action) + game.step() + + assert client_1.file_system.get_file("folder123", "cat.pdf") is None