Files
PrimAITE/src/primaite/notebooks/Terminal-Processing.ipynb

497 lines
15 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Terminal Processing\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This notebook serves as a guide on the functionality and use of the new Terminal simulation component.\n",
"\n",
"The Terminal service comes pre-installed on most Nodes (The exception being Switches, as these are currently dumb). "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from primaite.simulator.system.services.terminal.terminal import Terminal\n",
"from primaite.simulator.network.container import Network\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript\n",
"from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection\n",
"\n",
"def basic_network() -> Network:\n",
" \"\"\"Utility function for creating a default network to demonstrate Terminal functionality\"\"\"\n",
" network = Network()\n",
" node_a = Computer(hostname=\"node_a\", ip_address=\"192.168.0.10\", subnet_mask=\"255.255.255.0\", start_up_duration=0)\n",
" node_a.power_on()\n",
" node_b = Computer(hostname=\"node_b\", ip_address=\"192.168.0.11\", subnet_mask=\"255.255.255.0\", start_up_duration=0)\n",
" node_b.power_on()\n",
" network.connect(node_a.network_interface[1], node_b.network_interface[1])\n",
" return network"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The terminal can be accessed from a `Node` via the `software_manager` as demonstrated below. \n",
"\n",
"In the example, we have a basic network consisting of two computers, connected to form a basic network."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"network: Network = basic_network()\n",
"computer_a: Computer = network.get_node_by_hostname(\"node_a\")\n",
"terminal_a: Terminal = computer_a.software_manager.software.get(\"Terminal\")\n",
"computer_b: Computer = network.get_node_by_hostname(\"node_b\")\n",
"terminal_b: Terminal = computer_b.software_manager.software.get(\"Terminal\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To be able to send commands from `node_a` to `node_b`, you will need to `login` to `node_b` first, using valid user credentials. In the example below, we are remotely logging in to the 'admin' account on `node_b`, from `node_a`. \n",
"If you are not logged in, any commands sent will be rejected by the remote.\n",
"\n",
"Remote Logins return a RemoteTerminalConnection object, which can be used for sending commands to the remote node. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Login to the remote (node_b) from local (node_a)\n",
"term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username=\"admin\", password=\"admin\", ip_address=\"192.168.0.11\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can view all active connections to a terminal through use of the `show()` method."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"terminal_b.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The new connection object allows us to forward commands to be executed on the target node. The example below demonstrates how you can remotely install an application on the target node."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"RansomwareScript\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"computer_b.software_manager.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The code block below demonstrates how the Terminal class allows the user of `terminal_a`, on `computer_a`, to send a command to `computer_b` to create a downloads folder. \n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Display the current state of the file system on computer_b\n",
"computer_b.file_system.show()\n",
"\n",
"# Send command\n",
"term_a_term_b_remote_connection.execute([\"file_system\", \"create\", \"folder\", \"downloads\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The resultant call to `computer_b.file_system.show()` shows that the new folder has been created."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"computer_b.file_system.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"When finished, the connection can be closed by calling the `disconnect` function of the Remote Client object"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Display active connection\n",
"terminal_a.show()\n",
"terminal_b.show()\n",
"\n",
"term_a_term_b_remote_connection.disconnect()\n",
"\n",
"terminal_a.show()\n",
"terminal_b.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Disconnected Terminal sessions will no longer show in the node's Terminal connection list, but will be under the historic sessions in the `user_session_manager`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"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": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}