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

Merged PR 618: Port core notebooks

Browse Source 
## Summary
This PR covers changes made for the core notebooks to run correctly.

## Test process
All notebooks were tested to:

- ensure they run to completion
- the output (rewards etc) matches or is similar to the PrimAITE 3.3.1-internal notebooks

All changes run against test suite and successful pipeline builds.

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

Related work items: #3075
This commit is contained in:
Nick Todd
2025-02-26 10:49:04 +00:00
parent 5a7caa33f0 1605b09899
commit 502d163818
21 changed files with 315 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
.azure/azure-ci-build-pipeline.yaml
View File

@@ -113,10 +113,8 @@ stages:
- script: |
pytest --nbmake -n=auto src/primaite/notebooks --junit-xml=./notebook-tests/notebooks.xml
notebooks_exit_code=$?
pytest --nbmake -n=auto src/primaite/simulator/_package_data --junit-xml=./notebook-tests/package-notebooks.xml
package_notebooks_exit_code=$?
# Fail step if either of these do not have exit code 0
if [ $notebooks_exit_code -ne 0 ] || [ $package_notebooks_exit_code -ne 0 ]; then
# Fail step if exit code not equal to 0
if [ $notebooks_exit_code -ne 0 ]; then
exit 1
fi
displayName: 'Run notebooks on Linux and macOS'
@@ -126,11 +124,8 @@ stages:
- script: |
pytest --nbmake -n=auto src/primaite/notebooks --junit-xml=./notebook-tests/notebooks.xml
set notebooks_exit_code=%ERRORLEVEL%
pytest --nbmake -n=auto src/primaite/simulator/_package_data --junit-xml=./notebook-tests/package-notebooks.xml
set package_notebooks_exit_code=%ERRORLEVEL%
rem Fail step if either of these do not have exit code 0
rem Fail step if exit code not equal to 0
if %notebooks_exit_code% NEQ 0 exit /b 1
if %package_notebooks_exit_code% NEQ 0 exit /b 1
displayName: 'Run notebooks on Windows'
condition: eq(variables['Agent.OS'], 'Windows_NT')

1
MANIFEST.in
View File

@@ -1,3 +1,2 @@
include src/primaite/setup/_package_data/primaite_config.yaml
include src/primaite/config/_package_data/*.yaml
include src/primaite/simulator/_package_data/*.ipynb

2
src/primaite/config/_package_data/data_manipulation.yaml
View File

@@ -84,7 +84,7 @@ agents:
action: node-application-execute
options:
node_name: client_1
application_name: web-browser
application_name: database-client
reward_function:
reward_components:

6
src/primaite/notebooks/Action-masking.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Action Masking\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"PrimAITE environments support action masking. The action mask shows which of the agent's actions are applicable with the current environment state. For example, a node can only be turned on if it is currently turned off."
]
@@ -204,7 +204,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -218,7 +218,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

339
src/primaite/notebooks/Command-and-Control-E2E-Demonstration.ipynb
View File

@@ -6,11 +6,20 @@
"source": [
"# Command and Control Application Suite E2E Demonstration\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook demonstrates the current implementation of the command and control (C2) server and beacon applications in primAITE."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -50,109 +59,79 @@
"custom_c2_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: ProxyAgent\n",
" type: proxy-agent\n",
"\n",
" action_space:\n",
" options:\n",
" nodes:\n",
" - node_name: web_server\n",
" applications:\n",
" - application_name: C2Beacon\n",
" - node_name: client_1\n",
" applications:\n",
" - application_name: C2Server\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: do_nothing\n",
" action: do-nothing\n",
" options: {}\n",
" 1:\n",
" action: node_application_install\n",
" action: node-application-install\n",
" options:\n",
" node_id: 0\n",
" application_name: C2Beacon\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
" 2:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" config:\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\n",
" node_name: web_server\n",
" c2_server_ip_address: 192.168.10.21\n",
" 3:\n",
" action: node_application_execute\n",
" action: node-application-execute\n",
" options:\n",
" node_id: 0\n",
" application_id: 0\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
" 4:\n",
" action: c2_server_terminal_command\n",
" action: c2-server-terminal-command\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" ip_address:\n",
" account:\n",
" username: admin\n",
" password: admin\n",
" username: admin\n",
" password: admin\n",
" commands:\n",
" -\n",
" - software_manager\n",
" - application\n",
" - install\n",
" - RansomwareScript\n",
" - ransomware-script\n",
" 5:\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_id: 1\n",
" config:\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" node_name: client_1\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" 6:\n",
" action: c2_server_data_exfiltrate\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" target_file_name: \"database.db\"\n",
" target_folder_name: \"database\"\n",
" exfiltration_folder_name: \"spoils\"\n",
" target_ip_address: 192.168.1.14\n",
" account:\n",
" username: admin\n",
" password: admin\n",
" username: admin\n",
" password: admin\n",
"\n",
" 7:\n",
" action: c2_server_ransomware_launch\n",
" action: c2-server-ransomware-launch\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" 8:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" config:\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency: 10\n",
" masquerade_protocol: TCP\n",
" masquerade_port: DNS\n",
" node_name: web_server\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency: 10\n",
" masquerade_protocol: tcp\n",
" masquerade_port: dns\n",
" 9:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" config:\n",
" c2_server_ip_address: 192.168.10.22\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\n",
" node_name: web_server\n",
" c2_server_ip_address: 192.168.10.22\n",
"\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
" - type: dummy\n",
"\"\"\"\n",
"c2_agent_yaml = yaml.safe_load(custom_c2_agent)"
]
@@ -193,7 +172,7 @@
"source": [
"client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_1.software_manager.install(C2Server)\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.run()\n",
"client_1.software_manager.show()"
]
@@ -225,15 +204,15 @@
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: C2Beacon\n",
" - application_name: c2-beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: node_application_install \n",
" action: node-application-install \n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" application_name: C2Beacon\n",
" application_name: c2-beacon\n",
"```"
]
},
@@ -259,23 +238,18 @@
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
" ...\n",
" ...\n",
" action_map:\n",
" ...\n",
" 2:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0 # Node Index\n",
" config: # Further information about these config options can be found at the bottom of this notebook.\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\n",
" node_name: web_server\n",
" c2_server_ip_address: 192.168.10.21 # Further information about these config options can be found at the bottom of this notebook.\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\n",
"\n",
"```"
]
},
@@ -286,7 +260,7 @@
"outputs": [],
"source": [
"env.step(2)\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
"web_server.software_manager.show()\n",
"c2_beacon.show()"
]
@@ -297,26 +271,18 @@
"source": [
"### **Command and Control** | C2 Beacon Actions | node_application_execute\n",
"\n",
"The final action is ``node_application_execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n",
"The final action is ``node-application-execute`` which is used to establish a connection for the C2 application. This action can be called by the Red Agent via action ``3`` in it's action map. \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: C2Beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
" ...\n",
" 3:\n",
" action: node_application_execute\n",
" action: node-application-execute\n",
" options:\n",
" node_id: 0\n",
" application_id: 0\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
"```"
]
},
@@ -355,34 +321,26 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | C2_SERVER_TERMINAL_COMMAND\n",
"### **Command and Control** | C2 Server Actions | c2-server-terminal-command\n",
"\n",
"The C2 Server's terminal action: ``C2_SERVER_TERMINAL_COMMAND`` is indexed at ``4`` in it's action map. \n",
"The C2 Server's terminal action: ``c2-server-terminal-command`` is indexed at ``4`` in it's action map. \n",
"\n",
"This action leverages the terminal service that is installed by default on all nodes to grant red agents a lot more configurability. If you're unfamiliar with terminals then it's recommended that you refer to the ``Terminal Processing`` notebook.\n",
"\n",
"It's worth noting that an additional benefit a red agent has when using the terminal service via the C2 Server is that you can execute multiple commands in one action. \n",
"\n",
"In this notebook, the ``C2_SERVER_TERMINAL_COMMAND`` is used to install a RansomwareScript application on the ``web_server`` node.\n",
"In this notebook, the ``c2-server-terminal-command`` is used to install a RansomwareScript application on the ``web_server`` node.\n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 4:\n",
" action: C2_SERVER_TERMINAL_COMMAND\n",
" action: c2-server-terminal-command\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" ip_address:\n",
" account:\n",
" username: admin\n",
" password: admin\n",
" commands:\n",
@@ -390,7 +348,7 @@
" - software_manager\n",
" - application\n",
" - install\n",
" - RansomwareScript\n",
" - ransomware-script\n",
"```"
]
},
@@ -426,21 +384,13 @@
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 5:\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_id: 1\n",
" config:\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" node_name: client_1\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
"```\n"
]
},
@@ -459,7 +409,7 @@
"metadata": {},
"outputs": [],
"source": [
"ransomware_script: RansomwareScript = web_server.software_manager.software[\"RansomwareScript\"]\n",
"ransomware_script: RansomwareScript = web_server.software_manager.software[\"ransomware-script\"]\n",
"web_server.software_manager.show()\n",
"ransomware_script.show()"
]
@@ -468,9 +418,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2_server_data_exfiltrate\n",
"### **Command and Control** | C2 Server Actions | c2-server-data-exfiltrate\n",
"\n",
"The second to last action available is the ``c2_server_data_exfiltrate`` which is indexed as action ``6`` in the action map.\n",
"The second to last action available is the ``c2-server-data-exfiltrate`` which is indexed as action ``6`` in the action map.\n",
"\n",
"This action can be used to exfiltrate a target file on a remote node to the C2 Beacon and the C2 Server's host file system via the ``FTP`` services.\n",
"\n",
@@ -478,25 +428,17 @@
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 6:\n",
" action: c2_server_data_exfiltrate\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
" node_id: 1\n",
" target_file_name: \"database.db\"\n",
" target_folder_name: \"database\"\n",
" exfiltration_folder_name: \"spoils\"\n",
" target_ip_address: \"192.168.1.14\"\n",
" account:\n",
" username: \"admin\",\n",
" password: \"admin\"\n",
" username: \"admin\"\n",
" password: \"admin\"\n",
"\n",
"```"
]
@@ -534,9 +476,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2_server_ransomware_launch\n",
"### **Command and Control** | C2 Server Actions | c2-server-ransomware-launch\n",
"\n",
"Finally, the last available action is for the c2_server_ransomware_launch to start the ransomware script installed on the same node as the C2 beacon.\n",
"Finally, the last available action is for the c2-server-ransomware-launch to start the ransomware script installed on the same node as the C2 beacon.\n",
"\n",
"This action is indexed as action ``7``.\n",
"\n",
@@ -544,16 +486,9 @@
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 7:\n",
" action: c2_server_ransomware_launch\n",
" action: c2-server-ransomware-launch\n",
" options:\n",
" node_id: 1\n",
"```\n"
@@ -598,20 +533,20 @@
"custom_blue_agent_yaml = \"\"\"\n",
" - ref: defender\n",
" team: BLUE\n",
" type: ProxyAgent\n",
" type: proxy-agent\n",
"\n",
" observation_space:\n",
" type: CUSTOM\n",
" type: custom\n",
" options:\n",
" components:\n",
" - type: NODES\n",
" label: NODES\n",
" - type: nodes\n",
" label: nodes\n",
" options:\n",
" hosts:\n",
" - hostname: web_server\n",
" applications:\n",
" - application_name: C2Beacon\n",
" - application_name: RansomwareScript\n",
" - application_name: c2-beacon\n",
" - application_name: ransomware-script\n",
" folders:\n",
" - folder_name: exfiltration_folder\n",
" files:\n",
@@ -661,7 +596,7 @@
" - UDP\n",
" num_rules: 10\n",
"\n",
" - type: LINKS\n",
" - type: links\n",
" label: LINKS\n",
" options:\n",
" link_references:\n",
@@ -675,67 +610,41 @@
" - switch_2:eth-1<->client_1:eth-1\n",
" - switch_2:eth-2<->client_2:eth-1\n",
" - switch_2:eth-7<->security_suite:eth-2\n",
" - type: \"NONE\"\n",
" - type: \"none\"\n",
" label: ICS\n",
" options: {}\n",
"\n",
" action_space:\n",
" action_map:\n",
" 0:\n",
" action: do_nothing\n",
" action: do-nothing\n",
" options: {}\n",
" 1:\n",
" action: node_application_remove\n",
" action: node-application-remove\n",
" options:\n",
" node_id: 0\n",
" application_name: C2Beacon\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
" 2:\n",
" action: node_shutdown\n",
" action: node-shutdown\n",
" options:\n",
" node_id: 0\n",
" node_name: web_server\n",
" 3:\n",
" action: router_acl_add_rule\n",
" action: router-acl-add-rule\n",
" options:\n",
" target_router: router_1\n",
" position: 1\n",
" permission: 2\n",
" source_ip_id: 2\n",
" dest_ip_id: 3\n",
" source_port_id: 2\n",
" dest_port_id: 2\n",
" protocol_id: 1\n",
" source_wildcard_id: 0\n",
" dest_wildcard_id: 0\n",
" permission: DENY\n",
" src_ip: 192.168.10.21\n",
" dst_ip: 192.168.1.12\n",
" src_port: HTTP\n",
" dst_port: HTTP\n",
" protocol_name: ALL\n",
" src_wildcard: 0.0.0.1\n",
" dst_wildcard: 0.0.0.1\n",
"\n",
"\n",
" options:\n",
" nodes:\n",
" - node_name: web_server\n",
" applications:\n",
" - application_name: C2Beacon\n",
"\n",
" - node_name: database_server\n",
" folders:\n",
" - folder_name: database\n",
" files:\n",
" - file_name: database.db\n",
" services:\n",
" - service_name: DatabaseService\n",
" - node_name: router_1\n",
"\n",
" max_folders_per_node: 2\n",
" max_files_per_folder: 2\n",
" max_services_per_node: 2\n",
" max_nics_per_node: 8\n",
" max_acl_rules: 10\n",
" ip_list:\n",
" - 192.168.10.21\n",
" - 192.168.1.12\n",
" wildcard_list:\n",
" - 0.0.0.1\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
" - type: dummy\n",
"\n",
" agent_settings:\n",
" flatten_obs: False\n",
@@ -826,12 +735,12 @@
"\n",
"# Installing the C2 Server.\n",
"client_1.software_manager.install(C2Server)\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.run()\n",
"\n",
"# Installing the C2 Beacon.\n",
"web_server.software_manager.install(C2Beacon)\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
"c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
"c2_beacon.establish()"
]
@@ -875,7 +784,7 @@
"outputs": [],
"source": [
"# Installing RansomwareScript via C2 Terminal Commands\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"]],\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"ransomware-script\"]],\n",
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)\n"
@@ -1034,11 +943,11 @@
" web_server: Server = given_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
"\n",
" client_1.software_manager.install(C2Server)\n",
" c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
" c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
" c2_server.run()\n",
"\n",
" web_server.software_manager.install(C2Beacon)\n",
" c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
" c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
" c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
" c2_beacon.establish()\n",
"\n",
@@ -1136,7 +1045,7 @@
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
]
},
@@ -1224,7 +1133,7 @@
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
]
},
@@ -1387,18 +1296,20 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"As demonstrated earlier, red agents can use the ``configure_c2_beacon`` action to configure these settings mid episode through the configuration options:\n",
"As demonstrated earlier, red agents can use the ``configure-c2-beacon`` action to configure these settings mid episode through the configuration options:\n",
"\n",
"``` YAML\n",
"...\n",
" action: configure_c2_beacon\n",
" options:\n",
" node_id: 0\n",
" config:\n",
"```YAML\n",
"\n",
" action_space:\n",
" action_map:\n",
" 8:\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency: 10\n",
" masquerade_protocol: TCP\n",
" masquerade_port: DNS\n",
" masquerade_protocol: tcp\n",
" masquerade_port: dns\n",
"```"
]
},
@@ -1446,16 +1357,16 @@
"source": [
"web_server: Server = c2_config_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
"web_server.software_manager.install(C2Beacon)\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
"\n",
"client_1: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_1.software_manager.install(C2Server)\n",
"c2_server_1: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server_1: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server_1.run()\n",
"\n",
"client_2: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_2\")\n",
"client_2.software_manager.install(C2Server)\n",
"c2_server_2: C2Server = client_2.software_manager.software[\"C2Server\"]\n",
"c2_server_2: C2Server = client_2.software_manager.software[\"c2-server\"]\n",
"c2_server_2.run()"
]
},

94
src/primaite/notebooks/Data-Manipulation-Customising-Red-Agent.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Customising UC2 Red Agents\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook will go over some examples of how red agent behaviour can be varied by changing its configuration parameters.\n",
"\n",
@@ -76,11 +76,10 @@
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'do_nothing':\n",
" if red_action == 'do-nothing':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'node_application_execute':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" elif red_action == 'node-application-execute':\n",
" red_str = f\"ATTACK from {red_info.parameters['node_name']}\"\n",
" return red_str"
]
},
@@ -147,42 +146,15 @@
"```yaml\n",
" - ref: data_manipulation_attacker # name of agent\n",
" team: RED # not used, just for human reference\n",
" type: RedDatabaseCorruptingAgent # type of agent - this lets primaite know which agent class to use\n",
"\n",
" # Since the agent does not need to react to what is happening in the environment, the observation space is empty.\n",
" observation_space:\n",
" type: UC2RedObservation\n",
" options:\n",
" nodes: {}\n",
"\n",
" action_space:\n",
" \n",
" # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
" options:\n",
" nodes:\n",
" - node_name: client_1 # The network should have a node called client_1\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_1 should have DataManipulationBot configured on it\n",
" - node_name: client_2 # The network should have a node called client_2\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_2 should have DataManipulationBot configured on it\n",
"\n",
" # not important\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\n",
" # red agent does not need a reward function\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
" type: red-database-corrupting-agent # type of agent - this lets primaite know which agent class to use\n",
"\n",
" # These actions are passed to the RedDatabaseCorruptingAgent init method, they dictate the schedule of attacks\n",
" agent_settings:\n",
" start_settings:\n",
" start_step: 25 # first attack at step 25\n",
" frequency: 20 # attacks will happen every 20 steps (on average)\n",
" variance: 5 # the timing of attacks will vary by up to 5 steps earlier or later\n",
" possible_start_nodes: [client_1, client_2] # List of clients the attack can start from\n",
" target_application: data-manipulation-bot\n",
" start_step: 25 # first attack at step 25\n",
" frequency: 20 # attacks will happen every 20 steps (on average)\n",
" variance: 5 # the timing of attacks will vary by up to 5 steps earlier or later\n",
"```"
]
},
@@ -202,8 +174,7 @@
"simulation:\n",
" network:\n",
" nodes:\n",
" - ref: client_1\n",
" hostname: client_1\n",
" - hostname: client_1\n",
" type: computer\n",
" ip_address: 192.168.10.21\n",
" subnet_mask: 255.255.255.0\n",
@@ -211,15 +182,13 @@
" \n",
" # \n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
" data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
" payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
" server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient # Database client must be installed in order for DataManipulationBot to function\n",
" - type: database-client # Database client must be installed in order for DataManipulationBot to function\n",
" options:\n",
" db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
"```"
@@ -243,7 +212,8 @@
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"start_settings:\n",
" possible_start_nodes: [client_1]\n",
" target_application: DataManipulationBot\n",
" start_step: 25\n",
" frequency: 20\n",
" variance: 0\n",
@@ -253,7 +223,9 @@
" cfg = yaml.safe_load(f)\n",
" for agent in cfg['agents']:\n",
" if agent['ref'] == \"data_manipulation_attacker\":\n",
" print(f\"{agent['agent_settings']=}\")\n",
" agent['agent_settings'] = change\n",
" print(f\"{agent['agent_settings']=}\")\n",
"\n",
"env = PrimaiteGymEnv(env_config = cfg)\n",
"env.reset()\n",
@@ -310,9 +282,21 @@
"outputs": [],
"source": [
"change = yaml.safe_load(\"\"\"\n",
"# TODO:\n",
" agent_settings:\n",
" possible_start_nodes: [client_1]\n",
" target_application: DataManipulationBot\n",
"\n",
" action_space:\n",
" action_map:\n",
" 0:\n",
" action: do-nothing\n",
" options: {}\n",
" 1:\n",
" action: node-application-execute\n",
" options:\n",
" node_name: client_1\n",
" application_name: DataManipulationBot\n",
"\"\"\")\n",
"#TODO 2869 fix\n",
"\n",
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
@@ -355,18 +339,18 @@
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 1.0\n",
" data_manipulation_p_of_success: 1.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
@@ -400,18 +384,18 @@
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.0\n",
" data_manipulation_p_of_success: 0.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
@@ -438,7 +422,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -452,7 +436,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

14
src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Data Manipulation Scenario\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK"
]
},
{
@@ -429,9 +429,9 @@
" cfg = yaml.safe_load(f)\n",
" # set success probability to 1.0 to avoid rerunning cells.\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][1]['options']['port_scan_p_of_success'] = 1.0\n",
" # don't flatten observations so that we can see what is going on\n",
" cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n",
"\n",
@@ -458,10 +458,10 @@
" # parse the info dict form step output and write out what the red agent is doing\n",
" red_info : AgentHistoryItem = info['agent_actions']['data_manipulation_attacker']\n",
" red_action = red_info.action\n",
" if red_action == 'do_nothing':\n",
" if red_action == 'do-nothing':\n",
" red_str = 'DO NOTHING'\n",
" elif red_action == 'node_application_execute':\n",
" client = \"client 1\" if red_info.parameters['node_id'] == 0 else \"client 2\"\n",
" elif red_action == 'node-application-execute':\n",
" client = \"client 1\" if red_info.parameters['node_name'] == 0 else \"client 2\"\n",
" red_str = f\"ATTACK from {client}\"\n",
" return red_str"
]
@@ -600,7 +600,7 @@
"while abs(reward - 0.8) > 1e-5:\n",
" obs, reward, terminated, truncated, info = env.step(0) # do nothing\n",
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )\n",
" if env.game.step_counter > 10000:\n",
" if env.game.step_counter > 2000:\n",
" break # make sure there's no infinite loop if something went wrong"
]
},

6
src/primaite/notebooks/Getting-Information-Out-Of-PrimAITE.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Getting information out of PrimAITE\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n"
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n"
]
},
{
@@ -157,7 +157,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -171,7 +171,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.8"
"version": "3.10.12"
}
},
"nbformat": 4,

44
src/primaite/notebooks/Privilege-Escalation-and-Data-Loss-Example.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Simulating Privilege Escalation and Data Loss Using SSH and ACLs Manipulation\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"## Overview\n",
"\n",
@@ -51,6 +51,15 @@
"## The Scenario"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -68,7 +77,8 @@
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.system.applications.database_client import DatabaseClient\n",
"from primaite.simulator.system.applications.web_browser import WebBrowser\n",
"from primaite.simulator.system.services.database.database_service import DatabaseService"
"from primaite.simulator.system.services.database.database_service import DatabaseService\n",
"from primaite.simulator.network.hardware.nodes.network import firewall\n"
]
},
{
@@ -110,11 +120,11 @@
"outputs": [],
"source": [
"some_tech_jnr_dev_pc: Computer = game.simulation.network.get_node_by_hostname(\"some_tech_jnr_dev_pc\")\n",
"some_tech_jnr_dev_db_client: DatabaseClient = some_tech_jnr_dev_pc.software_manager.software[\"DatabaseClient\"]\n",
"some_tech_jnr_dev_web_browser: WebBrowser = some_tech_jnr_dev_pc.software_manager.software[\"WebBrowser\"]\n",
"some_tech_jnr_dev_db_client: DatabaseClient = some_tech_jnr_dev_pc.software_manager.software[\"database-client\"]\n",
"some_tech_jnr_dev_web_browser: WebBrowser = some_tech_jnr_dev_pc.software_manager.software[\"web-browser\"]\n",
"some_tech_rt: Router = game.simulation.network.get_node_by_hostname(\"some_tech_rt\")\n",
"some_tech_db_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_db_srv\")\n",
"some_tech_db_service: DatabaseService = some_tech_db_srv.software_manager.software[\"DatabaseService\"]\n",
"some_tech_db_service: DatabaseService = some_tech_db_srv.software_manager.software[\"database-service\"]\n",
"some_tech_storage_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_storage_srv\")\n",
"some_tech_web_srv: Server = game.simulation.network.get_node_by_hostname(\"some_tech_web_srv\")"
]
@@ -201,7 +211,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n",
" \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n",
"]\n",
"game.simulation.apply_request(caos_action)"
]
@@ -223,7 +233,7 @@
},
"outputs": [],
"source": [
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n",
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n",
"game.simulation.apply_request(caos_action)"
]
},
@@ -246,7 +256,7 @@
"metadata": {},
"outputs": [],
"source": [
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]"
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]"
]
},
{
@@ -259,7 +269,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n",
" \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_rt.network_interface[4].ip_address)\n",
"]\n",
"game.simulation.apply_request(caos_action)"
]
@@ -272,7 +282,7 @@
},
"outputs": [],
"source": [
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]"
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]"
]
},
{
@@ -305,7 +315,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"send_remote_command\", str(some_tech_rt.network_interface[4].ip_address),\n",
" \"service\", \"terminal\", \"send_remote_command\", str(some_tech_rt.network_interface[4].ip_address),\n",
" {\n",
" \"command\": [\n",
" \"acl\", \"add_rule\", \"PERMIT\", \"TCP\",\n",
@@ -358,7 +368,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"remote_logoff\", str(some_tech_rt.network_interface[4].ip_address)\n",
" \"service\", \"terminal\", \"remote_logoff\", str(some_tech_rt.network_interface[4].ip_address)\n",
"]\n",
"game.simulation.apply_request(caos_action)"
]
@@ -376,7 +386,7 @@
"metadata": {},
"outputs": [],
"source": [
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"UserSessionManager\"][\"active_remote_sessions\"]"
"game.get_sim_state()[\"network\"][\"nodes\"][\"some_tech_rt\"][\"services\"][\"user-session-manager\"][\"active_remote_sessions\"]"
]
},
{
@@ -396,7 +406,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"node_session_remote_login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n",
" \"service\", \"terminal\", \"node-session-remote-login\", \"admin\", \"admin\", str(some_tech_storage_srv.network_interface[1].ip_address)\n",
"]\n",
"game.simulation.apply_request(caos_action)"
]
@@ -411,7 +421,7 @@
"source": [
"caos_action = [\n",
" \"network\", \"node\", \"some_tech_jnr_dev_pc\", \n",
" \"service\", \"Terminal\", \"send_remote_command\", str(some_tech_storage_srv.network_interface[1].ip_address),\n",
" \"service\", \"terminal\", \"send_remote_command\", str(some_tech_storage_srv.network_interface[1].ip_address),\n",
" {\n",
" \"command\": [\n",
" \"file_system\", \"delete\", \"file\", db_backup_folder, \"database.db\"\n",
@@ -466,7 +476,7 @@
},
"outputs": [],
"source": [
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n",
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n",
"game.simulation.apply_request(caos_action)"
]
},
@@ -525,7 +535,7 @@
},
"outputs": [],
"source": [
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"WebBrowser\", \"execute\"]\n",
"caos_action = [\"network\", \"node\", \"some_tech_jnr_dev_pc\", \"application\", \"web-browser\", \"execute\"]\n",
"game.simulation.apply_request(caos_action)"
]
},

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

@@ -6,7 +6,7 @@
"source": [
"# Requests and Responses\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"Agents interact with the PrimAITE simulation via the Request system.\n"
]
@@ -52,12 +52,17 @@
"outputs": [],
"source": [
"sim = Simulation()\n",
"\n",
"sim.network.add_node(\n",
" HostNode(\n",
" hostname=\"client\",\n",
" ip_address='10.0.0.1',\n",
" subnet_mask='255.255.255.0',\n",
" operating_state=NodeOperatingState.ON)\n",
" HostNode.from_config(\n",
" config = {\n",
" 'type': \"host-node\",\n",
" 'hostname': \"client\",\n",
" 'ip_address': '10.0.0.1',\n",
" 'subnet_mask': '255.255.255.0',\n",
" 'operating_state': \"ON\",\n",
" }\n",
" )\n",
")\n",
"client = sim.network.get_node_by_hostname('client')\n"
]
@@ -94,7 +99,7 @@
"outputs": [],
"source": [
"response = sim.apply_request(\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"DNSClient\", \"stop\"],\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"dns-client\", \"stop\"],\n",
" context={}\n",
" )\n",
"print(response)"
@@ -114,7 +119,7 @@
"metadata": {},
"outputs": [],
"source": [
"print(f\"DNS Client state: {client.software_manager.software.get('DNSClient').operating_state.name}\")"
"print(f\"DNS Client state: {client.software_manager.software.get('dns-client').operating_state.name}\")"
]
},
{
@@ -138,7 +143,7 @@
"outputs": [],
"source": [
"response = sim.apply_request(\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"NonExistentApplication\", \"stop\"],\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"non-existent-application\", \"stop\"],\n",
" context={}\n",
" )\n",
"print(response)"
@@ -201,7 +206,7 @@
"outputs": [],
"source": [
"response = sim.apply_request(\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"DNSClient\", \"start\"],\n",
" request=[\"network\", \"node\", \"client\", \"service\", \"dns-client\", \"start\"],\n",
" context={}\n",
" )\n",
"print(response)"
@@ -210,7 +215,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},

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

@@ -6,7 +6,7 @@
"source": [
"# Terminal Processing\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK"
]
},
{
@@ -42,9 +42,26 @@
"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 = Computer.from_config(\n",
" config = {\n",
" \"type\": \"computer\",\n",
" \"hostname\": \"node_a\",\n",
" \"ip_address\": \"192.168.0.10\",\n",
" \"subnet_mask\": \"255.255.255.0\",\n",
" # \"startup_duration\": 0,\n",
" }\n",
" )\n",
" print(f\"{node_a=}\")\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 = Computer.from_config(\n",
" config = {\n",
" \"type\": \"computer\",\n",
" \"hostname\": \"node_b\",\n",
" \"ip_address\": \"192.168.0.11\",\n",
" \"subnet_mask\": \"255.255.255.0\",\n",
" # \"startup_duration\": 0,\n",
" }\n",
" )\n",
" node_b.power_on()\n",
" network.connect(node_a.network_interface[1], node_b.network_interface[1])\n",
" return network"
@@ -67,9 +84,9 @@
"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",
"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\")"
"terminal_b: Terminal = computer_b.software_manager.software.get(\"terminal\")"
]
},
{
@@ -121,7 +138,7 @@
"metadata": {},
"outputs": [],
"source": [
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"RansomwareScript\"])"
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"ransomware-script\"])"
]
},
{
@@ -211,7 +228,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -225,7 +242,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

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

@@ -6,7 +6,7 @@
"source": [
"# Train a Multi agent system using RLLIB\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\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."
]
@@ -40,6 +40,7 @@
"import ray\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",
"with open(PRIMAITE_PATHS.user_config_path / 'example_config/data_manipulation_marl.yaml', 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
@@ -110,21 +111,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"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,

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

@@ -6,7 +6,7 @@
"source": [
"# Train a Single agent system using RLLib\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\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."
]
@@ -32,6 +32,8 @@
"from primaite.session.ray_envs import PrimaiteRayEnv\n",
"import ray\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",
@@ -104,21 +106,9 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"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.11"
}
},
"nbformat": 4,

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

@@ -6,7 +6,7 @@
"source": [
"# Training an SB3 Agent\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\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."
]
@@ -35,6 +35,7 @@
"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",
"import yaml"
]
},
@@ -177,7 +178,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -191,7 +192,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

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

@@ -6,7 +6,7 @@
"source": [
"# Using Episode Schedules\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"PrimAITE supports the ability to use different variations on a scenario at different episodes. This can be used to increase \n",
"domain randomisation to prevent overfitting, or to set up curriculum learning to train agents to perform more complicated tasks.\n",
@@ -48,6 +48,7 @@
"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\""
]
},
@@ -409,7 +410,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -423,7 +424,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

49
src/primaite/simulator/_package_data/create-simulation_demo.ipynb → src/primaite/notebooks/create-simulation_demo.ipynb
View File

@@ -6,7 +6,7 @@
"source": [
"# Build a simulation using the Python API\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"Currently, this notebook manipulates the simulation by directly placing objects inside of the attributes of the network and domain. It should be refactored when proper methods exist for adding these objects."
]
@@ -70,9 +70,23 @@
"metadata": {},
"outputs": [],
"source": [
"my_pc = Computer(hostname=\"Computer\", ip_address=\"192.168.1.10\", subnet_mask=\"255.255.255.0\")\n",
"my_pc = Computer.from_config(\n",
" config={\n",
" \"type\": \"computer\",\n",
" \"hostname\":\"pc_1\",\n",
" \"ip_address\":\"192.168.1.10\",\n",
" \"subnet_mask\":\"255.255.255.0\",\n",
" }\n",
" )\n",
"net.add_node(my_pc)\n",
"my_server = Server(hostname=\"Server\", ip_address=\"192.168.1.11\", subnet_mask=\"255.255.255.0\")\n",
"my_server = Server.from_config(\n",
" config={\n",
" \"type\": \"server\",\n",
" \"hostname\":\"Server\",\n",
" \"ip_address\":\"192.168.1.11\",\n",
" \"subnet_mask\":\"255.255.255.0\"\n",
" }\n",
")\n",
"net.add_node(my_server)\n"
]
},
@@ -99,7 +113,13 @@
"metadata": {},
"outputs": [],
"source": [
"my_switch = Switch(hostname=\"switch1\", num_ports=12)\n",
"my_switch = Switch.from_config(\n",
" config = {\n",
" \"type\":\"switch\",\n",
" \"hostname\":\"switch1\",\n",
" \"num_ports\":12\n",
" }\n",
")\n",
"net.add_node(my_switch)\n",
"\n",
"pc_nic = NIC(ip_address=\"130.1.1.1\", gateway=\"130.1.1.255\", subnet_mask=\"255.255.255.0\")\n",
@@ -163,16 +183,30 @@
"metadata": {},
"outputs": [],
"source": [
"from pydantic import Field\n",
"\n",
"from pathlib import Path\n",
"from primaite.simulator.system.applications.application import Application, ApplicationOperatingState\n",
"from primaite.simulator.system.software import SoftwareHealthState, SoftwareCriticality\n",
"from primaite.simulator.file_system.file_system import FileSystem\n",
"from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP\n",
"from primaite.utils.validation.port import PORT_LOOKUP\n",
"from primaite.simulator.system.core.sys_log import SysLog\n",
"\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application, discriminator=\"MSPaint\"):\n",
" class ConfigSchema(Application.ConfigSchema):\n",
" type: str = \"MSPaint\"\n",
"\n",
" config: ConfigSchema = Field(default_factory=lambda: MSPaint.ConfigSchema())\n",
"\n",
" def __init__(self, **kwargs):\n",
" kwargs[\"name\"] = \"MSPaint\"\n",
" kwargs[\"port\"] = PORT_LOOKUP[\"HTTP\"]\n",
" kwargs[\"protocol\"] = PROTOCOL_LOOKUP[\"NONE\"]\n",
" super().__init__(**kwargs)\n",
"\n",
" def describe_state(self):\n",
" return super().describe_state()"
]
@@ -183,7 +217,8 @@
"metadata": {},
"outputs": [],
"source": [
"mspaint = MSPaint(name = \"mspaint\", health_state_actual=SoftwareHealthState.GOOD, health_state_visible=SoftwareHealthState.GOOD, criticality=SoftwareCriticality.MEDIUM, port=PORT_LOOKUP[\"HTTP\"], protocol = PROTOCOL_LOOKUP[\"NONE\"],operating_state=ApplicationOperatingState.RUNNING,execution_control_status='manual', file_system=FileSystem(sys_log=SysLog(hostname=\"Test\"), sim_root=Path(__name__).parent),)"
"my_pc.software_manager.install(MSPaint)\n",
"mspaint = my_pc.software_manager.software.get(\"MSPaint\")"
]
},
{
@@ -250,7 +285,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -264,7 +299,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

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

@@ -6,7 +6,7 @@
"source": [
"# Simple multi-processing demonstration\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This notebook uses SubprocVecEnv from SB3."
]
@@ -37,7 +37,6 @@
"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",
"\n",
"from primaite.session.environment import PrimaiteGymEnv\n"
]
},
@@ -138,7 +137,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -152,7 +151,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

16
src/primaite/simulator/_package_data/network_simulator_demo.ipynb → src/primaite/notebooks/network_simulator_demo.ipynb
View File

@@ -7,7 +7,7 @@
"source": [
"# PrimAITE Router Simulation Demo\n",
"\n",
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK\n",
"\n",
"This demo uses a modified version of the ARCD Use Case 2 Network (seen below) to demonstrate the capabilities of the Network simulator in PrimAITE."
]
@@ -650,21 +650,9 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"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.11"
}
},
"nbformat": 4,

2
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -313,7 +313,7 @@ class HostNode(Node, discriminator="host-node"):
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"HostARP": HostARP,
"host-arp": HostARP,
"icmp": ICMP,
"dns-client": DNSClient,
"ntp-client": NTPClient,

2
src/primaite/simulator/network/networks.py
View File

@@ -295,7 +295,7 @@ def arcd_uc2_network() -> Network:
# Security Suite
security_suite_cfg = {
"type": "server",
"hostname": "backup_server",
"hostname": "security_suite",
"ip_address": "192.168.1.110",
"subnet_mask": "255.255.255.0",
"default_gateway": "192.168.1.1",

2
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -52,6 +52,7 @@ class DataManipulationBot(Application, discriminator="data-manipulation-bot"):
payload: str = "DELETE"
port_scan_p_of_success: float = 0.1
data_manipulation_p_of_success: float = 0.1
repeat: bool = True
config: "DataManipulationBot.ConfigSchema" = Field(default_factory=lambda: DataManipulationBot.ConfigSchema())
@@ -76,6 +77,7 @@ class DataManipulationBot(Application, discriminator="data-manipulation-bot"):
self.payload = self.config.payload
self.port_scan_p_of_success = self.config.port_scan_p_of_success
self.data_manipulation_p_of_success = self.config.data_manipulation_p_of_success
self.repeat = self.config.repeat
def describe_state(self) -> Dict:
"""