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

Merge remote-tracking branch 'origin/4.0.0a1-dev' into 4.0.0-dev

Browse Source
This commit is contained in:
Marek Wolan
2025-02-26 15:33:50 +00:00
parent 9bb4442bbe 502d163818
commit 2bf118abc9
71 changed files with 464 additions and 310 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

3
src/primaite/config/_package_data/basic_lan_network_example.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
ports:
- ARP

3
src/primaite/config/_package_data/client_server_p2p_network_example.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
ports:
- ARP

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

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false
@@ -81,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:

3
src/primaite/config/_package_data/data_manipulation_marl.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

3
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/base_scenario.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
max_episode_length: 128
ports: []

3
src/primaite/config/_package_data/multi_lan_internet_network_example.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
ports:
- ARP

3
src/primaite/config/_package_data/scenario_with_placeholders/scenario.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

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

@@ -77,7 +77,7 @@ class AbstractAgent(BaseModel, ABC):
config: ConfigSchema = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
logger: AgentLog = AgentLog(agent_name="Abstract_Agent")
logger: AgentLog = None
history: List[AgentHistoryItem] = []
action_manager: ActionManager = Field(default_factory=lambda: ActionManager())
@@ -86,6 +86,11 @@ class AbstractAgent(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
def __init__(self, **kwargs):
"""Initialise and setup agent logger."""
super().__init__(**kwargs)
self.logger: AgentLog = AgentLog(agent_name=kwargs["config"]["ref"])
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if discriminator is None:

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,

164
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: ProxyAgent\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",
@@ -65,35 +74,34 @@
" 2:\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\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",
" 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",
" node_id: 1\n",
" config:\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
@@ -105,9 +113,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",
@@ -116,7 +123,7 @@
" 8:\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\n",
" node_id: 0\n",
" config:\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency: 10\n",
@@ -125,12 +132,16 @@
" 9:\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_name: web_server\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",
"\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
"\"\"\"\n",
"c2_agent_yaml = yaml.safe_load(custom_c2_agent)"
]
@@ -171,7 +182,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 +219,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 +248,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,7 +281,7 @@
"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",
@@ -285,16 +291,15 @@
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: c2-beacon\n",
" - application_name: C2Beacon\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,15 +338,15 @@
"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",
@@ -352,15 +357,14 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\n",
" - application_name: C2Server\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 +372,7 @@
" - software_manager\n",
" - application\n",
" - install\n",
" - RansomwareScript\n",
" - ransomware-script\n",
"```"
]
},
@@ -409,16 +413,15 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\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"
]
},
@@ -446,9 +449,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",
@@ -461,7 +464,7 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 6:\n",
@@ -472,9 +475,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 +514,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",
@@ -527,7 +529,7 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: c2-server\n",
" - application_name: C2Server\n",
" ...\n",
" action_map:\n",
" 7:\n",
@@ -582,7 +584,7 @@
" type: custom\n",
" options:\n",
" components:\n",
" - type: nodes\n",
" - type: NODES\n",
" label: NODES\n",
" options:\n",
" hosts:\n",
@@ -660,17 +662,17 @@
" 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",
" node_id: 0\n",
" application_name: C2Beacon\n",
" 2:\n",
" action: node-shutdown\n",
" options:\n",
" node_name: web-server\n",
" node_id: 0\n",
" 3:\n",
" action: router-acl-add-rule\n",
" options:\n",
@@ -685,6 +687,36 @@
" source_wildcard_id: 0\n",
" dest_wildcard_id: 0\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",
"\n",
" agent_settings:\n",
" flatten_obs: False\n",
"\"\"\"\n",
@@ -774,12 +806,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 +1367,18 @@
"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",
" 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",
" 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"
},

27
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"
]
},
{
@@ -49,9 +49,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"
@@ -491,7 +508,7 @@
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
@@ -505,7 +522,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,

4
src/primaite/simulator/network/airspace.py
View File

@@ -449,10 +449,8 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
`default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption.
"""
super().enable()
try:
if hasattr(self._connected_node, "default_gateway_hello"):
self._connected_node.default_gateway_hello()
except AttributeError:
pass
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:

70
src/primaite/simulator/network/container.py
View File

@@ -163,16 +163,6 @@ class Network(SimComponent):
:param links: Include link details in the output. Defaults to True.
:param markdown: Use Markdown style in table output. Defaults to False.
"""
nodes_type_map = {
"Router": self.router_nodes,
"Firewall": self.firewall_nodes,
"Switch": self.switch_nodes,
"Server": self.server_nodes,
"Computer": self.computer_nodes,
"Printer": self.printer_nodes,
"Wireless Router": self.wireless_router_nodes,
}
if nodes:
table = PrettyTable(["Node", "Type", "Operating State"])
if markdown:
@@ -189,21 +179,20 @@ class Network(SimComponent):
table.set_style(MARKDOWN)
table.align = "l"
table.title = "IP Addresses"
for nodes in nodes_type_map.values():
for node in nodes:
for i, port in node.network_interface.items():
if hasattr(port, "ip_address"):
if port.ip_address != IPv4Address("127.0.0.1"):
port_str = port.port_name if port.port_name else port.port_num
table.add_row(
[
node.config.hostname,
port_str,
port.ip_address,
port.subnet_mask,
node.config.default_gateway,
]
)
for node in self.nodes.values():
for i, port in node.network_interface.items():
if hasattr(port, "ip_address"):
if port.ip_address != IPv4Address("127.0.0.1"):
port_str = port.port_name if port.port_name else port.port_num
table.add_row(
[
node.config.hostname,
port_str,
port.ip_address,
port.subnet_mask,
node.config.default_gateway,
]
)
print(table)
if links:
@@ -215,22 +204,21 @@ class Network(SimComponent):
table.align = "l"
table.title = "Links"
links = list(self.links.values())
for nodes in nodes_type_map.values():
for node in nodes:
for link in links[::-1]:
if node in [link.endpoint_a.parent, link.endpoint_b.parent]:
table.add_row(
[
link.endpoint_a.parent.config.hostname,
str(link.endpoint_a),
link.endpoint_b.parent.config.hostname,
str(link.endpoint_b),
link.is_up,
link.bandwidth,
link.current_load_percent,
]
)
links.remove(link)
for node in self.nodes.values():
for link in links[::-1]:
if node in [link.endpoint_a.parent, link.endpoint_b.parent]:
table.add_row(
[
link.endpoint_a.parent.config.hostname,
str(link.endpoint_a),
link.endpoint_b.parent.config.hostname,
str(link.endpoint_b),
link.is_up,
link.bandwidth,
link.current_load_percent,
]
)
links.remove(link)
print(table)
def clear_links(self):

15
src/primaite/simulator/network/creation.py
View File

@@ -155,7 +155,12 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
# Create a core switch if more than one edge switch is needed
if num_of_switches > 1:
core_switch = Switch.from_config(
config={"type": "switch", "hostname": f"switch_core_{config.lan_name}", "start_up_duration": 0}
config={
"type": "switch",
"hostname": f"switch_core_{config.lan_name}",
"start_up_duration": 0,
"num_ports": 24,
}
)
core_switch.power_on()
network.add_node(core_switch)
@@ -183,7 +188,12 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
switch_port = 0
switch_n = 1
switch = Switch.from_config(
config={"type": "switch", "hostname": f"switch_edge_{switch_n}_{config.lan_name}", "start_up_duration": 0}
config={
"type": "switch",
"hostname": f"switch_edge_{switch_n}_{config.lan_name}",
"start_up_duration": 0,
"num_ports": 24,
}
)
switch.power_on()
network.add_node(switch)
@@ -207,6 +217,7 @@ class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
"type": "switch",
"hostname": f"switch_edge_{switch_n}_{config.lan_name}",
"start_up_duration": 0,
"num_ports": 24,
}
)
switch.power_on()

8
src/primaite/simulator/network/hardware/base.py
View File

@@ -639,10 +639,8 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC):
`default_gateway_hello` method is not defined, ignoring such errors to proceed without interruption.
"""
super().enable()
try:
if hasattr(self._connected_node, "default_gateway_hello"):
self._connected_node.default_gateway_hello()
except AttributeError:
pass
return True
@abstractmethod
@@ -1211,7 +1209,7 @@ class UserSessionManager(Service, discriminator="user-session-manager"):
"""Request should take the form [username, password, remote_ip_address]."""
username, password, remote_ip_address = request
response = RequestResponse.from_bool(self.remote_login(username, password, remote_ip_address))
response.data = {"remote_hostname": self.parent.hostname, "username": username}
response.data = {"remote_hostname": self.parent.config.hostname, "username": username}
return response
rm.add_request("remote_login", RequestType(func=_remote_login))
@@ -1244,7 +1242,7 @@ class UserSessionManager(Service, discriminator="user-session-manager"):
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.parent.hostname} User Sessions"
table.title = f"{self.parent.config.hostname} User Sessions"
def _add_session_to_table(user_session: UserSession):
"""

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

@@ -334,7 +334,7 @@ class HostNode(Node, discriminator="host-node"):
class ConfigSchema(Node.ConfigSchema):
"""Configuration Schema for HostNode class."""
type: Literal["host-node"]
type: Literal["host-node"] = "host-node"
hostname: str = "HostNode"
subnet_mask: IPV4Address = "255.255.255.0"
ip_address: IPV4Address
@@ -388,7 +388,7 @@ class HostNode(Node, discriminator="host-node"):
This method is invoked to ensure the host node can communicate with its default gateway, primarily to confirm
network connectivity and populate the ARP cache with the gateway's MAC address.
"""
if self.operating_state == NodeOperatingState.ON and self.default_gateway:
if self.operating_state == NodeOperatingState.ON and self.config.default_gateway:
self.software_manager.arp.get_default_gateway_mac_address()
def receive_frame(self, frame: Frame, from_network_interface: NIC):

1
src/primaite/simulator/network/hardware/nodes/network/router.py
View File

@@ -1352,7 +1352,6 @@ class Router(NetworkNode, discriminator="router"):
:return: A dictionary representing the current state.
"""
state = super().describe_state()
state["num_ports"] = self.config.num_ports
state["acl"] = self.acl.describe_state()
return state

5
src/primaite/simulator/network/hardware/nodes/network/switch.py
View File

@@ -103,8 +103,8 @@ class Switch(NetworkNode, discriminator="switch"):
type: Literal["switch"] = "switch"
hostname: str = "Switch"
num_ports: int = 24
"The number of ports on the switch. Default is 24."
num_ports: int = 8
"The number of ports on the switch."
config: ConfigSchema = Field(default_factory=lambda: Switch.ConfigSchema())
@@ -139,7 +139,6 @@ class Switch(NetworkNode, discriminator="switch"):
"""
state = super().describe_state()
state["ports"] = {port_num: port.describe_state() for port_num, port in self.network_interface.items()}
state["num_ports"] = self.config.num_ports # redundant?
state["mac_address_table"] = {mac: port.port_num for mac, port in self.mac_address_table.items()}
return state

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)

4
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -138,8 +138,8 @@ class ARP(Service, discriminator="arp"):
break
if use_default_gateway:
if self.software_manager.node.default_gateway:
target_ip_address = self.software_manager.node.default_gateway
if self.software_manager.node.config.default_gateway:
target_ip_address = self.software_manager.node.config.default_gateway
else:
return

2
src/primaite/simulator/system/software.py
View File

@@ -82,7 +82,7 @@ class Software(SimComponent, ABC):
"""Configurable options for all software."""
model_config = ConfigDict(extra="forbid")
starting_health_state: SoftwareHealthState = SoftwareHealthState.UNUSED
starting_health_state: SoftwareHealthState = SoftwareHealthState.GOOD
criticality: SoftwareCriticality = SoftwareCriticality.LOWEST
fixing_duration: int = 2

3
tests/assets/configs/action_penalty.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: false
save_step_metadata: false

3
tests/assets/configs/bad_primaite_session.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
ports:
- ARP

3
tests/assets/configs/basic_c2_setup.yaml
View File

@@ -4,6 +4,9 @@
# | node_a |------| switch_1 |------| node_b |
# -------------- -------------- --------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

2
tests/assets/configs/basic_firewall.yaml
View File

@@ -4,6 +4,8 @@
# | client_1 |------| switch_1 |------| client_2 |
# -------------- -------------- --------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false

3
tests/assets/configs/basic_node_with_software_listening_ports.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/basic_node_with_users.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/basic_switched_network.yaml
View File

@@ -4,6 +4,9 @@
# | client_1 |------| switch_1 |------| client_2 |
# -------------- -------------- --------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/data_manipulation.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

3
tests/assets/configs/dmz_network.yaml
View File

@@ -30,6 +30,9 @@
# | external_computer |------| switch_3 |------| external_server |
# ----------------------- -------------- ---------------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/eval_only_primaite_session.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
ports:
- ARP

3
tests/assets/configs/extended_config.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

3
tests/assets/configs/firewall_actions_network.yaml
View File

@@ -30,6 +30,9 @@
# | external_computer |------| switch_3 |------| external_server |
# ----------------------- -------------- ---------------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/fixing_duration_one_item.yaml
View File

@@ -4,6 +4,9 @@
# | client_1 |------| switch_1 |------| client_2 |
# -------------- -------------- --------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/install_and_configure_apps.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: false

3
tests/assets/configs/multi_agent_session.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: false
save_step_metadata: false

3
tests/assets/configs/nmap_network_service_recon_red_agent_config.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/nmap_ping_scan_red_agent_config.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/nmap_port_scan_red_agent_config.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/no_nodes_links_agents_network.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/scenario_with_placeholders/scenario.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

3
tests/assets/configs/shared_rewards.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: false
save_step_metadata: false

3
tests/assets/configs/software_fixing_duration.yaml
View File

@@ -4,6 +4,9 @@
# | client_1 |------| switch_1 |------| client_2 |
# -------------- -------------- --------------
#
metadata:
version: 3.0
io_settings:
save_step_metadata: false
save_pcap_logs: true

3
tests/assets/configs/test_application_install.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: false

3
tests/assets/configs/test_primaite_session.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
io_settings:
save_agent_actions: true
save_step_metadata: true

3
tests/assets/configs/wireless_wan_network_config.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
max_episode_length: 256
ports:

3
tests/assets/configs/wireless_wan_network_config_freq_max_override.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
max_episode_length: 256
ports:

3
tests/assets/configs/wireless_wan_network_config_freq_max_override_blocked.yaml
View File

@@ -1,3 +1,6 @@
metadata:
version: 3.0
game:
max_episode_length: 256
ports:

21
tests/integration_tests/game_layer/test_action_shapes.py
View File

@@ -1,21 +0,0 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from typing import Tuple
from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from tests import TEST_ASSETS_ROOT
FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.yaml"
def test_router_acl_add_rule_action_shape(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test to check ROUTER_ADD_ACL_RULE has the expected action shape."""
game, agent = game_and_agent
# assert that the shape of the actions is correct
router_acl_add_rule_action = agent.action_manager.actions.get("ROUTER_ACL_ADDRULE")
assert router_acl_add_rule_action.shape.get("source_ip_id") == len(agent.action_manager.ip_address_list)
assert router_acl_add_rule_action.shape.get("dest_ip_id") == len(agent.action_manager.ip_address_list)
assert router_acl_add_rule_action.shape.get("source_port_id") == len(agent.action_manager.ports)
assert router_acl_add_rule_action.shape.get("dest_port_id") == len(agent.action_manager.ports)
assert router_acl_add_rule_action.shape.get("protocol_id") == len(agent.action_manager.protocols)

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

@@ -117,7 +117,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
server_1 = game.simulation.network.get_node_by_hostname("server_1")
server_2 = game.simulation.network.get_node_by_hostname("server_2")
router = game.simulation.network.get_node_by_hostname("router")
assert router.acl.num_rules == 3
assert router.acl.num_rules == 4
assert client_1.ping("10.0.2.3") # client_1 can ping server_2
assert server_2.ping("10.0.1.2") # server_2 can ping client_1
@@ -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
@@ -167,8 +167,8 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
agent.store_action(action)
game.step()
# 5: Check that the ACL now has 5 rules, but that server_1 can still ping server_2
assert router.acl.num_rules == 5
# 5: Check that the ACL now has 6 rules, but that server_1 can still ping server_2
assert router.acl.num_rules == 6
assert server_1.ping("10.0.2.3") # Can ping server_2
@@ -199,7 +199,7 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P
game.step()
# 3: Check that the ACL now has 2 rules, and that client 1 cannot access example.com
assert router.acl.num_rules == 2
assert router.acl.num_rules == 3
assert not browser.get_webpage()
client_1.software_manager.software.get("dns-client").dns_cache.clear()
assert client_1.ping("10.0.2.2") # pinging still works because ICMP is allowed

2
tests/unit_tests/_primaite/_game/_agent/test_agent.py
View File

@@ -4,7 +4,7 @@ from primaite.game.agent.scripted_agents.random_agent import RandomAgent
def test_creating_empty_agent():
agent = RandomAgent()
agent = RandomAgent(config={"ref": "Empty Agent"})
assert len(agent.action_manager.action_map) == 0
assert isinstance(agent.observation_manager.obs, NullObservation)
assert len(agent.reward_function.reward_components) == 0

1
tests/unit_tests/_primaite/_simulator/_network/_hardware/nodes/test_switch.py
View File

@@ -17,4 +17,3 @@ def switch() -> Switch:
def test_describe_state(switch):
state = switch.describe_state()
assert len(state.get("ports")) is 8
assert state.get("num_ports") is 8

6
tests/unit_tests/_primaite/_simulator/_system/_applications/test_applications.py
View File

@@ -18,7 +18,7 @@ def test_scan(application):
def test_run_application(application):
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_actual == SoftwareHealthState.UNUSED
assert application.health_state_actual == SoftwareHealthState.GOOD
application.run()
assert application.operating_state == ApplicationOperatingState.RUNNING
@@ -37,9 +37,9 @@ def test_close_application(application):
def test_application_describe_states(application):
assert application.operating_state == ApplicationOperatingState.CLOSED
assert application.health_state_actual == SoftwareHealthState.UNUSED
assert application.health_state_actual == SoftwareHealthState.GOOD
assert SoftwareHealthState.UNUSED.value == application.describe_state().get("health_state_actual")
assert SoftwareHealthState.GOOD.value == application.describe_state().get("health_state_actual")
application.run()
assert SoftwareHealthState.GOOD.value == application.describe_state().get("health_state_actual")

12
tests/unit_tests/_primaite/_simulator/_system/_services/test_services.py
View File

@@ -22,7 +22,7 @@ def test_scan(service):
def test_start_service(service):
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.start()
assert service.operating_state == ServiceOperatingState.RUNNING
@@ -43,7 +43,7 @@ def test_pause_and_resume_service(service):
assert service.operating_state == ServiceOperatingState.STOPPED
service.resume()
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
@@ -58,11 +58,11 @@ def test_pause_and_resume_service(service):
def test_restart(service):
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.restart()
# Service is STOPPED. Restart will only work if the service was PAUSED or RUNNING
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.start()
assert service.operating_state == ServiceOperatingState.RUNNING
@@ -157,11 +157,11 @@ def test_service_fixing(service):
def test_enable_disable(service):
service.disable()
assert service.operating_state == ServiceOperatingState.DISABLED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
service.enable()
assert service.operating_state == ServiceOperatingState.STOPPED
assert service.health_state_actual == SoftwareHealthState.UNUSED
assert service.health_state_actual == SoftwareHealthState.GOOD
def test_overwhelm_service(service):

4
tests/unit_tests/_primaite/_simulator/_system/test_software.py
View File

@@ -39,6 +39,6 @@ def test_software_creation(software):
def test_software_set_health_state(software):
assert software.health_state_actual == SoftwareHealthState.UNUSED
software.set_health_state(SoftwareHealthState.GOOD)
assert software.health_state_actual == SoftwareHealthState.GOOD
software.set_health_state(SoftwareHealthState.COMPROMISED)
assert software.health_state_actual == SoftwareHealthState.COMPROMISED