Merged PR 534: #2840 NODE_SEND_LOCAL_COMMAND

## Summary
Enables agents to use a new CAOS action ``NODE_SEND_LOCAL_COMMAND``.

## Test process

Added a new unit test as well as tested manually via sandbox notebooks.

## Checklist
- [X] PR is linked to a **work item**
- [X] **acceptance criteria** of linked ticket are met
- [X] performed **self-review** of the code
- [X] written **tests** for any new functionality added with this PR
- [X] updated the **documentation** if this PR changes or adds functionality
- [X] written/updated **design docs** if this PR implements new functionality
- [X] updated the **change log**
- [X] ran **pre-commit** checks for code style
- [X] attended to any **TO-DOs** left in the code

Related work items: #2840

CHANGELOG.md

@@ -5,11 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+
 ## [3.4.0]
 
 ### Added
 -   Log observation space data by episode and step.
 -   Added `show_history` method to Agents, allowing you to view actions taken by an agent per step. By default, `DONOTHING` actions are omitted.
+-  New ``NODE_SEND_LOCAL_COMMAND`` action implemented which grants agents the ability to execute commands locally. (Previously limited to remote only)
 
 ### Changed
 -   ACL's are no longer applied to layer-2 traffic.

docs/source/request_system.rst

@@ -2,6 +2,8 @@
 
     © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 
+.. _request_system:
+
 Request System
 **************
 

docs/source/simulation_components/system/services/terminal.rst

@@ -23,6 +23,14 @@ Key capabilities
  - Simulates common Terminal processes/commands.
  - Leverages the Service base class for install/uninstall, status tracking etc.
 
+Usage
+"""""
+
+ - Pre-Installs on any `Node` component (with the exception of `Switches`).
+ - Terminal Clients connect, execute commands and disconnect from remote nodes.
+ - Ensures that users are logged in to the component before executing any commands.
+ - Service runs on SSH port 22 by default.
+ - Enables Agents to send commands both remotely and locally.
 
 Implementation
 """"""""""""""
@@ -30,19 +38,112 @@ Implementation
  - Manages remote connections in a dictionary by session ID.
  - Processes commands, forwarding to the ``RequestManager`` or ``SessionManager`` where appropriate.
  - Extends Service class.
-  - A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
+
+A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
+
+Command Format
+^^^^^^^^^^^^^^
+
+Terminals implement their commands through leveraging the pre-existing :ref:`request_system`.
+
+Due to this Terminals will only accept commands passed within the ``RequestFormat``.
+
+:py:class:`primaite.game.interface.RequestFormat`
+
+For example, ``terminal`` command actions when used in ``yaml`` format are formatted as follows:
+
+.. code-block:: yaml
+
+    command:
+        - "file_system"
+        - "create"
+        - "file"
+        - "downloads"
+        - "cat.png"
+        - "False
+
+This is then loaded from yaml into a dictionary containing the terminal command:
+
+.. code-block:: python
+
+    {"command":["file_system", "create", "file", "downloads", "cat.png", "False"]}
+
+Which is then passed to the ``Terminals`` Request Manager to be executed.
+
+Game Layer Usage (Agents)
+========================
+
+The below code examples demonstrate how to use terminal related actions in yaml files.
+
+yaml
+""""
+
+``NODE_SEND_LOCAL_COMMAND``
+"""""""""""""""""""""""""""
+
+Agents can execute local commands without needing to perform a separate remote login action (``SSH_TO_REMOTE``).
+
+.. code-block:: yaml
+
+    ...
+        ...
+          action: NODE_SEND_LOCAL_COMMAND
+          options:
+            node_id: 0
+            username: admin
+            password: admin
+            command: # Example command - Creates a file called 'cat.png' in the downloads folder.
+              - "file_system"
+              - "create"
+              - "file"
+              - "downloads"
+              - "cat.png"
+              - "False"
 
 
-Usage
+``SSH_TO_REMOTE``
-"""""
+"""""""""""""""""
 
- - Pre-Installs on all ``Nodes`` (with the exception of ``Switches``).
+Agents are able to use the terminal to login into remote nodes via ``SSH`` which allows for agents to execute commands on remote hosts.
- - Terminal Clients connect, execute commands and disconnect from remote nodes.
+
- - Ensures that users are logged in to the component before executing any commands.
+.. code-block:: yaml
- - Service runs on SSH port 22 by default.
+
+    ...
+        ...
+          action: SSH_TO_REMOTE
+          options:
+            node_id: 0
+            username: admin
+            password: admin
+            remote_ip: 192.168.0.10 # Example Ip Address. (The remote host's IP that will be used by ssh)
+
+
+``NODE_SEND_REMOTE_COMMAND``
+""""""""""""""""""""""""""""
+
+After remotely logging into another host, an agent can use the ``NODE_SEND_REMOTE_COMMAND`` to execute commands across the network remotely.
+
+.. code-block:: yaml
+
+    ...
+        ...
+          action: NODE_SEND_REMOTE_COMMAND
+          options:
+            node_id: 0
+            remote_ip: 192.168.0.10
+            command:
+              - "file_system"
+              - "create"
+              - "file"
+              - "downloads"
+              - "cat.png"
+              - "False"
+
+
+
+Simulation Layer Usage
+======================
 
-Usage
-=====
 
 The below code examples demonstrate how to create a terminal, a remote terminal, and how to send a basic application install command to a remote node.
 

src/primaite/game/agent/actions.py

@@ -1298,6 +1298,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."""
 
@@ -1406,6 +1428,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."""
 

src/primaite/notebooks/Terminal-Processing.ipynb

@@ -9,6 +9,13 @@
     "© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Simulation Layer Implementation."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -198,6 +205,271 @@
    "source": [
     "computer_b.user_session_manager.show(include_historic=True, include_session_id=True)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Game Layer Implementation\n",
+    "\n",
+    "This notebook section will detail the implementation of how the game layer utilises the terminal to support different agent actions.\n",
+    "\n",
+    "The ``Terminal`` is used in a variety of different ways in the game layer. Specifically, the terminal is leveraged to implement the following actions:\n",
+    "\n",
+    "\n",
+    "|  Game Layer Action                | Simulation Layer         |\n",
+    "|-----------------------------------|--------------------------|\n",
+    "| ``NODE_SEND_LOCAL_COMMAND``       | Uses the given user credentials, creates a ``LocalTerminalSession`` and executes the given command and returns the ``RequestResponse``.\n",
+    "| ``SSH_TO_REMOTE``                 | Uses the given user credentials and remote IP to create a ``RemoteTerminalSession``.\n",
+    "| ``NODE_SEND_REMOTE_COMMAND``      | Uses the given remote IP to locate the correct ``RemoteTerminalSession``, executes the given command and returns the ``RequestsResponse``."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Game Layer Setup\n",
+    "\n",
+    "Similar to other notebooks, the next code cells create a custom proxy agent to demonstrate how these commands can be leveraged by agents in the ``UC2`` network environment.\n",
+    "\n",
+    "If you're unfamiliar with ``UC2`` then please refer to the [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import yaml\n",
+    "from primaite.config.load import data_manipulation_config_path\n",
+    "from primaite.session.environment import PrimaiteGymEnv"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "custom_terminal_agent = \"\"\"\n",
+    "  - ref: CustomC2Agent\n",
+    "    team: RED\n",
+    "    type: ProxyAgent\n",
+    "    observation_space: null\n",
+    "    action_space:\n",
+    "      action_list:\n",
+    "        - type: DONOTHING\n",
+    "        - type: NODE_SEND_LOCAL_COMMAND\n",
+    "        - type: SSH_TO_REMOTE\n",
+    "        - type: NODE_SEND_REMOTE_COMMAND\n",
+    "      options:\n",
+    "        nodes:\n",
+    "          - node_name: client_1\n",
+    "        max_folders_per_node: 1\n",
+    "        max_files_per_folder: 1\n",
+    "        max_services_per_node: 2\n",
+    "        max_nics_per_node: 8\n",
+    "        max_acl_rules: 10\n",
+    "        ip_list:\n",
+    "          - 192.168.1.21\n",
+    "          - 192.168.1.14\n",
+    "        wildcard_list:\n",
+    "          - 0.0.0.1\n",
+    "      action_map:\n",
+    "        0:\n",
+    "          action: DONOTHING\n",
+    "          options: {}\n",
+    "        1:\n",
+    "          action: NODE_SEND_LOCAL_COMMAND\n",
+    "          options:\n",
+    "            node_id: 0\n",
+    "            username: admin\n",
+    "            password: admin\n",
+    "            command:\n",
+    "                - file_system\n",
+    "                - create\n",
+    "                - file\n",
+    "                - downloads\n",
+    "                - dog.png\n",
+    "                - False\n",
+    "        2:\n",
+    "          action: SSH_TO_REMOTE\n",
+    "          options:\n",
+    "            node_id: 0\n",
+    "            username: admin\n",
+    "            password: admin\n",
+    "            remote_ip: 192.168.10.22\n",
+    "        3:\n",
+    "          action: NODE_SEND_REMOTE_COMMAND\n",
+    "          options:\n",
+    "            node_id: 0\n",
+    "            remote_ip: 192.168.10.22\n",
+    "            command:\n",
+    "                - file_system\n",
+    "                - create\n",
+    "                - file\n",
+    "                - downloads\n",
+    "                - cat.png\n",
+    "                - False\n",
+    "    reward_function:\n",
+    "      reward_components:\n",
+    "        - type: DUMMY\n",
+    "\"\"\"\n",
+    "custom_terminal_agent_yaml = yaml.safe_load(custom_terminal_agent)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(data_manipulation_config_path()) as f:\n",
+    "    cfg = yaml.safe_load(f)\n",
+    "    # removing all agents & adding the custom agent.\n",
+    "    cfg['agents'] = {}\n",
+    "    cfg['agents'] = custom_terminal_agent_yaml\n",
+    "    \n",
+    "env = PrimaiteGymEnv(env_config=cfg)\n",
+    "\n",
+    "client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
+    "client_2: Computer = env.game.simulation.network.get_node_by_hostname(\"client_2\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Terminal Action | ``NODE_SEND_LOCAL_COMMAND`` \n",
+    "\n",
+    "The yaml snippet below shows all the relevant agent options for this action:\n",
+    "\n",
+    "```yaml\n",
+    "\n",
+    "    action_space:\n",
+    "      action_list:\n",
+    "      ...\n",
+    "        - type: NODE_SEND_LOCAL_COMMAND\n",
+    "      ...\n",
+    "      options:\n",
+    "        nodes: # Node List\n",
+    "          - node_name: client_1\n",
+    "        ...\n",
+    "    ...\n",
+    "      action_map:\n",
+    "        1:\n",
+    "          action: NODE_SEND_LOCAL_COMMAND\n",
+    "          options:\n",
+    "            node_id: 0 # Index 0 at the node list.\n",
+    "            username: admin\n",
+    "            password: admin\n",
+    "            command:\n",
+    "                - file_system\n",
+    "                - create\n",
+    "                - file\n",
+    "                - downloads\n",
+    "                - dog.png\n",
+    "                - False\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "env.step(1)\n",
+    "client_1.file_system.show(full=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Terminal Action | ``SSH_TO_REMOTE``  \n",
+    "\n",
+    "The yaml snippet below shows all the relevant agent options for this action:\n",
+    "\n",
+    "```yaml\n",
+    "\n",
+    "    action_space:\n",
+    "      action_list:\n",
+    "      ...\n",
+    "        - type: SSH_TO_REMOTE\n",
+    "      ...\n",
+    "      options:\n",
+    "        nodes: # Node List\n",
+    "          - node_name: client_1\n",
+    "        ...\n",
+    "    ...\n",
+    "      action_map:\n",
+    "        2:\n",
+    "          action: SSH_TO_REMOTE\n",
+    "          options:\n",
+    "            node_id: 0 # Index 0 at the node list.\n",
+    "            username: admin\n",
+    "            password: admin\n",
+    "            remote_ip: 192.168.10.22 # client_2's ip address.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "env.step(2)\n",
+    "client_2.session_manager.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Terminal Action |  ``NODE_SEND_REMOTE_COMMAND``\n",
+    "\n",
+    "The yaml snippet below shows all the relevant agent options for this action:\n",
+    "\n",
+    "```yaml\n",
+    "\n",
+    "    action_space:\n",
+    "      action_list:\n",
+    "      ...\n",
+    "        - type: NODE_SEND_REMOTE_COMMAND\n",
+    "      ...\n",
+    "      options:\n",
+    "        nodes: # Node List\n",
+    "          - node_name: client_1\n",
+    "        ...\n",
+    "    ...\n",
+    "      action_map:\n",
+    "        1:\n",
+    "          action: NODE_SEND_REMOTE_COMMAND\n",
+    "          options:\n",
+    "            node_id: 0 # Index 0 at the node list.\n",
+    "            remote_ip: 192.168.10.22\n",
+    "            commands:\n",
+    "                - file_system\n",
+    "                - create\n",
+    "                - file\n",
+    "                - downloads\n",
+    "                - cat.png\n",
+    "                - False\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "env.step(3)\n",
+    "client_2.file_system.show(full=True)"
+   ]
   }
  ],
  "metadata": {
@@ -216,7 +488,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

src/primaite/simulator/system/core/session_manager.py

@@ -413,5 +413,5 @@ class SessionManager:
         table.align = "l"
         table.title = f"{self.sys_log.hostname} Session Manager"
         for session in self.sessions_by_key.values():
-            table.add_row([session.dst_ip_address, session.dst_port.value, session.protocol.name])
+            table.add_row([session.with_ip_address, session.dst_port.value, session.protocol.name])
         print(table)

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

@@ -208,17 +208,37 @@ class Terminal(Service):
                         status="success",
                         data={},
                     )
-            else:
+            return RequestResponse(
-                return RequestResponse(
+                status="failure",
-                    status="failure",
+                data={},
-                    data={},
+            )
-                )
 
         rm.add_request(
             "send_remote_command",
             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]:

tests/conftest.py

@@ -466,6 +466,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(

tests/integration_tests/game_layer/actions/test_terminal_actions.py

@@ -164,3 +164,55 @@ 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
+    client_1.session_manager.show()