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

Merge branch 'feature/#3110-userguide-fixes' of ssh.dev.azure.com:v3/ma-dev-uk/PrimAITE/PrimAITE into feature/#3110-userguide-fixes

Browse Source
This commit is contained in:
Nick Todd
2025-03-12 12:54:21 +00:00
parent 9d800b6c45 68db549217
commit 5d21bed557
8 changed files with 141 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
docs/_templates/custom-class-template.rst vendored
View File

@@ -12,7 +12,8 @@
.. autoclass:: {{ objname }}
:members:
:show-inheritance:
:inherited-members:
:inherited-members: BaseModel
:exclude-members: model_computed_fields, model_config, model_fields
:special-members: __init__, __call__, __add__, __mul__
{% block methods %}
@@ -22,7 +23,14 @@
.. autosummary::
:nosignatures:
{% for item in methods %}
{%- if not item.startswith('_') %}
{%- if not item.startswith('_') and item not in [
'construct', 'copy', 'dict', 'from_orm', 'json', 'model_construct',
'model_copy', 'model_dump', 'model_dump_json', 'model_json_schema',
'model_parametrized_name', 'model_post_init', 'model_rebuild', '',
'model_validate', 'model_validate_json', 'model_validate_strings',
'parse_file', 'parse_obj', 'parse_raw', 'schema', 'schema_json',
'update_forward_refs', 'validate',
] %}
~{{ name }}.{{ item }}
{%- endif -%}
{%- endfor %}
@@ -35,7 +43,12 @@
.. autosummary::
{% for item in attributes %}
{%- if not item.startswith('_') and item not in [
'model_computed_fields', 'model_config', 'model_extra', 'model_fields',
'model_fields_set',
] %}
~{{ name }}.{{ item }}
{%- endif -%}
{%- endfor %}
{% endif %}
{% endblock %}

50
src/primaite/notebooks/Requests-and-Responses.ipynb
View File

@@ -8,21 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"Agents interact with the PrimAITE simulation via the Request system.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sending a request"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's set up a minimal network simulation and send some requests to see how it works."
"This notebook demonstrates how agents interact with the PrimAITE simulation via the Request system.\n"
]
},
{
@@ -45,6 +31,20 @@
"from primaite.simulator.sim_container import Simulation\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sending a request"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Before we can send some requests we need to set up a minimal network simulation. The code snippet below creates a PrimAITE simulation with a singular generic host called `client`."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -71,9 +71,16 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"A request is structured in a similar way to a command line interface - a list of strings with positional args. It's also possible to supply an optional `context` dictionary. We will craft a request that stops the pre-installed DNSClient service on the client node.\n",
"Now we can simulation component to interact with, we can start sending some requests."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"A request is structured in a similar way to a command line interface - a list of strings with positional args. It's also possible to supply an optional `context` dictionary. We will craft a request that stops the pre-installed `dns-client` service on the client node.\n",
"\n",
"First let's verify that the DNS Client is running on the client.\n"
"First let's verify that the `dns-client` is running on the client."
]
},
{
@@ -82,14 +89,15 @@
"metadata": {},
"outputs": [],
"source": [
"client.software_manager.show()"
"client.software_manager.show()\n",
"client.software_manager.software['dns-client'].operating_state.name"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Send a request to the simulator to stop the DNSClient."
"Send a request to the simulator to stop the `dns-client`."
]
},
{
@@ -110,7 +118,7 @@
"metadata": {},
"source": [
"\n",
"The request returns a `RequestResponse` object which tells us that the request was successfully executed. Let's verify that the DNS client is in a stopped state now."
"The request returns a `RequestResponse` object which tells us that the request was successfully executed. Let's verify that the `dns-client` is in a stopped state now."
]
},
{
@@ -196,7 +204,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, if we try to start the DNSClient back up, we get a failure because we cannot start software on a node that is turned off."
"Now, if we try to start the `dns-client` back up, we get a failure because we cannot start software on a node that is turned off."
]
},
{

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

@@ -13,23 +13,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"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). "
"This notebook serves as a guide on the functionality and use of the `terminal` service from both the simulation and game layers."
]
},
{
@@ -51,8 +35,47 @@
"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",
"from primaite.simulator.system.services.terminal.terminal import RemoteTerminalConnection\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `terminal` service comes pre-installed on most node types. \n",
"\n",
"_The only exception to this being `switches` network nodes, this is because PrimAITE currently only implements 'dumb' switches. `routers` and `firewalls` however all support the `terminal`._"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this notebook, the terminal is demoed on a basic network consisting of two computers, connected together via a link to form a basic LAN network which can be seen by the `basic_network()` method defined below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def basic_network() -> Network:\n",
" \"\"\"Utility function for creating a default network to demonstrate Terminal functionality\"\"\"\n",
" network = Network()\n",
@@ -65,7 +88,6 @@
" # \"startup_duration\": 0,\n",
" }\n",
" )\n",
" print(f\"{node_a=}\")\n",
" node_a.power_on()\n",
" node_b = Computer.from_config(\n",
" config = {\n",
@@ -85,9 +107,7 @@
"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."
"After setting up the network, the terminal can be accessed from a `Node` via the `software_manager`:"
]
},
{
@@ -107,10 +127,11 @@
"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",
"However, before we're able to send commands from `node_a` to `node_b`, you will need to `login` to `node_b` first, using valid user credentials. \n",
"\n",
"Remote Logins return a RemoteTerminalConnection object, which can be used for sending commands to the remote node. "
"After providing successful credentials, the login method will return type of `TerminalClientConnection` object which can then be used for sending commands to the node. \n",
"\n",
"In the example below, we are remotely logging in to the default ***'admin'*** account on `node_b`, from `node_a` (If you are not logged in, any commands sent will be rejected by the remote).\n"
]
},
{
@@ -143,7 +164,9 @@
"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."
"As we logged into a remote node, the login method return a `RemoteTerminalConnection` which allows us to forward commands to be executed on the target node. \n",
"\n",
"The example below demonstrates how you can remotely install an application on the target node."
]
},
{
@@ -168,7 +191,7 @@
"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"
"As the terminal allows us to leverage the [request system](./Requests-and-Responses.ipynb) we have full access to the request manager on any simulation component. For example, the code snippet below demonstrates how we the `terminal` allows the user of `terminal_a`, on `computer_a`, to send a command (in the form of a request) to `computer_b` to create a downloads folder. \n"
]
},
{
@@ -263,14 +286,16 @@
"\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",
"The ``terminal`` is directly leveraged to implement the following agent 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",
"| ``node-session-remote-login`` | 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``."
"| ``node-send-remote-command`` | Uses the given remote IP to locate the correct ``RemoteTerminalSession``, executes the given command and returns the ``RequestsResponse``.\n",
"\n",
"Additionally, the `terminal` is utilised extensively by the [c2 suite](./Command-and-Control-E2E-Demonstration.ipynb) and it's related actions. "
]
},
{
@@ -302,7 +327,7 @@
"outputs": [],
"source": [
"custom_terminal_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" - ref: CustomTerminalAgent\n",
" team: RED\n",
" type: proxy-agent\n",
" action_space:\n",
@@ -373,22 +398,11 @@
"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",
" node_name: client_1\n",
" username: admin\n",
" password: admin\n",
" command:\n",
@@ -420,22 +434,11 @@
"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-session-remote-login\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 2:\n",
" action: node-session-remote-login\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" node_name: client_1\n",
" username: admin\n",
" password: admin\n",
" remote_ip: 192.168.10.22 # client_2's ip address.\n",
@@ -461,24 +464,13 @@
"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",
" 3:\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",
" node_name: client_1\n",
" remote_ip: 192.168.10.22 # client_2's ip address.\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
@@ -501,7 +493,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},

31
src/primaite/notebooks/Training-an-RLLIB-MARL-System.ipynb
View File

@@ -8,7 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use the `PrimaiteRayMARLEnv` to train a very basic system with two PPO agents."
"This notebook will demonstrate how to use the `PrimaiteRayMARLEnv` to train a very basic system with two PPO agents on the [UC2 scenario](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
@@ -34,17 +34,20 @@
"outputs": [],
"source": [
"import yaml\n",
"\n",
"from primaite import PRIMAITE_PATHS\n",
"\n",
"import ray\n",
"from primaite import PRIMAITE_PATHS\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n",
"from primaite.session.ray_envs import PrimaiteRayMARLEnv\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"\n",
"from primaite.session.ray_envs import PrimaiteRayMARLEnv"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"ray.init(local_mode=True)"
]
},
@@ -114,6 +117,18 @@
"display_name": "Python 3 (ipykernel)",
"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,

9
src/primaite/notebooks/Training-an-RLLib-Agent.ipynb
View File

@@ -8,7 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use PrimaiteRayEnv to train a basic PPO agent."
"This notebook demonstrates how to use the ``PrimaiteRayEnv`` to train a basic PPO agent on the [UC2 scenario](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
@@ -27,13 +27,10 @@
"outputs": [],
"source": [
"import yaml\n",
"from primaite.config.load import data_manipulation_config_path\n",
"\n",
"from primaite.session.ray_envs import PrimaiteRayEnv\n",
"import ray\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.ray_envs import PrimaiteRayEnv\n",
"from ray.rllib.algorithms.ppo import PPOConfig\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"\n",
"\n",
"# If you get an error saying this config file doesn't exist, you may need to run `primaite setup` in your command line\n",
"# to copy the files to your user data path.\n",

14
src/primaite/notebooks/Training-an-SB3-Agent.ipynb
View File

@@ -8,7 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file."
"This notebook will demonstrate how to use primaite to create and train a PPO agent, using a pre-defined configuration file using the [UC2 scenario](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
@@ -33,21 +33,11 @@
"metadata": {},
"outputs": [],
"source": [
"from primaite.game.game import PrimaiteGame\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent\n",
"from primaite.config.load import data_manipulation_config_path\n",
"import yaml"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from primaite.config.load import data_manipulation_config_path"
]
},
{
"cell_type": "code",
"execution_count": null,

9
src/primaite/notebooks/Using-Episode-Schedules.ipynb
View File

@@ -48,7 +48,6 @@
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite import PRIMAITE_PATHS\n",
"from prettytable import PrettyTable\n",
"from primaite.game.agent.scripted_agents import probabilistic_agent, data_manipulation_bot\n",
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/scenario_with_placeholders\""
]
},
@@ -239,7 +238,7 @@
"### Episode 2\n",
"When we reset the environment again, it moves onto episode 2, where it will bring in greens_1 and reds_1 for green and red agent definitions. Let's verify the agent names and that they take actions at the defined frequency.\n",
"\n",
"Most green actions will be `node_application_execute` while red will `DONOTHING` except at steps 10 and 20."
"Most green actions will be `node-application-execute` while red will `do-nothing` except at steps 10 and 20."
]
},
{
@@ -270,7 +269,7 @@
"### Episode 3\n",
"When we reset the environment again, it moves onto episode 3, where it will bring in greens_2 and reds_2 for green and red agent definitions. Let's verify the agent names and that they take actions at the defined frequency.\n",
"\n",
"Now, green will perform `node_application_execute` only 5% of the time, while red will perform `node_application_execute` more frequently than before."
"Now, green will perform `node-application-execute` only 5% of the time, while red will perform `node-application-execute` more frequently than before."
]
},
{
@@ -353,7 +352,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"In the 0th episode, simulation_variant_1.yaml is loaded in and the server gets a `DatabaseService`, while client_1 gets `DatabaseClient`."
"In the 0th episode, `simulation_variant_1.yaml` is loaded in and the server gets a `database-service`, while client_1 gets `database-client`."
]
},
{
@@ -382,7 +381,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"In the 1st episode, `simulation_variant_2.yaml` is loaded in, therefore the server gets a `FTPServer` and client_1 gets a `RansomwareScript`."
"In the 1st episode, `simulation_variant_2.yaml` is loaded in, therefore the server gets a `ftp-server` and client_1 gets a `ransomware-script`."
]
},
{

15
src/primaite/notebooks/multi-processing.ipynb
View File

@@ -8,7 +8,7 @@
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook uses SubprocVecEnv from SB3."
"This notebook uses **SubprocVecEnv** from **SB3** to train an agent on the [UC2 scenario](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
@@ -37,15 +37,7 @@
"from stable_baselines3 import PPO\n",
"from stable_baselines3.common.utils import set_random_seed\n",
"from stable_baselines3.common.vec_env import SubprocVecEnv\n",
"from primaite.session.environment import PrimaiteGymEnv\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.config.load import data_manipulation_config_path"
]
},
@@ -72,12 +64,11 @@
"metadata": {},
"outputs": [],
"source": [
"\n",
"EPISODE_LEN = 128\n",
"NUM_EPISODES = 10\n",
"NO_STEPS = EPISODE_LEN * NUM_EPISODES\n",
"BATCH_SIZE = 32\n",
"LEARNING_RATE = 3e-4\n"
"LEARNING_RATE = 3e-4"
]
},
{