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

Merge branch 'dev' into UC7-migration

Browse Source
This commit is contained in:
Archer Bowen
2025-02-27 11:36:27 +00:00
committed by Marek Wolan
parent 2662540965 7b092158db
commit 0b7cdd451b
35 changed files with 331 additions and 337 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')

17
CHANGELOG.md
View File

@@ -42,6 +42,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [3.3.0] - 2024-09-04
## [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)
- Added ability to set the observation threshold for NMNE, file access and application executions
### Changed
- ACL's are no longer applied to layer-2 traffic.
- Random number seed values are recorded in simulation/seed.log if the seed is set in the config file
or `generate_seed_value` is set to `true`.
- ARP .show() method will now include the port number associated with each entry.
- Added `services_requires_scan` and `applications_requires_scan` to agent observation space config to allow the agents to be able to see actual health states of services and applications without requiring scans (Default `True`, set to `False` to allow agents to see actual health state without scanning).
- Updated the `Terminal` class to provide response information when sending remote command execution.
## [3.3.0] - 2024-09-04
### Added
- Random Number Generator Seeding by specifying a random number seed in the config file.
- Implemented Terminal service class, providing a generic terminal simulation.

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

6
docs/source/configuration/agents.rst
View File

@@ -20,12 +20,6 @@ Agents can be scripted (deterministic and stochastic), or controlled by a reinfo
- ref: green_agent_example
team: GREEN
type: probabilistic-agent
observation_space:
type: UC2GreenObservation # TODO: what
action_space:
reward_function:
reward_components:
- type: dummy
agent_settings:
start_settings:

20
docs/source/simulation_components/system/applications/data_manipulation_bot.rst
View File

@@ -97,26 +97,6 @@ If not using the data manipulation bot manually, it needs to be used with a data
team: RED
type: red-database-corrupting-agent
observation_space:
type: uc2-red-observation #TODO what
options:
nodes:
- node_name: client_1
observations:
- logon_status
- operating_status
applications:
- application_ref: data_manipulation_bot
observations:
operating_status
health_status
folders: {}
action_space:
reward_function:
reward_components:
- type: dummy
agent_settings:
start_settings:
start_step: 25

2
docs/source/simulation_components/system/applications/database_client.rst
View File

@@ -59,7 +59,7 @@ Python
# install DatabaseClient
client.software_manager.install(DatabaseClient)
database_client: DatabaseClient = client.software_manager.software.get("database-sclient")
database_client: DatabaseClient = client.software_manager.software.get("database-client")
# Configure the DatabaseClient
database_client.configure(server_ip_address=IPv4Address("192.168.0.1")) # address of the DatabaseService

2
src/primaite/VERSION
View File

@@ -1 +1 @@
4.0.0a1-dev
4.0.0-dev

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:

28
src/primaite/game/agent/interface.py
View File

@@ -107,35 +107,27 @@ class AbstractAgent(BaseModel, ABC):
self.reward_function = RewardFunction(config=self.config.reward_function)
return super().model_post_init(__context)
def add_agent_action(self, item: AgentHistoryItem, table: PrettyTable) -> PrettyTable:
"""Update the given table with information from given AgentHistoryItem."""
node, application = "unknown", "unknown"
if (node_id := item.parameters.get("node_id")) is not None:
node = self.action_manager.node_names[node_id]
if (application_id := item.parameters.get("application_id")) is not None:
application = self.action_manager.application_names[node_id][application_id]
if (application_name := item.parameters.get("application_name")) is not None:
application = application_name
table.add_row([item.timestep, item.action, node, application, item.response.status])
return table
def show_history(self, ignored_actions: Optional[list] = None):
"""
Print an agent action provided it's not the DONOTHING action.
Print an agent action provided it's not the do-nothing action.
:param ignored_actions: OPTIONAL: List of actions to be ignored when displaying the history.
If not provided, defaults to ignore DONOTHING actions.
If not provided, defaults to ignore do-nothing actions.
"""
if not ignored_actions:
ignored_actions = ["DONOTHING"]
ignored_actions = ["do-nothing"]
table = PrettyTable()
table.field_names = ["Step", "Action", "Node", "Application", "Response"]
print(f"Actions for '{self.agent_name}':")
table.field_names = ["Step", "Action", "Params", "Response", "Response Data"]
print(f"Actions for '{self.config.ref}':")
for item in self.history:
if item.action in ignored_actions:
pass
else:
table = self.add_agent_action(item=item, table=table)
# format dict by putting each key-value entry on a separate line and putting a blank line on the end.
param_string = "\n".join([*[f"{k}: {v:.30}" for k, v in item.parameters.items()], ""])
data_string = "\n".join([*[f"{k}: {v:.30}" for k, v in item.response.data], ""])
table.add_row([item.timestep, item.action, param_string, item.response.status, data_string])
print(table)
def update_observation(self, state: Dict) -> ObsType:

4
src/primaite/game/agent/observations/observation_manager.py
View File

@@ -230,7 +230,7 @@ class ObservationManager(BaseModel):
return self.obs.space
@classmethod
def from_config(cls, config: Optional[Dict]) -> "ObservationManager":
def from_config(cls, config: Optional[Dict], thresholds: Optional[Dict] = {}) -> "ObservationManager":
"""
Create observation space from a config.
@@ -241,6 +241,8 @@ class ObservationManager(BaseModel):
AbstractObservation
options: this must adhere to the chosen observation type's ConfigSchema nested class.
:type config: Dict
:param thresholds: Dictionary containing the observation thresholds.
:type thresholds: Optional[Dict]
"""
if config is None:
return cls(NullObservation())

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,

211
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,12 +59,12 @@
"custom_c2_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: proxy-a.gent\n",
" type: proxy-agent\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-install\n",
@@ -66,37 +75,31 @@
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\n",
" config:\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\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_name: web_server\n",
" application_name: c2-beacon\n",
" 4:\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_name: client_1\n",
" config:\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" 6:\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
@@ -105,9 +108,8 @@
" 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",
@@ -117,20 +119,19 @@
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\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",
" 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",
" options:\n",
" node_name: web_server\n",
" config:\n",
" c2_server_ip_address: 192.168.10.22\n",
" keep_alive_frequency:\n",
" masquerade_protocol:\n",
" masquerade_port:\n",
" c2_server_ip_address: 192.168.10.22\n",
"\n",
" reward_function:\n",
" reward_components:\n",
" - type: dummy\n",
"\"\"\"\n",
"c2_agent_yaml = yaml.safe_load(custom_c2_agent)"
]
@@ -171,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()"
]
@@ -208,7 +209,7 @@
" ...\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: c2-beacon\n",
@@ -237,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",
" 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",
"```"
]
},
@@ -275,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: c2-beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
" ...\n",
" 3:\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",
"```"
]
},
@@ -333,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: c2-server\n",
" ...\n",
" action_map:\n",
" 4:\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",
@@ -368,7 +348,7 @@
" - software_manager\n",
" - application\n",
" - install\n",
" - RansomwareScript\n",
" - ransomware-script\n",
"```"
]
},
@@ -404,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: c2-server\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"
]
},
@@ -446,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",
@@ -456,13 +428,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 6:\n",
" action: c2-server-data-exfiltrate\n",
@@ -472,9 +437,8 @@
" 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",
"```"
]
@@ -512,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",
@@ -522,13 +486,6 @@
"\n",
"``` yaml\n",
" action_space:\n",
" options:\n",
" nodes: # Node List\n",
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 7:\n",
" action: c2-server-ransomware-launch\n",
@@ -583,7 +540,7 @@
" options:\n",
" components:\n",
" - type: nodes\n",
" label: NODES\n",
" label: nodes\n",
" options:\n",
" hosts:\n",
" - hostname: web_server\n",
@@ -660,30 +617,34 @@
" 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",
" options:\n",
" node_name: web-server\n",
" application_name: C2Beacon\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
" 2:\n",
" action: node-shutdown\n",
" options:\n",
" node_name: web-server\n",
" node_name: web_server\n",
" 3:\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",
" reward_function:\n",
" reward_components:\n",
" - type: dummy\n",
"\n",
" agent_settings:\n",
" flatten_obs: False\n",
@@ -774,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()"
]
@@ -1335,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",
"```"
]
},

70
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",
@@ -79,8 +79,7 @@
" 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",
" red_str = f\"ATTACK from {red_info.parameters['node_name']}\"\n",
" return red_str"
]
},
@@ -149,18 +148,13 @@
" team: RED # not used, just for human reference\n",
" type: red-database-corrupting-agent # 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: uc2-red-observation # TODO: what\n",
" options:\n",
" nodes: {}\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",
"```"
]
},
@@ -180,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",
@@ -189,13 +182,13 @@
" \n",
" # \n",
" applications:\n",
" - type: data-manipulation-bot\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",
" - type: database-client # 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",
"```"
@@ -219,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",
@@ -229,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",
@@ -286,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",
@@ -330,16 +338,19 @@
"# Make attack always succeed.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - type: data-manipulation-bot\n",
" - ref: data_manipulation_bot\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",
" - type: web-browser\n",
" - ref: client_1_web_browser\n",
" type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - type: database-client\n",
" - ref: client_1_database_client\n",
" type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
@@ -372,16 +383,19 @@
"# Make attack always fail.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - type: data-manipulation-bot\n",
" - ref: data_manipulation_bot\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",
" - type: web-browser\n",
" - ref: client_1_web_browser\n",
" type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - type: database-client\n",
" - ref: client_1_database_client\n",
" type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
@@ -408,7 +422,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -422,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"
]
},
{
@@ -200,7 +200,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -214,7 +214,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"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)"
]
},

25
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)"
@@ -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"
},

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

@@ -6,7 +6,14 @@
"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"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
@@ -49,9 +56,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"
@@ -281,21 +305,7 @@
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: proxy-agent\n",
" observation_space: null\n",
" action_space:\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: do-nothing\n",
@@ -491,7 +501,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "venv",
"language": "python",
"name": "python3"
},
@@ -505,7 +515,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

@@ -350,7 +350,7 @@ class HostNode(Node, discriminator="host-node"):
self.connect_nic(NIC(ip_address=kwargs["config"].ip_address, subnet_mask=kwargs["config"].subnet_mask))
for folder in self.config.folders:
# handle empty folder defined by just a string
# handle empty foler defined by just a string
self.file_system.create_folder(folder["folder_name"])
for file in folder.get("files", []):

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:
"""

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

@@ -415,5 +415,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, session.protocol])
table.add_row([session.with_ip_address, session.dst_port, session.protocol])
print(table)

2
tests/integration_tests/configuration_file_parsing/test_node_file_system_config.py
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from pathlib import Path
from typing import Union

2
tests/integration_tests/game_layer/actions/test_user_account_actions.py
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
import pytest
from primaite.simulator.network.hardware.nodes.host.computer import Computer

8
tests/integration_tests/game_layer/observations/test_nic_observations.py
View File

@@ -85,6 +85,14 @@ def test_nic(simulation):
# in the NICObservation class so we set it now.
nic_obs.capture_nmne = True
# The Simulation object created by the fixture also creates the
# NICObservation class with the NICObservation.capture_nmnme class variable
# set to False. Under normal (non-test) circumstances this class variable
# is set from a config file such as data_manipulation.yaml. So although
# capture_nmne is set to True in the NetworkInterface class it's still False
# in the NICObservation class so we set it now.
nic_obs.capture_nmne = True
# Set the NMNE configuration to capture DELETE/ENCRYPT queries as MNEs
nmne_config = {
"capture_nmne": True, # Enable the capture of MNEs

2
tests/integration_tests/game_layer/observations/test_obs_data_capture.py
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
import json
from primaite.session.environment import PrimaiteGymEnv

6
tests/integration_tests/game_layer/test_actions.py
View File

@@ -108,7 +108,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
"""
Test that the RouterACLAddRuleAction can form a request and that it is accepted by the simulation.
The ACL starts off with 4 rules, and we add a rule, and check that the ACL now has 5 rules.
The ACL starts off with 3 rules, and we add a rule, and check that the ACL now has 4 rules.
"""
game, agent = game_and_agent
@@ -140,7 +140,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
agent.store_action(action)
game.step()
# 3: Check that the acl now has 5 rules, and that client 1 cannot ping server 2
# 3: Check that the acl now has 6 rules, and that client 1 cannot ping server 2
assert router.acl.num_rules == 5
assert not client_1.ping("10.0.2.3") # Cannot ping server_2
assert client_1.ping("10.0.2.2") # Can ping server_1
@@ -198,7 +198,7 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P
agent.store_action(action)
game.step()
# 3: Check that the ACL now has 3 rules, and that client 1 cannot access example.com
# 3: Check that the ACL now has 2 rules, and that client 1 cannot access example.com
assert router.acl.num_rules == 3
assert not browser.get_webpage()
client_1.software_manager.software.get("dns-client").dns_cache.clear()

2
tests/integration_tests/system/test_arp.py
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router, RouterARP
from primaite.simulator.system.services.arp.arp import ARP
from primaite.utils.validation.port import PORT_LOOKUP