oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merged PR 534: #2840 NODE_SEND_LOCAL_COMMAND

Browse Source 
## 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
This commit is contained in:
Archer Bowen
2024-09-24 10:47:08 +00:00
parent 1a800ff83d e29815305d
commit b0363a6b48
9 changed files with 489 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CHANGELOG.md
View File

@@ -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.

2
docs/source/request_system.rst
View File

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

119
docs/source/simulation_components/system/services/terminal.rst
View File

@@ -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``).
- 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.
Agents are able to use the terminal to login into remote nodes via ``SSH`` which allows for agents to execute commands on remote hosts.
.. code-block:: yaml
...
...
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.

23
src/primaite/game/agent/actions.py
View File

@@ -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."""

274
src/primaite/notebooks/Terminal-Processing.ipynb
View File

@@ -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,

2
src/primaite/simulator/system/core/session_manager.py
View File

@@ -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)

30
src/primaite/simulator/system/services/terminal/terminal.py
View File

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

1
tests/conftest.py
View File

@@ -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(

52
tests/integration_tests/game_layer/actions/test_terminal_actions.py
View File

@@ -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()