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

3086 UC7 Migration - All YAMLS, tests and notebooks. A few lingering issues such as the OS-SCAN not working and agent logs not appearing.

Browse Source
This commit is contained in:
Marek Wolan
2025-02-10 14:39:28 +00:00
parent 0d1edf0362
commit d8c8aa40a4
118 changed files with 21789 additions and 368 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
CHANGELOG.md
View File

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [4.0.0] = TBC
### 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, `do-nothing` 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
- Agents now follow a common configuration format, simplifying the configuration of agents and their extensibilty.
@@ -24,6 +28,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Updated tests that don't use YAMLs to still use the new action and agent schemas
- Nodes now use a config schema and are extensible, allowing for plugin support.
- Node tests have been updated to use the new node config schemas when not using YAML files.
- ACLs 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.
### Fixed
- DNS client no longer fails to check its cache if a DNS server address is missing.

2
docs/Makefile
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# Minimal makefile for Sphinx documentation
# You can set these variables from the command line, and also
# from the environment for the first two.

2
docs/build-sphinx-docs-to-github-pages.sh
View File

@@ -1,4 +1,4 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
#!/bin/bash
set -x

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

@@ -21,7 +21,7 @@ Agents can be scripted (deterministic and stochastic), or controlled by a reinfo
team: GREEN
type: probabilistic-agent
observation_space:
type: UC2GreenObservation
type: UC2GreenObservation # TODO: what
action_space:
reward_function:
reward_components:
@@ -160,3 +160,4 @@ If ``True``, gymnasium flattening will be performed on the observation space bef
-----------------
Agents will record their action log for each step. This is a summary of what the agent did, along with response information from requests within the simulation.
A summary of the actions taken by the agent can be viewed using the `show_history()` function. By default, this will display all actions taken apart from ``DONOTHING``.

33
docs/source/configuration/simulation/nodes/common/common_node_attributes.rst
View File

@@ -54,6 +54,39 @@ Optional. Default value is ``3``.
The number of time steps required to occur in order for the node to cycle from ``ON`` to ``SHUTTING_DOWN`` and then finally ``OFF``.
``file_system``
---------------
Optional.
The file system of the node. This configuration allows nodes to be initialised with files and/or folders.
The file system takes a list of folders and files.
Example:
.. code-block:: yaml
simulation:
network:
nodes:
- hostname: client_1
type: computer
ip_address: 192.168.10.11
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
file_system:
- empty_folder # example of an empty folder
- downloads:
- "test_1.txt" # files in the downloads folder
- "test_2.txt"
- root:
- passwords: # example of file with size and type
size: 69 # size in bytes
type: TXT # See FileType for list of available file types
List of file types: :py:mod:`primaite.simulator.file_system.file_type.FileType`
``users``
---------

12
docs/source/configuration/simulation/nodes/network_examples.rst
View File

@@ -1177,8 +1177,8 @@ ACLs permitting or denying traffic as per our configured ACL rules.
some_tech_storage_srv = network.get_node_by_hostname("some_tech_storage_srv")
some_tech_storage_srv.file_system.create_file(file_name="test.png")
pc_1_ftp_client: FTPClient = network.get_node_by_hostname("pc_1").software_manager.software["FTPClient"]
pc_2_ftp_client: FTPClient = network.get_node_by_hostname("pc_2").software_manager.software["FTPClient"]
pc_1_ftp_client: FTPClient = network.get_node_by_hostname("pc_1").software_manager.software["ftp-client"]
pc_2_ftp_client: FTPClient = network.get_node_by_hostname("pc_2").software_manager.software["ftp-client"]
assert not pc_1_ftp_client.request_file(
dest_ip_address=some_tech_storage_srv.network_interface[1].ip_address,
@@ -1224,7 +1224,7 @@ ACLs permitting or denying traffic as per our configured ACL rules.
web_server: Server = network.get_node_by_hostname("some_tech_web_srv")
web_ftp_client: FTPClient = web_server.software_manager.software["FTPClient"]
web_ftp_client: FTPClient = web_server.software_manager.software["ftp-client"]
assert not web_ftp_client.request_file(
dest_ip_address=some_tech_storage_srv.network_interface[1].ip_address,
@@ -1269,7 +1269,7 @@ ACLs permitting or denying traffic as per our configured ACL rules.
some_tech_storage_srv.file_system.create_file(file_name="test.png")
some_tech_snr_dev_pc: Computer = network.get_node_by_hostname("some_tech_snr_dev_pc")
snr_dev_ftp_client: FTPClient = some_tech_snr_dev_pc.software_manager.software["FTPClient"]
snr_dev_ftp_client: FTPClient = some_tech_snr_dev_pc.software_manager.software["ftp-client"]
assert snr_dev_ftp_client.request_file(
dest_ip_address=some_tech_storage_srv.network_interface[1].ip_address,
@@ -1294,7 +1294,7 @@ ACLs permitting or denying traffic as per our configured ACL rules.
some_tech_storage_srv.file_system.create_file(file_name="test.png")
some_tech_jnr_dev_pc: Computer = network.get_node_by_hostname("some_tech_jnr_dev_pc")
jnr_dev_ftp_client: FTPClient = some_tech_jnr_dev_pc.software_manager.software["FTPClient"]
jnr_dev_ftp_client: FTPClient = some_tech_jnr_dev_pc.software_manager.software["ftp-client"]
assert not jnr_dev_ftp_client.request_file(
dest_ip_address=some_tech_storage_srv.network_interface[1].ip_address,
@@ -1337,7 +1337,7 @@ ACLs permitting or denying traffic as per our configured ACL rules.
some_tech_storage_srv.file_system.create_file(file_name="test.png")
some_tech_hr_pc: Computer = network.get_node_by_hostname("some_tech_hr_1")
hr_ftp_client: FTPClient = some_tech_hr_pc.software_manager.software["FTPClient"]
hr_ftp_client: FTPClient = some_tech_hr_pc.software_manager.software["ftp-client"]
assert not hr_ftp_client.request_file(
dest_ip_address=some_tech_storage_srv.network_interface[1].ip_address,

6
docs/source/configuration/simulation/nodes/router.rst
View File

@@ -74,7 +74,7 @@ The subnet mask setting for the port.
``acl``
-------
Sets up the ACL rules for the router.
Sets up the ACL rules for the router to apply to layer-3 traffic. These are not applied to layer-2 traffic such as ARP.
e.g.
@@ -85,10 +85,6 @@ e.g.
...
acl:
1:
action: PERMIT
src_port: ARP
dst_port: ARP
2:
action: PERMIT
protocol: ICMP

6
docs/source/how_to_guides/extensible_agents.rst
View File

@@ -46,17 +46,13 @@ The core features that should be implemented in any new agent are detailed below
- ref: example_green_agent
team: GREEN
type: ExampleAgent
type: example-agent
action_space:
action_map:
0:
action: do-nothing
options: {}
reward_function:
reward_components:
- type: dummy
agent_settings:
start_step: 25
frequency: 20

8
docs/source/how_to_guides/extensible_nodes.rst
View File

@@ -26,9 +26,9 @@ class Router(NetworkNode, identifier="router"):
""" Represents a network router within the simulation, managing routing and forwarding of IP packets across network interfaces."""
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"UserSessionManager": UserSessionManager,
"UserManager": UserManager,
"Terminal": Terminal,
"user-session-manager": UserSessionManager,
"user-manager": UserManager,
"terminal": Terminal,
}
network_interfaces: Dict[str, RouterInterface] = {}
@@ -52,4 +52,4 @@ class Router(NetworkNode, identifier="router"):
Changes to YAML file.
=====================
While effort has been made to ensure that nodes defined within configuration YAML files for use with PrimAITE 3.X remain compatible with PrimAITE v4+, it is encouraged to review for minor changes needed.
While effort has been made to ensure that nodes defined within configuration YAML files for use with PrimAITE 3.X remain compatible with PrimAITE v4+, it is encouraged to review for minor changes needed.

2
docs/source/request_system.rst
View File

@@ -2,6 +2,8 @@
© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
.. _request_system:
Request System
**************

8
docs/source/simulation_components/network/network.rst
View File

@@ -97,19 +97,19 @@ we'll use the following Network that has a client, server, two switches, and a r
network.connect(endpoint_a=switch_2.network_interface[1], endpoint_b=client_1.network_interface[1])
network.connect(endpoint_a=switch_1.network_interface[1], endpoint_b=server_1.network_interface[1])
8. Add ACL rules on the Router to allow ARP and ICMP traffic.
8. Add an ACL rule on the Router to allow ICMP traffic.
.. code-block:: python
router_1.acl.add_rule(
action=ACLAction.PERMIT,
src_port=Port["ARP"],
dst_port=Port["ARP"],
src_port=PORT_LOOKUP["ARP"],
dst_port=PORT_LOOKUP["ARP"],
position=22
)
router_1.acl.add_rule(
action=ACLAction.PERMIT,
protocol=IPProtocol["ICMP"],
protocol=PROTOCOL_LOOKUP["ICMP"],
position=23
)

4
docs/source/simulation_components/network/nodes/wireless_router.rst
View File

@@ -102,8 +102,8 @@ ICMP traffic, ensuring basic network connectivity and ping functionality.
network.connect(pc_a.network_interface[1], router_1.router_interface)
# Configure Router 1 ACLs
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port["ARP"], dst_port=Port["ARP"], position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol["ICMP"], position=23)
router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=PORT_LOOKUP["ARP"], dst_port=PORT_LOOKUP["ARP"], position=22)
router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=PROTOCOL_LOOKUP["ICMP"], position=23)
# Configure PC B
pc_b = Computer(

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

@@ -183,7 +183,7 @@ Python
# Example command: Installing and configuring Ransomware:
ransomware_installation_command = { "commands": [
["software_manager","application","install","RansomwareScript"],
["software_manager","application","install","ransomware-script"],
],
"username": "admin",
"password": "admin",

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

@@ -77,7 +77,7 @@ Python
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
client_1.software_manager.install(DataManipulationBot)
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
data_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("data-manipulation-bot")
data_manipulation_bot.configure(server_ip_address=IPv4Address("192.168.1.14"), payload="DELETE")
data_manipulation_bot.run()
@@ -98,7 +98,7 @@ If not using the data manipulation bot manually, it needs to be used with a data
type: red-database-corrupting-agent
observation_space:
type: UC2RedObservation
type: uc2-red-observation #TODO what
options:
nodes:
- node_name: client_1

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("DatabaseClient")
database_client: DatabaseClient = client.software_manager.software.get("database-sclient")
# Configure the DatabaseClient
database_client.configure(server_ip_address=IPv4Address("192.168.0.1")) # address of the DatabaseService

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

@@ -62,7 +62,7 @@ Python
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
client_1.software_manager.install(RansomwareScript)
RansomwareScript: RansomwareScript = client_1.software_manager.software.get("RansomwareScript")
RansomwareScript: RansomwareScript = client_1.software_manager.software.get("ransomware-script")
RansomwareScript.configure(server_ip_address=IPv4Address("192.168.1.14"))
RansomwareScript.execute()

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

@@ -61,7 +61,7 @@ The :ref:`DNSClient` must be configured to use the :ref:`DNSServer`. The :ref:`D
# Install WebBrowser on computer
computer.software_manager.install(WebBrowser)
web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
web_browser: WebBrowser = computer.software_manager.software.get("web-browser")
web_browser.run()
# configure the WebBrowser

2
docs/source/simulation_components/system/services/database_service.rst
View File

@@ -66,7 +66,7 @@ Python
# Install DatabaseService on server
server.software_manager.install(DatabaseService)
db_service: DatabaseService = server.software_manager.software.get("DatabaseService")
db_service: DatabaseService = server.software_manager.software.get("database-service")
db_service.start()
# configure DatabaseService

2
docs/source/simulation_components/system/services/dns_client.rst
View File

@@ -56,7 +56,7 @@ Python
# Install DNSClient on server
server.software_manager.install(DNSClient)
dns_client: DNSClient = server.software_manager.software.get("DNSClient")
dns_client: DNSClient = server.software_manager.software.get("dns-client")
dns_client.start()
# configure DatabaseService

2
docs/source/simulation_components/system/services/dns_server.rst
View File

@@ -53,7 +53,7 @@ Python
# Install DNSServer on server
server.software_manager.install(DNSServer)
dns_server: DNSServer = server.software_manager.software.get("DNSServer")
dns_server: DNSServer = server.software_manager.software.get("dns-server")
dns_server.start()
# configure DatabaseService

2
docs/source/simulation_components/system/services/ftp_client.rst
View File

@@ -60,7 +60,7 @@ Python
# Install FTPClient on server
server.software_manager.install(FTPClient)
ftp_client: FTPClient = server.software_manager.software.get("FTPClient")
ftp_client: FTPClient = server.software_manager.software.get("ftp-client")
ftp_client.start()

2
docs/source/simulation_components/system/services/ftp_server.rst
View File

@@ -55,7 +55,7 @@ Python
# Install FTPServer on server
server.software_manager.install(FTPServer)
ftp_server: FTPServer = server.software_manager.software.get("FTPServer")
ftp_server: FTPServer = server.software_manager.software.get("ftp-server")
ftp_server.start()
ftp_server.server_password = "test"

2
docs/source/simulation_components/system/services/ntp_client.rst
View File

@@ -53,7 +53,7 @@ Python
# Install NTPClient on server
server.software_manager.install(NTPClient)
ntp_client: NTPClient = server.software_manager.software.get("NTPClient")
ntp_client: NTPClient = server.software_manager.software.get("ntp-client")
ntp_client.start()
ntp_client.configure(ntp_server_ip_address=IPv4Address("192.168.0.10"))

2
docs/source/simulation_components/system/services/ntp_server.rst
View File

@@ -55,7 +55,7 @@ Python
# Install NTPServer on server
server.software_manager.install(NTPServer)
ntp_server: NTPServer = server.software_manager.software.get("NTPServer")
ntp_server: NTPServer = server.software_manager.software.get("ntp-server")
ntp_server.start()

131
docs/source/simulation_components/system/services/terminal.rst
View File

@@ -23,6 +23,14 @@ Key capabilities
- Simulates common Terminal processes/commands.
- Leverages the Service base class for install/uninstall, status tracking etc.
Usage
"""""
- Pre-Installs on any `Node` component (with the exception of `Switches`).
- Terminal Clients connect, execute commands and disconnect from remote nodes.
- Ensures that users are logged in to the component before executing any commands.
- Service runs on SSH port 22 by default.
- Enables Agents to send commands both remotely and locally.
Implementation
""""""""""""""
@@ -30,19 +38,112 @@ Implementation
- Manages remote connections in a dictionary by session ID.
- Processes commands, forwarding to the ``RequestManager`` or ``SessionManager`` where appropriate.
- Extends Service class.
- A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
A detailed guide on the implementation and functionality of the Terminal class can be found in the "Terminal-Processing" jupyter notebook.
Command Format
^^^^^^^^^^^^^^
Terminals implement their commands through leveraging the pre-existing :ref:`request_system`.
Due to this Terminals will only accept commands passed within the ``RequestFormat``.
:py:class:`primaite.game.interface.RequestFormat`
For example, ``terminal`` command actions when used in ``yaml`` format are formatted as follows:
.. code-block:: yaml
command:
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False
This is then loaded from yaml into a dictionary containing the terminal command:
.. code-block:: python
{"command":["file_system", "create", "file", "downloads", "cat.png", "False"]}
Which is then passed to the ``Terminals`` Request Manager to be executed.
Game Layer Usage (Agents)
========================
The below code examples demonstrate how to use terminal related actions in yaml files.
yaml
""""
``node-send-local-command``
"""""""""""""""""""""""""""
Agents can execute local commands without needing to perform a separate remote login action (``node-session-remote-login``).
.. code-block:: yaml
...
...
action: node-send-local-command
options:
node_id: 0
username: admin
password: admin
command: # Example command - Creates a file called 'cat.png' in the downloads folder.
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False"
Usage
"""""
``node-session-remote-login``
"""""""""""""""""
- Pre-Installs on all ``Nodes`` (with the exception of ``Switches``).
- Terminal Clients connect, execute commands and disconnect from remote nodes.
- Ensures that users are logged in to the component before executing any commands.
- Service runs on SSH port 22 by default.
Agents are able to use the terminal to login into remote nodes via ``SSH`` which allows for agents to execute commands on remote hosts.
.. code-block:: yaml
...
...
action: node-session-remote-login
options:
node_id: 0
username: admin
password: admin
remote_ip: 192.168.0.10 # Example Ip Address. (The remote host's IP that will be used by ssh)
``node-send-remote-command``
""""""""""""""""""""""""""""
After remotely logging into another host, an agent can use the ``node-send-remote-command`` to execute commands across the network remotely.
.. code-block:: yaml
...
...
action: node-send-remote-command
options:
node_id: 0
remote_ip: 192.168.0.10
command:
- "file_system"
- "create"
- "file"
- "downloads"
- "cat.png"
- "False"
Simulation Layer Usage
======================
Usage
=====
The below code examples demonstrate how to create a terminal, a remote terminal, and how to send a basic application install command to a remote node.
@@ -65,7 +166,7 @@ Python
operating_state=NodeOperatingState.ON,
)
terminal: Terminal = client.software_manager.software.get("Terminal")
terminal: Terminal = client.software_manager.software.get("terminal")
Creating Remote Terminal Connection
"""""""""""""""""""""""""""""""""""
@@ -86,7 +187,7 @@ Creating Remote Terminal Connection
node_b.power_on()
network.connect(node_a.network_interface[1], node_b.network_interface[1])
terminal_a: Terminal = node_a.software_manager.software.get("Terminal")
terminal_a: Terminal = node_a.software_manager.software.get("terminal")
term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username="admin", password="Admin123!", ip_address="192.168.0.11")
@@ -112,12 +213,12 @@ Executing a basic application install command
node_b.power_on()
network.connect(node_a.network_interface[1], node_b.network_interface[1])
terminal_a: Terminal = node_a.software_manager.software.get("Terminal")
terminal_a: Terminal = node_a.software_manager.software.get("terminal")
term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username="admin", password="Admin123!", ip_address="192.168.0.11")
term_a_term_b_remote_connection.execute(["software_manager", "application", "install", "RansomwareScript"])
term_a_term_b_remote_connection.execute(["software_manager", "application", "install", "ransomware-script"])
@@ -140,7 +241,7 @@ Creating a folder on a remote node
node_b.power_on()
network.connect(node_a.network_interface[1], node_b.network_interface[1])
terminal_a: Terminal = node_a.software_manager.software.get("Terminal")
terminal_a: Terminal = node_a.software_manager.software.get("terminal")
term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username="admin", password="Admin123!", ip_address="192.168.0.11")
@@ -167,7 +268,7 @@ Disconnect from Remote Node
node_b.power_on()
network.connect(node_a.network_interface[1], node_b.network_interface[1])
terminal_a: Terminal = node_a.software_manager.software.get("Terminal")
terminal_a: Terminal = node_a.software_manager.software.get("terminal")
term_a_term_b_remote_connection: RemoteTerminalConnection = terminal_a.login(username="admin", password="Admin123!", ip_address="192.168.0.11")

2
docs/source/simulation_components/system/services/web_server.rst
View File

@@ -56,7 +56,7 @@ Python
# Install WebServer on server
server.software_manager.install(WebServer)
web_server: WebServer = server.software_manager.software.get("WebServer")
web_server: WebServer = server.software_manager.software.get("web-server")
web_server.start()
Via Configuration

2
docs/source/simulation_components/system/software.rst
View File

@@ -30,7 +30,7 @@ See :ref:`Node Start up and Shut down`
node.software_manager.install(WebServer)
web_server: WebServer = node.software_manager.software.get("WebServer")
web_server: WebServer = node.software_manager.software.get("web-server")
assert web_server.operating_state is ServiceOperatingState.RUNNING # service is immediately ran after install
node.power_off()

3253
src/primaite/config/_package_data/uc7_config.yaml Normal file
View File

File diff suppressed because it is too large Load Diff

3303
src/primaite/config/_package_data/uc7_config_tap003.yaml Normal file
View File

File diff suppressed because it is too large Load Diff

43
src/primaite/config/_package_data/uc7_multiple_attack_variants/TAP001_PC1.yaml Normal file
View File

@@ -0,0 +1,43 @@
red: &red
- ref: attacker
team: RED
type: tap-001
agent_settings:
start_step: 1
frequency: 5
variance: 0
repeat_kill_chain: false
repeat_kill_chain_stages: true
default_target_ip: 192.168.220.3
default_starting_node: "ST-PROJ-C-PRV-PC-1"
# starting_nodes: ["ST-PROJ-A-PRV-PC-1", "ST-PROJ-B-PRV-PC-2", "ST-PROJ-C-PRV-PC-3"]
starting_nodes:
kill_chain:
ACTIVATE:
probability: 1
PROPAGATE:
probability: 1
scan_attempts: 20
repeat_scan: false
network_addresses:
- 192.168.230.0/29 # ST Project A
- 192.168.10.0/26 # Remote Site
- 192.168.20.0/30 # Remote DMZ
# - 192.168.240.0/29 # ST Project B
# - 192.168.250.0/29 # ST Project C
- 192.168.220.0/29 # ST Data (Contains Target)
COMMAND_AND_CONTROL:
probability: 1
keep_alive_frequency: 5
masquerade_port: HTTP
masquerade_protocol: TCP
c2_server_name: ISP-PUB-SRV-DNS
c2_server_ip: 8.8.8.8
PAYLOAD:
probability: 1
exfiltrate: true
corrupt: true
exfiltration_folder_name:
target_username: admin
target_password: admin
continue_on_failed_exfil: True

43
src/primaite/config/_package_data/uc7_multiple_attack_variants/TAP001_PC2.yaml Normal file
View File

@@ -0,0 +1,43 @@
red: &red
- ref: attacker
team: RED
type: tap-001
agent_settings:
start_step: 1
frequency: 5
variance: 0
repeat_kill_chain: false
repeat_kill_chain_stages: true
default_target_ip: 192.168.220.3
default_starting_node: "ST-PROJ-B-PRV-PC-2"
# starting_nodes: ["ST-PROJ-A-PRV-PC-1", "ST-PROJ-B-PRV-PC-2", "ST-PROJ-C-PRV-PC-3"]
starting_nodes:
kill_chain:
ACTIVATE:
probability: 1
PROPAGATE:
probability: 1
scan_attempts: 20
repeat_scan: false
network_addresses:
- 192.168.230.0/29 # ST Project A
- 192.168.10.0/26 # Remote Site
- 192.168.20.0/30 # Remote DMZ
# - 192.168.240.0/29 # ST Project B
# - 192.168.250.0/29 # ST Project C
- 192.168.220.0/29 # ST Data (Contains Target)
COMMAND_AND_CONTROL:
probability: 1
keep_alive_frequency: 5
masquerade_port: HTTP
masquerade_protocol: TCP
c2_server_name: ISP-PUB-SRV-DNS
c2_server_ip: 8.8.8.8
PAYLOAD:
probability: 1
exfiltrate: true
corrupt: true
exfiltration_folder_name:
target_username: admin
target_password: admin
continue_on_failed_exfil: True

43
src/primaite/config/_package_data/uc7_multiple_attack_variants/TAP001_PC3.yaml Normal file
View File

@@ -0,0 +1,43 @@
red: &red
- ref: attacker
team: RED
type: tap-001
agent_settings:
start_step: 1
frequency: 5
variance: 0
repeat_kill_chain: false
repeat_kill_chain_stages: true
default_target_ip: 192.168.220.3
default_starting_node: "ST-PROJ-C-PRV-PC-3"
# starting_nodes: ["ST-PROJ-A-PRV-PC-1", "ST-PROJ-B-PRV-PC-2", "ST-PROJ-C-PRV-PC-3"]
starting_nodes:
kill_chain:
ACTIVATE:
probability: 1
PROPAGATE:
probability: 1
scan_attempts: 20
repeat_scan: false
network_addresses:
- 192.168.230.0/29 # ST Project A
- 192.168.10.0/26 # Remote Site
- 192.168.20.0/30 # Remote DMZ
# - 192.168.240.0/29 # ST Project B
# - 192.168.250.0/29 # ST Project C
- 192.168.220.0/29 # ST Data (Contains Target)
COMMAND_AND_CONTROL:
probability: 1
keep_alive_frequency: 5
masquerade_port: HTTP
masquerade_protocol: TCP
c2_server_name: ISP-PUB-SRV-DNS
c2_server_ip: 8.8.8.8
PAYLOAD:
probability: 1
exfiltrate: true
corrupt: true
exfiltration_folder_name:
target_username: admin
target_password: admin
continue_on_failed_exfil: True

94
src/primaite/config/_package_data/uc7_multiple_attack_variants/TAP003.yaml Normal file
View File

@@ -0,0 +1,94 @@
red: &red
- ref: attacker
team: RED
type: tap-003
observation_space: {}
action_space: {}
agent_settings:
start_step: 1
frequency: 3
variance: 0
repeat_kill_chain: false
repeat_kill_chain_stages: true
default_starting_node: "ST-PROJ-A-PRV-PC-1"
starting_nodes:
# starting_nodes: ["ST-PROJ-A-PRV-PC-1", "ST-PROJ-B-PRV-PC-2", "ST-PROJ-C-PRV-PC-3"]
kill_chain:
PLANNING:
probability: 1
starting_network_knowledge:
credentials:
ST-PROJ-A-PRV-PC-1:
username: admin
password: admin
ST-PROJ-B-PRV-PC-2:
username: admin
password: admin
ST-PROJ-C-PRV-PC-3:
username: admin
password: admin
ST-INTRA-PRV-RT-DR-1:
ip_address: 192.168.230.1
username: admin
password: admin
ST-INTRA-PRV-RT-CR:
ip_address: 192.168.160.1
username: admin
password: admin
REM-PUB-RT-DR:
ip_address: 192.168.10.2
username: admin
password: admin
ACCESS:
probability: 1
MANIPULATION:
probability: 1
account_changes:
- host: ST-INTRA-PRV-RT-DR-1
ip_address: 192.168.230.1 # ST-INTRA-PRV-RT-DR-1
action: change_password
username: admin
new_password: "red_pass"
- host: ST-INTRA-PRV-RT-CR
ip_address: 192.168.160.1 # ST-INTRA-PRV-RT-CR
action: change_password
username: "admin"
new_password: "red_pass"
- host: REM-PUB-RT-DR
ip_address: 192.168.10.2 # REM-PUB-RT-DR
action: change_password
username: "admin"
new_password: "red_pass"
EXPLOIT:
probability: 1
malicious_acls:
- target_router: ST-INTRA-PRV-RT-DR-1
position: 1
permission: DENY
src_ip: ALL
src_wildcard: 0.0.255.255
dst_ip: ALL
dst_wildcard: 0.0.255.255
src_port: POSTGRES_SERVER
dst_port: POSTGRES_SERVER
protocol_name: TCP
- target_router: ST-INTRA-PRV-RT-CR
position: 1
permission: DENY
src_ip: ALL
src_wildcard: 0.0.255.255
dst_ip: ALL
dst_wildcard: 0.0.255.255
src_port: HTTP
dst_port: HTTP
protocol_name: TCP
- target_router: REM-PUB-RT-DR
position: 1
permission: DENY
src_ip: ALL
src_wildcard: 0.0.255.255
dst_ip: ALL
dst_wildcard: 0.0.255.255
src_port: DNS
dst_port: DNS
protocol_name: TCP

42
src/primaite/config/_package_data/uc7_multiple_attack_variants/schedule.yaml Normal file
View File

@@ -0,0 +1,42 @@
base_scenario: uc7_config_no_red.yaml
schedule:
0:
- TAP001_PC1.yaml
1:
- TAP001_PC2.yaml
2:
- TAP001_PC3.yaml
3:
- TAP001_PC1.yaml
4:
- TAP001_PC2.yaml
5:
- TAP003.yaml
6:
- TAP003.yaml
7:
- TAP003.yaml
8:
- TAP003.yaml
9:
- TAP003.yaml
10:
- TAP001_PC1.yaml
11:
- TAP003.yaml
12:
- TAP001_PC1.yaml
13:
- TAP003.yaml
14:
- TAP001_PC2.yaml
15:
- TAP003.yaml
16:
- TAP001_PC3.yaml
17:
- TAP003.yaml
18:
- TAP001_PC1.yaml
19:
- TAP003.yaml

3210
src/primaite/config/_package_data/uc7_multiple_attack_variants/uc7_config_no_red.yaml Normal file
View File

File diff suppressed because it is too large Load Diff

71
src/primaite/game/agent/actions/node.py
View File

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC, abstractmethod
from typing import ClassVar, List, Optional, Union
from typing import ClassVar, List, Literal, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
@@ -153,8 +153,6 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, discriminator="node-nmap-po
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, discriminator="node-network-service-recon"):
"""Action which performs an nmap network service recon (ping scan followed by port scan)."""
config: "NodeNetworkServiceReconAction.ConfigSchema"
class ConfigSchema(NodeNMAPAbstractAction.ConfigSchema):
"""Configuration schema for NodeNetworkServiceReconAction."""
@@ -179,3 +177,70 @@ class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, discriminator="node-
"show": config.show,
},
]
class NodeAccountsAddUserAction(AbstractAction, discriminator="node-account-add-user"):
class ConfigSchema(AbstractAction.ConfigSchema):
type: Literal["node-account-add-user"] = "node-account-add-user"
node_name: str
username: str
password: str
is_admin: bool
@classmethod
@staticmethod
def form_request(config: ConfigSchema) -> RequestFormat:
return [
"network",
"node",
config.node_name,
"service",
"user-manager",
"add_user",
config.username,
config.password,
config.is_admin,
]
class NodeAccountsDisableUserAction(AbstractAction, discriminator="node-account-disable-user"):
class ConfigSchema(AbstractAction.ConfigSchema):
type: Literal["node-account-disable-user"] = "node-account-disable-user"
node_name: str
username: str
@classmethod
@staticmethod
def form_request(config: ConfigSchema) -> RequestFormat:
return [
"network",
"node",
config.node_name,
"service",
"user-manager",
"disable_user",
config.username,
]
class NodeSendLocalCommandAction(AbstractAction, discriminator="node-send-local-command"):
class ConfigSchema(AbstractAction.ConfigSchema):
type: Literal["node-send-local-command"] = "node-send-local-command"
node_name: str
username: str
password: str
command: RequestFormat
@staticmethod
def form_request(config: ConfigSchema) -> RequestFormat:
return [
"network",
"node",
config.node_name,
"service",
"terminal",
"send_local_command",
config.username,
config.password,
{"command": config.command},
]

13
src/primaite/game/agent/actions/session.py
View File

@@ -34,8 +34,6 @@ class NodeSessionAbstractAction(AbstractAction, ABC):
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="node-session-remote-login"):
"""Action which performs a remote session login."""
config: "NodeSessionsRemoteLoginAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
"""Configuration schema for NodeSessionsRemoteLoginAction."""
@@ -53,7 +51,7 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="no
config.node_name,
"service",
"terminal",
"node-session-remote-login",
"node_session_remote_login",
config.username,
config.password,
config.remote_ip,
@@ -63,8 +61,6 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="no
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, discriminator="node-session-remote-logoff"):
"""Action which performs a remote session logout."""
config: "NodeSessionsRemoteLogoutAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
"""Configuration schema for NodeSessionsRemoteLogoutAction."""
@@ -78,14 +74,13 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, discriminator="n
return ["network", "node", config.node_name, "service", "terminal", config.verb, config.remote_ip]
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, discriminator="node-account-change-password"):
class NodeAccountChangePasswordAction(AbstractAction, discriminator="node-account-change-password"):
"""Action which changes the password for a user."""
config: "NodeAccountChangePasswordAction.ConfigSchema"
class ConfigSchema(NodeSessionAbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for NodeAccountsChangePasswordAction."""
node_name: str
username: str
current_password: str
new_password: str

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

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, TYPE_CHECKING
from gymnasium.core import ActType, ObsType
from prettytable import PrettyTable
from pydantic import BaseModel, ConfigDict, Field
from primaite.game.agent.actions import ActionManager
@@ -42,6 +43,9 @@ class AgentHistoryItem(BaseModel):
reward_info: Dict[str, Any] = {}
observation: Optional[ObsType] = None
"""The observation space data for this step."""
class AbstractAgent(BaseModel, ABC):
"""Base class for scripted and RL agents."""
@@ -67,6 +71,9 @@ class AbstractAgent(BaseModel, ABC):
default_factory=lambda: ObservationManager.ConfigSchema()
)
reward_function: RewardFunction.ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema())
thresholds: Optional[Dict] = {}
# TODO: this is only relevant to some observations, need to refactor the way thresholds are dealt with (#3085)
"""A dict containing the observation thresholds."""
config: ConfigSchema = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
@@ -90,10 +97,42 @@ class AbstractAgent(BaseModel, ABC):
def model_post_init(self, __context: Any) -> None:
"""Overwrite the default empty action, observation, and rewards with ones defined through the config."""
self.action_manager = ActionManager(config=self.config.action_space)
self.config.observation_space.options.thresholds = self.config.thresholds
self.observation_manager = ObservationManager(config=self.config.observation_space)
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.
:param ignored_actions: OPTIONAL: List of actions to be ignored when displaying the history.
If not provided, defaults to ignore DONOTHING actions.
"""
if not ignored_actions:
ignored_actions = ["DONOTHING"]
table = PrettyTable()
table.field_names = ["Step", "Action", "Node", "Application", "Response"]
print(f"Actions for '{self.agent_name}':")
for item in self.history:
if item.action in ignored_actions:
pass
else:
table = self.add_agent_action(item=item, table=table)
print(table)
def update_observation(self, state: Dict) -> ObsType:
"""
Convert a state from the simulator into an observation for the agent using the observation space.
@@ -140,12 +179,23 @@ class AbstractAgent(BaseModel, ABC):
return request
def process_action_response(
self, timestep: int, action: str, parameters: Dict[str, Any], request: RequestFormat, response: RequestResponse
self,
timestep: int,
action: str,
parameters: Dict[str, Any],
request: RequestFormat,
response: RequestResponse,
observation: ObsType,
) -> None:
"""Process the response from the most recent action."""
self.history.append(
AgentHistoryItem(
timestep=timestep, action=action, parameters=parameters, request=request, response=response
timestep=timestep,
action=action,
parameters=parameters,
request=request,
response=response,
observation=observation,
)
)

53
src/primaite/game/agent/observations/file_system_observations.py
View File

@@ -26,7 +26,13 @@ class FileObservation(AbstractObservation, discriminator="file"):
file_system_requires_scan: Optional[bool] = None
"""If True, the file must be scanned to update the health state. Tf False, the true state is always shown."""
def __init__(self, where: WhereType, include_num_access: bool, file_system_requires_scan: bool) -> None:
def __init__(
self,
where: WhereType,
include_num_access: bool,
file_system_requires_scan: bool,
thresholds: Optional[Dict] = {},
) -> None:
"""
Initialise a file observation instance.
@@ -48,10 +54,36 @@ class FileObservation(AbstractObservation, discriminator="file"):
if self.include_num_access:
self.default_observation["num_access"] = 0
# TODO: allow these to be configured in yaml
self.high_threshold = 10
self.med_threshold = 5
self.low_threshold = 0
if thresholds.get("file_access") is None:
self.low_file_access_threshold = 0
self.med_file_access_threshold = 5
self.high_file_access_threshold = 10
else:
self._set_file_access_threshold(
thresholds=[
thresholds.get("file_access")["low"],
thresholds.get("file_access")["medium"],
thresholds.get("file_access")["high"],
]
)
def _set_file_access_threshold(self, thresholds: List[int]):
"""
Method that validates and then sets the file access threshold.
:param: thresholds: The file access threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=[
thresholds[0],
thresholds[1],
thresholds[2],
],
threshold_identifier="file_access",
):
self.low_file_access_threshold = thresholds[0]
self.med_file_access_threshold = thresholds[1]
self.high_file_access_threshold = thresholds[2]
def _categorise_num_access(self, num_access: int) -> int:
"""
@@ -60,11 +92,11 @@ class FileObservation(AbstractObservation, discriminator="file"):
:param num_access: Number of file accesses.
:return: Bin number corresponding to the number of accesses.
"""
if num_access > self.high_threshold:
if num_access > self.high_file_access_threshold:
return 3
elif num_access > self.med_threshold:
elif num_access > self.med_file_access_threshold:
return 2
elif num_access > self.low_threshold:
elif num_access > self.low_file_access_threshold:
return 1
return 0
@@ -122,6 +154,7 @@ class FileObservation(AbstractObservation, discriminator="file"):
where=parent_where + ["files", config.file_name],
include_num_access=config.include_num_access,
file_system_requires_scan=config.file_system_requires_scan,
thresholds=config.thresholds,
)
@@ -149,6 +182,7 @@ class FolderObservation(AbstractObservation, discriminator="folder"):
num_files: int,
include_num_access: bool,
file_system_requires_scan: bool,
thresholds: Optional[Dict] = {},
) -> None:
"""
Initialise a folder observation instance.
@@ -177,6 +211,7 @@ class FolderObservation(AbstractObservation, discriminator="folder"):
where=None,
include_num_access=include_num_access,
file_system_requires_scan=self.file_system_requires_scan,
thresholds=thresholds,
)
)
while len(self.files) > num_files:
@@ -253,6 +288,7 @@ class FolderObservation(AbstractObservation, discriminator="folder"):
for file_config in config.files:
file_config.include_num_access = config.include_num_access
file_config.file_system_requires_scan = config.file_system_requires_scan
file_config.thresholds = config.thresholds
files = [FileObservation.from_config(config=f, parent_where=where) for f in config.files]
return cls(
@@ -261,4 +297,5 @@ class FolderObservation(AbstractObservation, discriminator="folder"):
num_files=config.num_files,
include_num_access=config.include_num_access,
file_system_requires_scan=config.file_system_requires_scan,
thresholds=config.thresholds,
)

46
src/primaite/game/agent/observations/host_observations.py
View File

@@ -54,7 +54,15 @@ class HostObservation(AbstractObservation, discriminator="host"):
"""
If True, files and folders must be scanned to update the health state. If False, true state is always shown.
"""
include_users: Optional[bool] = None
services_requires_scan: Optional[bool] = None
"""
If True, services must be scanned to update the health state. If False, true state is always shown.
"""
applications_requires_scan: Optional[bool] = None
"""
If True, applications must be scanned to update the health state. If False, true state is always shown.
"""
include_users: Optional[bool] = True
"""If True, report user session information."""
def __init__(
@@ -73,6 +81,8 @@ class HostObservation(AbstractObservation, discriminator="host"):
monitored_traffic: Optional[Dict],
include_num_access: bool,
file_system_requires_scan: bool,
services_requires_scan: bool,
applications_requires_scan: bool,
include_users: bool,
) -> None:
"""
@@ -108,6 +118,12 @@ class HostObservation(AbstractObservation, discriminator="host"):
:param file_system_requires_scan: If True, the files and folders must be scanned to update the health state.
If False, the true state is always shown.
:type file_system_requires_scan: bool
:param services_requires_scan: If True, services must be scanned to update the health state.
If False, the true state is always shown.
:type services_requires_scan: bool
:param applications_requires_scan: If True, applications must be scanned to update the health state.
If False, the true state is always shown.
:type applications_requires_scan: bool
:param include_users: If True, report user session information.
:type include_users: bool
"""
@@ -121,7 +137,7 @@ class HostObservation(AbstractObservation, discriminator="host"):
# Ensure lists have lengths equal to specified counts by truncating or padding
self.services: List[ServiceObservation] = services
while len(self.services) < num_services:
self.services.append(ServiceObservation(where=None))
self.services.append(ServiceObservation(where=None, services_requires_scan=services_requires_scan))
while len(self.services) > num_services:
truncated_service = self.services.pop()
msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
@@ -129,7 +145,9 @@ class HostObservation(AbstractObservation, discriminator="host"):
self.applications: List[ApplicationObservation] = applications
while len(self.applications) < num_applications:
self.applications.append(ApplicationObservation(where=None))
self.applications.append(
ApplicationObservation(where=None, applications_requires_scan=applications_requires_scan)
)
while len(self.applications) > num_applications:
truncated_application = self.applications.pop()
msg = f"Too many applications in Node observation space for node. Truncating {truncated_application.where}"
@@ -153,7 +171,13 @@ class HostObservation(AbstractObservation, discriminator="host"):
self.nics: List[NICObservation] = network_interfaces
while len(self.nics) < num_nics:
self.nics.append(NICObservation(where=None, include_nmne=include_nmne, monitored_traffic=monitored_traffic))
self.nics.append(
NICObservation(
where=None,
include_nmne=include_nmne,
monitored_traffic=monitored_traffic,
)
)
while len(self.nics) > num_nics:
truncated_nic = self.nics.pop()
msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_nic.where}"
@@ -269,8 +293,15 @@ class HostObservation(AbstractObservation, discriminator="host"):
folder_config.include_num_access = config.include_num_access
folder_config.num_files = config.num_files
folder_config.file_system_requires_scan = config.file_system_requires_scan
folder_config.thresholds = config.thresholds
for nic_config in config.network_interfaces:
nic_config.include_nmne = config.include_nmne
nic_config.thresholds = config.thresholds
for service_config in config.services:
service_config.services_requires_scan = config.services_requires_scan
for application_config in config.applications:
application_config.applications_requires_scan = config.applications_requires_scan
application_config.thresholds = config.thresholds
services = [ServiceObservation.from_config(config=c, parent_where=where) for c in config.services]
applications = [ApplicationObservation.from_config(config=c, parent_where=where) for c in config.applications]
@@ -281,7 +312,10 @@ class HostObservation(AbstractObservation, discriminator="host"):
count = 1
while len(nics) < config.num_nics:
nic_config = NICObservation.ConfigSchema(
nic_num=count, include_nmne=config.include_nmne, monitored_traffic=config.monitored_traffic
nic_num=count,
include_nmne=config.include_nmne,
monitored_traffic=config.monitored_traffic,
thresholds=config.thresholds,
)
nics.append(NICObservation.from_config(config=nic_config, parent_where=where))
count += 1
@@ -301,5 +335,7 @@ class HostObservation(AbstractObservation, discriminator="host"):
monitored_traffic=config.monitored_traffic,
include_num_access=config.include_num_access,
file_system_requires_scan=config.file_system_requires_scan,
services_requires_scan=config.services_requires_scan,
applications_requires_scan=config.applications_requires_scan,
include_users=config.include_users,
)

49
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -1,13 +1,14 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, List, Optional
from typing import ClassVar, Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.nmne import NMNEConfig
from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.port import Port
@@ -15,6 +16,9 @@ from primaite.utils.validation.port import Port
class NICObservation(AbstractObservation, discriminator="network-interface"):
"""Status information about a network interface within the simulation environment."""
capture_nmne: ClassVar[bool] = NMNEConfig().capture_nmne
"A Boolean specifying whether malicious network events should be captured."
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for NICObservation."""
@@ -25,7 +29,13 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
monitored_traffic: Optional[Dict[IPProtocol, List[Port]]] = None
"""A dict containing which traffic types are to be included in the observation."""
def __init__(self, where: WhereType, include_nmne: bool, monitored_traffic: Optional[Dict] = None) -> None:
def __init__(
self,
where: WhereType,
include_nmne: bool,
monitored_traffic: Optional[Dict] = None,
thresholds: Dict = {},
) -> None:
"""
Initialise a network interface observation instance.
@@ -45,10 +55,18 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
self.nmne_inbound_last_step: int = 0
self.nmne_outbound_last_step: int = 0
# TODO: allow these to be configured in yaml
self.high_nmne_threshold = 10
self.med_nmne_threshold = 5
self.low_nmne_threshold = 0
if thresholds.get("nmne") is None:
self.low_nmne_threshold = 0
self.med_nmne_threshold = 5
self.high_nmne_threshold = 10
else:
self._set_nmne_threshold(
thresholds=[
thresholds.get("nmne")["low"],
thresholds.get("nmne")["medium"],
thresholds.get("nmne")["high"],
]
)
self.monitored_traffic = monitored_traffic
if self.monitored_traffic:
@@ -105,6 +123,20 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
bandwidth_utilisation = traffic_value / nic_max_bandwidth
return int(bandwidth_utilisation * 9) + 1
def _set_nmne_threshold(self, thresholds: List[int]):
"""
Method that validates and then sets the NMNE threshold.
:param: thresholds: The NMNE threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=thresholds,
threshold_identifier="nmne",
):
self.low_nmne_threshold = thresholds[0]
self.med_nmne_threshold = thresholds[1]
self.high_nmne_threshold = thresholds[2]
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
@@ -116,7 +148,7 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
"""
nic_state = access_from_nested_dict(state, self.where)
if nic_state is NOT_PRESENT_IN_STATE:
if nic_state is NOT_PRESENT_IN_STATE or self.where is None:
return self.default_observation
obs = {"nic_status": 1 if nic_state["enabled"] else 2}
@@ -164,7 +196,7 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
for port in self.monitored_traffic[protocol]:
obs["TRAFFIC"][protocol][port] = {"inbound": 0, "outbound": 0}
if self.include_nmne:
if self.capture_nmne and self.include_nmne:
obs.update({"NMNE": {}})
direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
@@ -224,6 +256,7 @@ class NICObservation(AbstractObservation, discriminator="network-interface"):
where=parent_where + ["NICs", config.nic_num],
include_nmne=config.include_nmne,
monitored_traffic=config.monitored_traffic,
thresholds=config.thresholds,
)

18
src/primaite/game/agent/observations/node_observations.py
View File

@@ -48,7 +48,13 @@ class NodesObservation(AbstractObservation, discriminator="nodes"):
include_num_access: Optional[bool] = None
"""Flag to include the number of accesses."""
file_system_requires_scan: bool = True
"""If True, the folder must be scanned to update the health state. Tf False, the true state is always shown."""
"""If True, the folder must be scanned to update the health state. If False, the true state is always shown."""
services_requires_scan: bool = True
"""If True, the services must be scanned to update the health state.
If False, the true state is always shown."""
applications_requires_scan: bool = True
"""If True, the applications must be scanned to update the health state.
If False, the true state is always shown."""
include_users: Optional[bool] = True
"""If True, report user session information."""
num_ports: Optional[int] = None
@@ -196,8 +202,14 @@ class NodesObservation(AbstractObservation, discriminator="nodes"):
host_config.include_num_access = config.include_num_access
if host_config.file_system_requires_scan is None:
host_config.file_system_requires_scan = config.file_system_requires_scan
if host_config.services_requires_scan is None:
host_config.services_requires_scan = config.services_requires_scan
if host_config.applications_requires_scan is None:
host_config.applications_requires_scan = config.applications_requires_scan
if host_config.include_users is None:
host_config.include_users = config.include_users
if not host_config.thresholds:
host_config.thresholds = config.thresholds
for router_config in config.routers:
if router_config.num_ports is None:
@@ -214,6 +226,8 @@ class NodesObservation(AbstractObservation, discriminator="nodes"):
router_config.num_rules = config.num_rules
if router_config.include_users is None:
router_config.include_users = config.include_users
if not router_config.thresholds:
router_config.thresholds = config.thresholds
for firewall_config in config.firewalls:
if firewall_config.ip_list is None:
@@ -228,6 +242,8 @@ class NodesObservation(AbstractObservation, discriminator="nodes"):
firewall_config.num_rules = config.num_rules
if firewall_config.include_users is None:
firewall_config.include_users = config.include_users
if not firewall_config.thresholds:
firewall_config.thresholds = config.thresholds
hosts = [HostObservation.from_config(config=c, parent_where=where) for c in config.hosts]
routers = [RouterObservation.from_config(config=c, parent_where=where) for c in config.routers]

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

@@ -114,7 +114,9 @@ class NestedObservation(AbstractObservation, discriminator="custom"):
instances = dict()
for component in config.components:
obs_class = AbstractObservation._registry[component.type]
obs_instance = obs_class.from_config(config=obs_class.ConfigSchema(**component.options))
obs_instance = obs_class.from_config(
config=obs_class.ConfigSchema(**component.options, thresholds=config.thresholds)
)
instances[component.label] = obs_instance
return cls(components=instances)
@@ -242,8 +244,5 @@ class ObservationManager(BaseModel):
"""
if config is None:
return cls(NullObservation())
obs_type = config["type"]
obs_class = AbstractObservation._registry[obs_type]
observation = obs_class.from_config(config=obs_class.ConfigSchema(**config["options"]))
obs_manager = cls(observation)
obs_manager = cls(config=config)
return obs_manager

36
src/primaite/game/agent/observations/observations.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
"""Manages the observation space for the agent."""
from abc import ABC, abstractmethod
from typing import Any, Dict, Iterable, Optional, Type, Union
from typing import Any, Dict, Iterable, List, Optional, Type, Union
from gymnasium import spaces
from gymnasium.core import ObsType
@@ -19,6 +19,9 @@ class AbstractObservation(ABC):
class ConfigSchema(ABC, BaseModel):
"""Config schema for observations."""
thresholds: Optional[Dict] = {}
"""A dict containing the observation thresholds."""
model_config = ConfigDict(extra="forbid")
_registry: Dict[str, Type["AbstractObservation"]] = {}
@@ -69,3 +72,34 @@ class AbstractObservation(ABC):
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> "AbstractObservation":
"""Create this observation space component form a serialised format."""
return cls()
def _validate_thresholds(self, thresholds: List[int] = None, threshold_identifier: Optional[str] = "") -> bool:
"""
Method that checks if the thresholds are non overlapping and in the correct (ascending) order.
Pass in the thresholds from low to high e.g.
thresholds=[low_threshold, med_threshold, ..._threshold, high_threshold]
Throws an error if the threshold is not valid
:param: thresholds: List of thresholds in ascending order.
:type: List[int]
:param: threshold_identifier: The name of the threshold option.
:type: Optional[str]
:returns: bool
"""
if thresholds is None or len(thresholds) < 2:
raise Exception(f"{threshold_identifier} thresholds are invalid {thresholds}")
for idx in range(1, len(thresholds)):
if not isinstance(thresholds[idx], int):
raise Exception(f"{threshold_identifier} threshold ({thresholds[idx]}) is not a valid int.")
if not isinstance(thresholds[idx - 1], int):
raise Exception(f"{threshold_identifier} threshold ({thresholds[idx]}) is not a valid int.")
if thresholds[idx] <= thresholds[idx - 1]:
raise Exception(
f"{threshold_identifier} threshold ({thresholds[idx - 1]}) "
f"is greater than or equal to ({thresholds[idx]}.)"
)
return True

78
src/primaite/game/agent/observations/software_observation.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
@@ -19,7 +19,10 @@ class ServiceObservation(AbstractObservation, discriminator="service"):
service_name: str
"""Name of the service, used for querying simulation state dictionary"""
def __init__(self, where: WhereType) -> None:
services_requires_scan: Optional[bool] = None
"""If True, services must be scanned to update the health state. If False, true state is always shown."""
def __init__(self, where: WhereType, services_requires_scan: bool) -> None:
"""
Initialise a service observation instance.
@@ -28,6 +31,7 @@ class ServiceObservation(AbstractObservation, discriminator="service"):
:type where: WhereType
"""
self.where = where
self.services_requires_scan = services_requires_scan
self.default_observation = {"operating_status": 0, "health_status": 0}
def observe(self, state: Dict) -> ObsType:
@@ -44,7 +48,9 @@ class ServiceObservation(AbstractObservation, discriminator="service"):
return self.default_observation
return {
"operating_status": service_state["operating_state"],
"health_status": service_state["health_state_visible"],
"health_status": service_state["health_state_visible"]
if self.services_requires_scan
else service_state["health_state_actual"],
}
@property
@@ -70,7 +76,9 @@ class ServiceObservation(AbstractObservation, discriminator="service"):
:return: Constructed service observation instance.
:rtype: ServiceObservation
"""
return cls(where=parent_where + ["services", config.service_name])
return cls(
where=parent_where + ["services", config.service_name], services_requires_scan=config.services_requires_scan
)
class ApplicationObservation(AbstractObservation, discriminator="application"):
@@ -82,7 +90,12 @@ class ApplicationObservation(AbstractObservation, discriminator="application"):
application_name: str
"""Name of the application, used for querying simulation state dictionary"""
def __init__(self, where: WhereType) -> None:
applications_requires_scan: Optional[bool] = None
"""
If True, applications must be scanned to update the health state. If False, true state is always shown.
"""
def __init__(self, where: WhereType, applications_requires_scan: bool, thresholds: Optional[Dict] = {}) -> None:
"""
Initialise an application observation instance.
@@ -92,25 +105,52 @@ class ApplicationObservation(AbstractObservation, discriminator="application"):
:type where: WhereType
"""
self.where = where
self.applications_requires_scan = applications_requires_scan
self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
# TODO: allow these to be configured in yaml
self.high_threshold = 10
self.med_threshold = 5
self.low_threshold = 0
if thresholds.get("app_executions") is None:
self.low_app_execution_threshold = 0
self.med_app_execution_threshold = 5
self.high_app_execution_threshold = 10
else:
self._set_application_execution_thresholds(
thresholds=[
thresholds.get("app_executions")["low"],
thresholds.get("app_executions")["medium"],
thresholds.get("app_executions")["high"],
]
)
def _set_application_execution_thresholds(self, thresholds: List[int]):
"""
Method that validates and then sets the application execution threshold.
:param: thresholds: The application execution threshold to validate and set.
"""
if self._validate_thresholds(
thresholds=[
thresholds[0],
thresholds[1],
thresholds[2],
],
threshold_identifier="app_executions",
):
self.low_app_execution_threshold = thresholds[0]
self.med_app_execution_threshold = thresholds[1]
self.high_app_execution_threshold = thresholds[2]
def _categorise_num_executions(self, num_executions: int) -> int:
"""
Represent number of file accesses as a categorical variable.
Represent number of application executions as a categorical variable.
:param num_access: Number of file accesses.
:param num_access: Number of application executions.
:return: Bin number corresponding to the number of accesses.
"""
if num_executions > self.high_threshold:
if num_executions > self.high_app_execution_threshold:
return 3
elif num_executions > self.med_threshold:
elif num_executions > self.med_app_execution_threshold:
return 2
elif num_executions > self.low_threshold:
elif num_executions > self.low_app_execution_threshold:
return 1
return 0
@@ -128,7 +168,9 @@ class ApplicationObservation(AbstractObservation, discriminator="application"):
return self.default_observation
return {
"operating_status": application_state["operating_state"],
"health_status": application_state["health_state_visible"],
"health_status": application_state["health_state_visible"]
if self.applications_requires_scan
else application_state["health_state_actual"],
"num_executions": self._categorise_num_executions(application_state["num_executions"]),
}
@@ -161,4 +203,8 @@ class ApplicationObservation(AbstractObservation, discriminator="application"):
:return: Constructed application observation instance.
:rtype: ApplicationObservation
"""
return cls(where=parent_where + ["applications", config.application_name])
return cls(
where=parent_where + ["applications", config.application_name],
applications_requires_scan=config.applications_requires_scan,
thresholds=config.thresholds,
)

26
src/primaite/game/game.py
View File

@@ -7,12 +7,14 @@ from pydantic import BaseModel, ConfigDict
from primaite import DEFAULT_BANDWIDTH, getLogger
from primaite.game.agent.interface import AbstractAgent, ProxyAgent
from primaite.game.agent.observations import NICObservation
from primaite.game.agent.rewards import SharedReward
from primaite.game.science import graph_has_cycle, topological_sort
from primaite.simulator import SIM_OUTPUT
from primaite.simulator.network.creation import NetworkNodeAdder
from primaite.simulator.network.hardware.base import NetworkInterface, Node, NodeOperatingState, UserManager
from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall # noqa: F401
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from primaite.simulator.network.nmne import NMNEConfig
@@ -44,15 +46,15 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
SERVICE_TYPES_MAPPING = {
"DNSClient": DNSClient,
"DNSServer": DNSServer,
"DatabaseService": DatabaseService,
"WebServer": WebServer,
"FTPClient": FTPClient,
"FTPServer": FTPServer,
"NTPClient": NTPClient,
"NTPServer": NTPServer,
"Terminal": Terminal,
"dns-client": DNSClient,
"dns-server": DNSServer,
"database-service": DatabaseService,
"web-server": WebServer,
"ftp-client": FTPClient,
"ftp-server": FTPServer,
"ntp-client": NTPClient,
"ntp-server": NTPServer,
"terminal": Terminal,
}
"""List of available services that can be installed on nodes in the PrimAITE Simulation."""
@@ -68,6 +70,8 @@ class PrimaiteGameOptions(BaseModel):
seed: int = None
"""Random number seed for RNGs."""
generate_seed_value: bool = False
"""Internally generated seed value."""
max_episode_length: int = 256
"""Maximum number of episodes for the PrimAITE game."""
ports: List[Port]
@@ -175,6 +179,7 @@ class PrimaiteGame:
parameters=parameters,
request=request,
response=response,
observation=obs,
)
def pre_timestep(self) -> None:
@@ -263,6 +268,7 @@ class PrimaiteGame:
node_sets_cfg = network_config.get("node_sets", [])
# Set the NMNE capture config
NetworkInterface.nmne_config = NMNEConfig(**network_config.get("nmne_config", {}))
NICObservation.capture_nmne = NMNEConfig(**network_config.get("nmne_config", {})).capture_nmne
for node_cfg in nodes_cfg:
n_type = node_cfg["type"]
@@ -293,6 +299,7 @@ class PrimaiteGame:
if "users" in node_cfg and new_node.software_manager.software.get("user-manager"):
user_manager: UserManager = new_node.software_manager.software["user-manager"] # noqa
for user_cfg in node_cfg["users"]:
user_manager.add_user(**user_cfg, bypass_can_perform_action=True)
@@ -407,6 +414,7 @@ class PrimaiteGame:
agents_cfg = cfg.get("agents", [])
for agent_cfg in agents_cfg:
agent_cfg = {**agent_cfg, "thresholds": game.options.thresholds}
new_agent = AbstractAgent.from_config(agent_cfg)
game.agents[agent_cfg["ref"]] = new_agent
if isinstance(new_agent, ProxyAgent):

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

@@ -50,40 +50,22 @@
"custom_c2_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: ProxyAgent\n",
" type: proxy-a.gent\n",
"\n",
" action_space:\n",
" options:\n",
" nodes:\n",
" - node_name: web_server\n",
" applications:\n",
" - application_name: C2Beacon\n",
" - node_name: client_1\n",
" applications:\n",
" - application_name: C2Server\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 2\n",
" max_nics_per_node: 8\n",
" max_acl_rules: 10\n",
" ip_list:\n",
" - 192.168.1.21\n",
" - 192.168.1.14\n",
" wildcard_list:\n",
" - 0.0.0.1\n",
" action_map:\n",
" 0:\n",
" action: do_nothing\n",
" options: {}\n",
" 1:\n",
" action: node_application_install\n",
" action: node-application-install\n",
" options:\n",
" node_id: 0\n",
" application_name: C2Beacon\n",
" node_name: web_server\n",
" application_name: c2-beacon\n",
" 2:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" node_name: web_server\n",
" config:\n",
" c2_server_ip_address: 192.168.10.21\n",
" keep_alive_frequency:\n",
@@ -92,10 +74,10 @@
" 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",
" 4:\n",
" action: c2_server_terminal_command\n",
" action: c2-server-terminal-command\n",
" options:\n",
" node_id: 1\n",
" ip_address:\n",
@@ -111,14 +93,14 @@
" 5:\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" config:\n",
" server_ip_address: 192.168.1.14\n",
" payload: ENCRYPT\n",
" 6:\n",
" action: c2_server_data_exfiltrate\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" target_file_name: \"database.db\"\n",
" target_folder_name: \"database\"\n",
" exfiltration_folder_name: \"spoils\"\n",
@@ -128,31 +110,27 @@
" password: admin\n",
"\n",
" 7:\n",
" action: c2_server_ransomware_launch\n",
" action: c2-server-ransomware-launch\n",
" options:\n",
" node_id: 1\n",
" node_name: client_1\n",
" 8:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" 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",
" 9:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\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",
"\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
"\"\"\"\n",
"c2_agent_yaml = yaml.safe_load(custom_c2_agent)"
]
@@ -225,7 +203,7 @@
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: C2Beacon\n",
" - application_name: c2-beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
@@ -233,7 +211,7 @@
" action: node_application_install \n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" application_name: C2Beacon\n",
" application_name: c2-beacon\n",
"```"
]
},
@@ -268,7 +246,7 @@
" action_map:\n",
" ...\n",
" 2:\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0 # Node Index\n",
" config: # Further information about these config options can be found at the bottom of this notebook.\n",
@@ -286,7 +264,7 @@
"outputs": [],
"source": [
"env.step(2)\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
"web_server.software_manager.show()\n",
"c2_beacon.show()"
]
@@ -307,13 +285,13 @@
" nodes: # Node List\n",
" - node_name: web_server\n",
" applications: \n",
" - application_name: C2Beacon\n",
" - application_name: c2-beacon\n",
" ...\n",
" ...\n",
" action_map:\n",
" ...\n",
" 3:\n",
" action: node_application_execute\n",
" action: node-application-execute\n",
" options:\n",
" node_id: 0\n",
" application_id: 0\n",
@@ -374,11 +352,11 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 4:\n",
" action: C2_SERVER_TERMINAL_COMMAND\n",
" action: c2-server-terminal-command\n",
" options:\n",
" node_id: 1\n",
" ip_address:\n",
@@ -431,7 +409,7 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 5:\n",
@@ -459,7 +437,7 @@
"metadata": {},
"outputs": [],
"source": [
"ransomware_script: RansomwareScript = web_server.software_manager.software[\"RansomwareScript\"]\n",
"ransomware_script: RansomwareScript = web_server.software_manager.software[\"ransomware-script\"]\n",
"web_server.software_manager.show()\n",
"ransomware_script.show()"
]
@@ -483,11 +461,11 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 6:\n",
" action: c2_server_data_exfiltrate\n",
" action: c2-server-data-exfiltrate\n",
" options:\n",
" node_id: 1\n",
" target_file_name: \"database.db\"\n",
@@ -549,11 +527,11 @@
" ...\n",
" - node_name: client_1\n",
" applications: \n",
" - application_name: C2Server\n",
" - application_name: c2-server\n",
" ...\n",
" action_map:\n",
" 7:\n",
" action: c2_server_ransomware_launch\n",
" action: c2-server-ransomware-launch\n",
" options:\n",
" node_id: 1\n",
"```\n"
@@ -598,20 +576,20 @@
"custom_blue_agent_yaml = \"\"\"\n",
" - ref: defender\n",
" team: BLUE\n",
" type: ProxyAgent\n",
" type: proxy-agent\n",
"\n",
" observation_space:\n",
" type: CUSTOM\n",
" type: custom\n",
" options:\n",
" components:\n",
" - type: NODES\n",
" - type: nodes\n",
" label: NODES\n",
" options:\n",
" hosts:\n",
" - hostname: web_server\n",
" applications:\n",
" - application_name: C2Beacon\n",
" - application_name: RansomwareScript\n",
" - application_name: c2-beacon\n",
" - application_name: ransomware-script\n",
" folders:\n",
" - folder_name: exfiltration_folder\n",
" files:\n",
@@ -661,7 +639,7 @@
" - UDP\n",
" num_rules: 10\n",
"\n",
" - type: LINKS\n",
" - type: links\n",
" label: LINKS\n",
" options:\n",
" link_references:\n",
@@ -675,7 +653,7 @@
" - switch_2:eth-1<->client_1:eth-1\n",
" - switch_2:eth-2<->client_2:eth-1\n",
" - switch_2:eth-7<->security_suite:eth-2\n",
" - type: \"NONE\"\n",
" - type: \"none\"\n",
" label: ICS\n",
" options: {}\n",
"\n",
@@ -685,16 +663,16 @@
" action: do_nothing\n",
" options: {}\n",
" 1:\n",
" action: node_application_remove\n",
" action: node-application-remove\n",
" options:\n",
" node_id: 0\n",
" node_name: web-server\n",
" application_name: C2Beacon\n",
" 2:\n",
" action: node_shutdown\n",
" action: node-shutdown\n",
" options:\n",
" node_id: 0\n",
" node_name: web-server\n",
" 3:\n",
" action: router_acl_add_rule\n",
" action: router-acl-add-rule\n",
" options:\n",
" target_router: router_1\n",
" position: 1\n",
@@ -707,36 +685,6 @@
" 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",
@@ -875,7 +823,7 @@
"outputs": [],
"source": [
"# Installing RansomwareScript via C2 Terminal Commands\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"]],\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"ransomware-script\"]],\n",
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)\n"
@@ -1034,11 +982,11 @@
" web_server: Server = given_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
"\n",
" client_1.software_manager.install(C2Server)\n",
" c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
" c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
" c2_server.run()\n",
"\n",
" web_server.software_manager.install(C2Beacon)\n",
" c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
" c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
" c2_beacon.configure(c2_server_ip_address=\"192.168.10.21\")\n",
" c2_beacon.establish()\n",
"\n",
@@ -1132,11 +1080,11 @@
"outputs": [],
"source": [
"# Attempting to install the C2 RansomwareScript\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"]],\n",
"ransomware_install_command = {\"commands\":[[\"software_manager\", \"application\", \"install\", \"ransomware-script\"]],\n",
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
]
},
@@ -1220,11 +1168,11 @@
"outputs": [],
"source": [
"# Attempting to install the C2 RansomwareScript\n",
"ransomware_install_command = {\"commands\":[\"software_manager\", \"application\", \"install\", \"RansomwareScript\"],\n",
"ransomware_install_command = {\"commands\":[\"software_manager\", \"application\", \"install\", \"ransomware-script\"],\n",
" \"username\": \"admin\",\n",
" \"password\": \"admin\"}\n",
"\n",
"c2_server: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server.send_command(C2Command.TERMINAL, command_options=ransomware_install_command)"
]
},
@@ -1345,7 +1293,7 @@
"metadata": {},
"outputs": [],
"source": [
"database_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"database_server\")\n",
"database_server: Server = blue_env.game.simulation.network.get_node_by_hostname(\"database-server\")\n",
"database_server.software_manager.file_system.show(full=True)"
]
},
@@ -1391,7 +1339,7 @@
"\n",
"``` YAML\n",
"...\n",
" action: configure_c2_beacon\n",
" action: configure-c2-beacon\n",
" options:\n",
" node_id: 0\n",
" config:\n",
@@ -1446,16 +1394,16 @@
"source": [
"web_server: Server = c2_config_env.game.simulation.network.get_node_by_hostname(\"web_server\")\n",
"web_server.software_manager.install(C2Beacon)\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"C2Beacon\"]\n",
"c2_beacon: C2Beacon = web_server.software_manager.software[\"c2-beacon\"]\n",
"\n",
"client_1: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_1.software_manager.install(C2Server)\n",
"c2_server_1: C2Server = client_1.software_manager.software[\"C2Server\"]\n",
"c2_server_1: C2Server = client_1.software_manager.software[\"c2-server\"]\n",
"c2_server_1.run()\n",
"\n",
"client_2: Computer = c2_config_env.game.simulation.network.get_node_by_hostname(\"client_2\")\n",
"client_2.software_manager.install(C2Server)\n",
"c2_server_2: C2Server = client_2.software_manager.software[\"C2Server\"]\n",
"c2_server_2: C2Server = client_2.software_manager.software[\"c2-server\"]\n",
"c2_server_2.run()"
]
},
@@ -1759,6 +1707,16 @@
"\n",
"display_obs_diffs(tcp_c2_obs, udp_c2_obs, blue_config_env.game.step_counter)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"env.game.agents[\"CustomC2Agent\"].show_history()"
]
}
],
"metadata": {

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

@@ -47,7 +47,7 @@
"source": [
"def make_cfg_have_flat_obs(cfg):\n",
" for agent in cfg['agents']:\n",
" if agent['type'] == \"ProxyAgent\":\n",
" if agent['type'] == \"proxy-agent\":\n",
" agent['agent_settings']['flatten_obs'] = False"
]
},
@@ -76,9 +76,9 @@
" # 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",
" 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",
" return red_str"
@@ -147,36 +147,14 @@
"```yaml\n",
" - ref: data_manipulation_attacker # name of agent\n",
" team: RED # not used, just for human reference\n",
" type: RedDatabaseCorruptingAgent # type of agent - this lets primaite know which agent class to use\n",
" 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: UC2RedObservation\n",
" type: uc2-red-observation # TODO: what\n",
" options:\n",
" nodes: {}\n",
"\n",
" action_space:\n",
" \n",
" # The agent has access to the DataManipulationBoth on clients 1 and 2.\n",
" options:\n",
" nodes:\n",
" - node_name: client_1 # The network should have a node called client_1\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_1 should have DataManipulationBot configured on it\n",
" - node_name: client_2 # The network should have a node called client_2\n",
" applications:\n",
" - application_name: DataManipulationBot # The node client_2 should have DataManipulationBot configured on it\n",
"\n",
" # not important\n",
" max_folders_per_node: 1\n",
" max_files_per_folder: 1\n",
" max_services_per_node: 1\n",
"\n",
" # red agent does not need a reward function\n",
" reward_function:\n",
" reward_components:\n",
" - type: DUMMY\n",
"\n",
" # These actions are passed to the RedDatabaseCorruptingAgent init method, they dictate the schedule of attacks\n",
" agent_settings:\n",
" start_settings:\n",
@@ -211,15 +189,13 @@
" \n",
" # \n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.8 # Probability that port scan is successful\n",
" data_manipulation_p_of_success: 0.8 # Probability that SQL attack is successful\n",
" payload: \"DELETE\" # The SQL query which causes the attack (this has to be DELETE)\n",
" server_ip: 192.168.1.14 # IP address of server hosting the database\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient # Database client must be installed in order for DataManipulationBot to function\n",
" - type: database-client # Database client must be installed in order for DataManipulationBot to function\n",
" options:\n",
" db_server_ip: 192.168.1.14 # IP address of server hosting the database\n",
"```"
@@ -354,19 +330,16 @@
"# Make attack always succeed.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 1.0\n",
" data_manipulation_p_of_success: 1.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" - type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" - type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",
@@ -399,19 +372,16 @@
"# Make attack always fail.\n",
"change = yaml.safe_load(\"\"\"\n",
" applications:\n",
" - ref: data_manipulation_bot\n",
" type: DataManipulationBot\n",
" - type: data-manipulation-bot\n",
" options:\n",
" port_scan_p_of_success: 0.0\n",
" data_manipulation_p_of_success: 0.0\n",
" payload: \"DELETE\"\n",
" server_ip: 192.168.1.14\n",
" - ref: client_1_web_browser\n",
" type: WebBrowser\n",
" - type: web-browser\n",
" options:\n",
" target_url: http://arcd.com/users/\n",
" - ref: client_1_database_client\n",
" type: DatabaseClient\n",
" - type: database-client\n",
" options:\n",
" db_server_ip: 192.168.1.14\n",
"\"\"\")\n",

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

@@ -684,6 +684,15 @@
" print(f\"step: {env.game.step_counter}, Red action: {info['agent_actions']['data_manipulation_attacker'].action}, Blue reward:{reward:.2f}\" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.agents[\"data_manipulation_attacker\"].show_history()"
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -717,7 +726,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
"version": "3.10.11"
}
},
"nbformat": 4,

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

@@ -153,6 +153,49 @@
"PRIMAITE_CONFIG[\"developer_mode\"][\"enabled\"] = was_enabled\n",
"PRIMAITE_CONFIG[\"developer_mode\"][\"output_sys_logs\"] = was_syslogs_enabled"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Viewing Agent history\n",
"\n",
"It's possible to view the actions carried out by an agent for a given training session using the `show_history()` method. By default, this will be all actions apart from DONOTHING actions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"\n",
"# Run the training session to generate some resultant data.\n",
"for i in range(100):\n",
" env.step(0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Calling `.show_history()` should show us when the Data Manipulation used the `NODE_APPLICATION_EXECUTE` action."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"attacker = env.game.agents[\"data_manipulation_attacker\"]\n",
"\n",
"attacker.show_history()"
]
}
],
"metadata": {
@@ -171,7 +214,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.8"
"version": "3.10.11"
}
},
"nbformat": 4,

479
src/primaite/notebooks/How-To-Use-Primaite-Dev-Mode.ipynb Normal file
View File

@@ -0,0 +1,479 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# PrimAITE Developer mode\n",
"\n",
"PrimAITE has built in developer tools.\n",
"\n",
"The dev-mode is designed to help make the development of PrimAITE easier.\n",
"\n",
"`NOTE: For the purposes of the notebook, the commands are preceeded by \"!\". When running the commands, run it without the \"!\".`\n",
"\n",
"To display the available dev-mode options, run the command below:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode --help"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Save the current PRIMAITE_CONFIG to restore after the notebook runs\n",
"\n",
"from primaite import PRIMAITE_CONFIG\n",
"\n",
"temp_config = PRIMAITE_CONFIG.copy()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Dev mode options"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### enable\n",
"\n",
"Enables the dev mode for PrimAITE.\n",
"\n",
"This will enable the developer mode for PrimAITE.\n",
"\n",
"By default, when developer mode is enabled, session logs will be generated in the PRIMAITE_ROOT/sessions folder unless configured to be generated in another location."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode enable"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### disable\n",
"\n",
"Disables the dev mode for PrimAITE.\n",
"\n",
"This will disable the developer mode for PrimAITE."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode disable"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### show\n",
"\n",
"Shows if PrimAITE is running in dev mode or production mode.\n",
"\n",
"The command will also show the developer mode configuration."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode show"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### config\n",
"\n",
"Configure the PrimAITE developer mode"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --help"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### path\n",
"\n",
"Set the path where generated session files will be output.\n",
"\n",
"By default, this value will be in PRIMAITE_ROOT/sessions.\n",
"\n",
"To reset the path to default, run:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config path -root\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config path --default"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --sys-log-level or -slevel\n",
"\n",
"Set the system log level.\n",
"\n",
"This will override the system log level in configurations and will make PrimAITE include the set log level and above.\n",
"\n",
"Available options are:\n",
"- `DEBUG`\n",
"- `INFO`\n",
"- `WARNING`\n",
"- `ERROR`\n",
"- `CRITICAL`\n",
"\n",
"Default value is `DEBUG`\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --sys-log-level DEBUG\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -slevel DEBUG"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --agent-log-level or -alevel\n",
"\n",
"Set the agent log level.\n",
"\n",
"This will override the agent log level in configurations and will make PrimAITE include the set log level and above.\n",
"\n",
"Available options are:\n",
"- `DEBUG`\n",
"- `INFO`\n",
"- `WARNING`\n",
"- `ERROR`\n",
"- `CRITICAL`\n",
"\n",
"Default value is `DEBUG`\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --agent-log-level DEBUG\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -alevel DEBUG"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --output-sys-logs or -sys\n",
"\n",
"If enabled, developer mode will output system logs.\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --output-sys-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -sys"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To disable outputting sys logs:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --no-sys-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -nsys"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --output-agent-logs or -agent\n",
"\n",
"If enabled, developer mode will output agent action logs.\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --output-agent-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -agent"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To disable outputting agent action logs:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --no-agent-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -nagent"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --output-pcap-logs or -pcap\n",
"\n",
"If enabled, developer mode will output PCAP logs.\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --output-pcap-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -pcap"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To disable outputting PCAP logs:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --no-pcap-logs\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -npcap"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### --output-to-terminal or -t\n",
"\n",
"If enabled, developer mode will output logs to the terminal.\n",
"\n",
"Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --output-to-terminal\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -t"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To disable terminal outputs:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config --no-terminal\n",
"\n",
"# or\n",
"\n",
"!primaite dev-mode config -nt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Combining commands\n",
"\n",
"It is possible to combine commands to set the configuration.\n",
"\n",
"This saves having to enter multiple commands and allows for a much more efficient setting of PrimAITE developer mode configurations."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Example of setting system log level and enabling the system logging:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config -slevel WARNING -sys"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Another example where the system log and agent action log levels are set and enabled and should be printed to terminal:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite dev-mode config -slevel ERROR -sys -alevel ERROR -agent -t"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Restore PRIMAITE_CONFIG\n",
"from primaite.utils.cli.primaite_config_utils import update_primaite_application_config\n",
"\n",
"\n",
"global PRIMAITE_CONFIG\n",
"PRIMAITE_CONFIG[\"developer_mode\"] = temp_config[\"developer_mode\"]\n",
"update_primaite_application_config(config=PRIMAITE_CONFIG)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

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

@@ -114,7 +114,7 @@
"metadata": {},
"outputs": [],
"source": [
"print(f\"DNS Client state: {client.software_manager.software.get('DNSClient').operating_state.name}\")"
"print(f\"DNS Client state: {client.software_manager.software.get('dns-client').operating_state.name}\")"
]
},
{

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

@@ -9,6 +9,13 @@
"© Crown-owned copyright 2024, Defence Science and Technology Laboratory UK"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -67,9 +74,9 @@
"source": [
"network: Network = basic_network()\n",
"computer_a: Computer = network.get_node_by_hostname(\"node_a\")\n",
"terminal_a: Terminal = computer_a.software_manager.software.get(\"Terminal\")\n",
"terminal_a: Terminal = computer_a.software_manager.software.get(\"terminal\")\n",
"computer_b: Computer = network.get_node_by_hostname(\"node_b\")\n",
"terminal_b: Terminal = computer_b.software_manager.software.get(\"Terminal\")"
"terminal_b: Terminal = computer_b.software_manager.software.get(\"terminal\")"
]
},
{
@@ -121,7 +128,7 @@
"metadata": {},
"outputs": [],
"source": [
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"RansomwareScript\"])"
"term_a_term_b_remote_connection.execute([\"software_manager\", \"application\", \"install\", \"ransomware-script\"])"
]
},
{
@@ -169,6 +176,22 @@
"computer_b.file_system.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Information about the latest response when executing a remote command can be seen by calling the `last_response` attribute within `Terminal`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(terminal_a.last_response)"
]
},
{
"cell_type": "markdown",
"metadata": {},
@@ -207,6 +230,263 @@
"source": [
"computer_b.user_session_manager.show(include_historic=True, include_session_id=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Game Layer Implementation\n",
"\n",
"This notebook section will detail the implementation of how the game layer utilises the terminal to support different agent actions.\n",
"\n",
"The ``Terminal`` is used in a variety of different ways in the game layer. Specifically, the terminal is leveraged to implement the following actions:\n",
"\n",
"\n",
"| Game Layer Action | Simulation Layer |\n",
"|-----------------------------------|--------------------------|\n",
"| ``node-send-local-command`` | Uses the given user credentials, creates a ``LocalTerminalSession`` and executes the given command and returns the ``RequestResponse``.\n",
"| ``node-session-remote-login`` | Uses the given user credentials and remote IP to create a ``RemoteTerminalSession``.\n",
"| ``node-send-remote-command`` | Uses the given remote IP to locate the correct ``RemoteTerminalSession``, executes the given command and returns the ``RequestsResponse``."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Game Layer Setup\n",
"\n",
"Similar to other notebooks, the next code cells create a custom proxy agent to demonstrate how these commands can be leveraged by agents in the ``UC2`` network environment.\n",
"\n",
"If you're unfamiliar with ``UC2`` then please refer to the [UC2-E2E-Demo notebook for further reference](./Data-Manipulation-E2E-Demonstration.ipynb)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import yaml\n",
"from primaite.config.load import data_manipulation_config_path\n",
"from primaite.session.environment import PrimaiteGymEnv"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"custom_terminal_agent = \"\"\"\n",
" - ref: CustomC2Agent\n",
" team: RED\n",
" type: 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",
" options: {}\n",
" 1:\n",
" action: node-send-local-command\n",
" options:\n",
" node_name: client_1\n",
" username: admin\n",
" password: admin\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - dog.png\n",
" - False\n",
" 2:\n",
" action: node-session-remote-login\n",
" options:\n",
" node_name: client_1\n",
" username: admin\n",
" password: admin\n",
" remote_ip: 192.168.10.22\n",
" 3:\n",
" action: node-send-remote-command\n",
" options:\n",
" node_name: client_1\n",
" remote_ip: 192.168.10.22\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - cat.png\n",
" - False\n",
"\"\"\"\n",
"custom_terminal_agent_yaml = yaml.safe_load(custom_terminal_agent)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(data_manipulation_config_path()) as f:\n",
" cfg = yaml.safe_load(f)\n",
" # removing all agents & adding the custom agent.\n",
" cfg['agents'] = {}\n",
" cfg['agents'] = custom_terminal_agent_yaml\n",
"\n",
"env = PrimaiteGymEnv(env_config=cfg)\n",
"\n",
"client_1: Computer = env.game.simulation.network.get_node_by_hostname(\"client_1\")\n",
"client_2: Computer = env.game.simulation.network.get_node_by_hostname(\"client_2\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``node-send-local-command`` \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: node-send-local-command\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: node-send-local-command\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" username: admin\n",
" password: admin\n",
" command:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - dog.png\n",
" - False\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(1)\n",
"client_1.file_system.show(full=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``node-session-remote-login`` \n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: node-session-remote-login\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 2:\n",
" action: node-session-remote-login\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" username: admin\n",
" password: admin\n",
" remote_ip: 192.168.10.22 # client_2's ip address.\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(2)\n",
"client_2.session_manager.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Terminal Action | ``node-send-remote-command``\n",
"\n",
"The yaml snippet below shows all the relevant agent options for this action:\n",
"\n",
"```yaml\n",
"\n",
" action_space:\n",
" action_list:\n",
" ...\n",
" - type: node-send-remote-command\n",
" ...\n",
" options:\n",
" nodes: # Node List\n",
" - node_name: client_1\n",
" ...\n",
" ...\n",
" action_map:\n",
" 1:\n",
" action: node-send-remote-command\n",
" options:\n",
" node_id: 0 # Index 0 at the node list.\n",
" remote_ip: 192.168.10.22\n",
" commands:\n",
" - file_system\n",
" - create\n",
" - file\n",
" - downloads\n",
" - cat.png\n",
" - False\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(3)\n",
"client_2.file_system.show(full=True)"
]
}
],
"metadata": {

1621
src/primaite/notebooks/UC7-E2E-Demo.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

1843
src/primaite/notebooks/UC7-TAP001-Kill-Chain-E2E.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

1687
src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

156
src/primaite/notebooks/UC7-Training.ipynb Normal file
View File

@@ -0,0 +1,156 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": [
"# Training an SB3 Agent\n",
"\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."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### First, we import the inital packages and read in our configuration file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import yaml\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite import PRIMAITE_PATHS\n",
"from prettytable import PrettyTable\n",
"from deepdiff.diff import DeepDiff\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.network.hardware.nodes.network.router import Router\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"\n",
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/uc7_config.yaml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"gym = PrimaiteGymEnv(env_config=scenario_path)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from stable_baselines3 import PPO\n",
"\n",
"# EPISODE_LEN = 128\n",
"EPISODE_LEN = 128\n",
"NUM_EPISODES = 10\n",
"NO_STEPS = EPISODE_LEN * NUM_EPISODES\n",
"BATCH_SIZE = 32\n",
"LEARNING_RATE = 3e-4"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = PPO('MlpPolicy', gym, learning_rate=LEARNING_RATE, n_steps=NO_STEPS, batch_size=BATCH_SIZE, verbose=0, tensorboard_log=\"./PPO_UC7/\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model.learn(total_timesteps=NO_STEPS)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model.save(\"PrimAITE-PPO-UC7-example-agent\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"eval_model = PPO(\"MlpPolicy\", gym)\n",
"eval_model = PPO.load(\"PrimAITE-PPO-UC7-example-agent\", gym)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from stable_baselines3.common.evaluation import evaluate_policy\n",
"\n",
"evaluate_policy(eval_model, gym, n_eval_episodes=1)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

586
src/primaite/notebooks/UC7-attack-variants.ipynb Normal file
View File

@@ -0,0 +1,586 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# UC7 with Attack Variability\n",
"\n",
"© Crown-owned copyright 2025, Defence Science and Technology Laboratory UK"
]
},
{
"cell_type": "markdown",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"source": [
"This notebook demonstrates the PrimAITE environment with the UC7 network laydown and multiple attack personas. The first attack persona is TAP001 which performs a ransomware attack against the database. The other one is TAP003 which is able to maliciously add ACL rules that block green pattern of life.\n",
"\n",
"The environment switches between these two attacks on a pre-defined schedule which is defined in the schedule.yaml file of the scenario folder."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup and Imports"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import yaml\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite import PRIMAITE_PATHS\n",
"from prettytable import PrettyTable\n",
"from deepdiff.diff import DeepDiff\n",
"from primaite.session.environment import PrimaiteGymEnv\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.network.hardware.nodes.network.router import Router\n",
"from primaite.simulator.system.services.dns.dns_server import DNSServer\n",
"from primaite.simulator.system.software import SoftwareHealthState\n",
"from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus\n",
"from primaite.simulator.network.hardware.nodes.network.switch import Switch\n",
"from primaite.simulator.system.applications.web_browser import WebBrowser\n",
"from primaite.simulator.network.container import Network\n",
"from primaite.simulator.system.services.service import ServiceOperatingState\n",
"from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState\n",
"from primaite.simulator.system.services.database.database_service import DatabaseService\n",
"from primaite.simulator.system.applications.database_client import DatabaseClient\n",
"from primaite.simulator.network.hardware.nodes.network.firewall import Firewall\n",
"from primaite.game.game import PrimaiteGame\n",
"from primaite.simulator.sim_container import Simulation\n",
"from primaite.config.load import load, _EXAMPLE_CFG\n",
"from primaite.simulator.network.hardware.nodes.host.server import Server\n",
"from primaite.simulator.network.hardware.nodes.network.router import Router\n",
"from primaite.simulator.network.hardware.nodes.host.computer import Computer\n",
"\n",
"scenario_path = PRIMAITE_PATHS.user_config_path / \"example_config/uc7_multiple_attack_variants\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env = PrimaiteGymEnv(env_config=scenario_path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Schedule\n",
"\n",
"Let's print the schedule so that we can see which attack we can expect on each episode.\n",
"\n",
"On episodes 0-4, the TAP001 agent will be used, and on episodes 5-9, the TAP003 agent will be used. Then, the environment will alternate between the two. Furthermore, the TAP001 agent will alternate between starting at `ST-PROJ-A-PRV-PC-1`, `ST-PROJ-B-PRV-PC-2`, `ST-PROJ-C-PRV-PC-3`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open(scenario_path / \"schedule.yaml\",'r') as f:\n",
" print(f.read())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP001 attack\n",
"\n",
"Let's first demonstrate the TAP001 attack. We will let the environment run for 30 steps and print out the red agent's actions.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#utils\n",
"def run_green_and_red_pol(num_steps):\n",
" for i in range(num_steps): # perform steps\n",
" env.step(0)\n",
"\n",
"def print_agent_actions_except_do_nothing(agent_name):\n",
" \"\"\"Get the agent's action history, filter out DONOTHING actions, print relevant data in a table.\"\"\"\n",
" table = PrettyTable()\n",
" table.field_names = [\"Step\", \"Action\", \"Node\", \"Application\", \"Target IP\", \"Response\"]\n",
" print(f\"Episode: {env.episode_counter}, Actions for '{agent_name}':\")\n",
" for item in env.game.agents[agent_name].history:\n",
" if item.action == \"do-nothing\":\n",
" continue\n",
"\n",
" node, application, target_ip = \"N/A\", \"N/A\", \"N/A\",\n",
"\n",
" if item.action.startswith(\"node-map\"):\n",
" node = item.parameters['source-node']\n",
" application = \"nmap\"\n",
" target_ip = str(item.parameters['target_ip_address'])\n",
" target_ip = (target_ip[:25]+'...') if len(target_ip)>25 else target_ip # truncate long string\n",
"\n",
" elif item.action == \"router-acl-add-rule\":\n",
" node = item.parameters.get(\"router_name\")\n",
" elif item.action == \"node-send-remote-command\":\n",
" node = item.parameters.get(\"node_name\")\n",
" target_ip = item.parameters.get(\"remote_ip\")\n",
" application = item.parameters.get(\"command\")\n",
" elif item.action == \"node-session-remote-login\":\n",
" node = item.parameters.get(\"node_name\")\n",
" target_ip = item.parameters.get(\"remote_ip\")\n",
" application = \"user-manager\"\n",
" elif item.action.startswith(\"c2-server\"):\n",
" application = \"c2-server\"\n",
" node = item.parameters.get('node_name')\n",
" elif item.action == \"configure-c2-beacon\":\n",
" application = \"c2-beacon\"\n",
" node = item.parameters.get('node_name')\n",
"\n",
" else:\n",
" if (node_id := item.parameters.get('node_id')) is not None:\n",
" node = env.game.agents[agent_name].action_manager.node_names[node_id]\n",
" if (application_id := item.parameters.get('application_id')) is not None:\n",
" application = env.game.agents[agent_name].action_manager.application_names[node_id][application_id]\n",
" if (application_name := item.parameters.get('application_name')) is not None:\n",
" application = application_name\n",
"\n",
" table.add_row([item.timestep, item.action, node, application, target_ip, item.response.status])\n",
"\n",
" print(table)\n",
" print(\"(Any DONOTHING actions are omitted)\")\n",
"\n",
"def finish_episode_and_print_reward():\n",
" while env.game.step_counter < 128:\n",
" env.step(0)\n",
" print(f\"Total reward this episode: {env.agent.reward_function.total_reward:2f}\")\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"run_green_and_red_pol(110)\n",
"print_agent_actions_except_do_nothing(\"attacker\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_data_prv_srv_db: Server = env.game.simulation.network.get_node_by_hostname(\"ST-DATA-PRV-SRV-DB\")\n",
"st_data_prv_srv_db.file_system.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP001 Prevention"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The blue agent should be able to prevent the ransomware attack by blocking the red agent's access to the database. Let's run the environment until the observation space shows symptoms of the attack starting.\n",
"\n",
"Because we are in episode index 1, the red agent will use `ST-PROJ-A-PRV-PC-1` to start the attack. On step 25, the red agent installs `RansomwareScript`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"obs, reward, term, trunc, info = env.step(0)\n",
"for i in range(25): # we know that the ransomware install happens at step 25\n",
" old = obs\n",
" obs, reward, term, trunc, info = env.step(0)\n",
" new = obs\n",
"\n",
"diff = DeepDiff(old,new)\n",
"print(f\"Step {env.game.step_counter}\") # it's step 26 now because the step counter is incremented after the step\n",
"for d,v in diff.get('values_changed', {}).items():\n",
" print(f\"{d}: {v['old_value']} -> {v['new_value']}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that on HOST0, application index 1 has gone from `operating_status` 0 to 3, meaning there wasn't an application before, but now there is an application in the `INSTALLING` state. The blue agent should be able to detect this and block the red agent's access to the database. Action 43 will block `ST-PROJ-A-PRV-PC-1` from sending POSTGRES traffic to the DB server.\n",
"\n",
"If this were a different episode, it could have been `ST-PROJ-B-PRV-PC-2` or `ST-PROJ-C-PRV-PC-3` that are affected, and a different defensive action would be required."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(43)\n",
"env.step(45)\n",
"env.step(47)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_intra_prv_rt_cr: Router = env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-CR\")\n",
"st_intra_prv_rt_cr.acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"st_intra_prv_rt_cr.acl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now TAP001 is unable to locate the database!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print_agent_actions_except_do_nothing(\"attacker\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TAP003 attack\n",
"\n",
"Let's skip until episode 5 and demonstrate the TAP003 attack. We will let the environment run and print out the red agent's actions.\n",
"\n",
"By default, TAP003 will add the following rules:\n",
"\n",
"|Target Router | Impact |\n",
"|----------------------|--------|\n",
"|`ST-INTRA-PRV-RT-DR-1`| Blocks all `POSTGRES_SERVER` that arrives at the `ST-INTRA-PRV-RT-DR-1` router. This rule will prevent all ST_PROJ_* hosts from accessing the database (`ST-DATA-PRV-SRV-DB`).|\n",
"|`ST-INTRA-PRV-RT-CR`| Blocks all `HTTP` traffic that arrives at the`ST-INTRA-PRV-RT-CR` router. This rule will prevent all SOME_TECH hosts from accessing the webserver (`ST-DMZ-PUB-SRV-WEB`)|\n",
"|`REM-PUB-RT-DR`| Blocks all `DNS` traffic that arrives at the `REM-PUB-RT-DR` router. This rule prevents any remote site works from accessing the DNS Server (`ISP-PUB-SRV-DNS`).|"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"while env.episode_counter < 5:\n",
" env.reset()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"run_green_and_red_pol(128)\n",
"print_agent_actions_except_do_nothing(\"attacker\")\n",
"obs, reward, term, trunc, info = env.step(0); # one more step so we can capture the value of `obs`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The agent selected to add ACL rules that will prevent green pattern of life by blocking a variety of different traffic. This has a negative impact on reward. Let's view the ACL list on the affected router."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-DR-1\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-CR\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that at indices 1-5, there are ACL rules that block all traffic. The blue agent can see this rule in the `ROUTERS` part of the observation space.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER0']['ACL'][1]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER1']['ACL'][1]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"obs['NODES']['ROUTER2']['ACL'][1]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Preventing TAP003 attack\n",
"\n",
"The blue agent can prevent the red agent from adding ACL rules. TAP003 relies on connecting to the router via SSH, and sending remote ACL_ADDRULE requests. The blue agent can prevent this by pre-emptively changing the admin password on the affected routers or by blocking SSH traffic between the red agent's starting node and the target routers."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"obs, reward, term, trunc, info = env.step(0)\n",
"old = obs\n",
"for i in range(128): \n",
" obs, reward, term, trunc, info = env.step(0)\n",
" new = obs\n",
"\n",
"diff = DeepDiff(old,new)\n",
"print(f\"Step {env.game.step_counter}\") # it's the next step now because the step counter is incremented after the step\n",
"for d,v in diff.get('values_changed', {}).items():\n",
" print(f\"{d}: {v['old_value']} -> {v['new_value']}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By printing the reward of each individual agent, we will see what green agents are affected the most. Of course, these green rewards count towards the blue reward so ultimately the blue agent should learn to remove the ACL rule."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The most effective option that the blue agent has against TAP003 is to prevent the red agent from ever adding the ACLs in the first place through blocking the SSH connection."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"env.step(51) # SSH Blocking ACL on ST-INRA-PRV-RT-R1\n",
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, another option the blue agent can take is to change the passwords of the different target routers that TAP003 will attack through the `NODE_ACCOUNTS_CHANGE_PASSWORD` action."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"env.step(50) # NODE_ACCOUNTS_CHANGE_PASSWORD | st-intra-prv-rt-cr\n",
"env.step(52) # NODE_ACCOUNTS_CHANGE_PASSWORD | st-intra-prv-rt-dr-1\n",
"env.step(54) # NODE_ACCOUNTS_CHANGE_PASSWORD | rem-pub-rt-dr\n",
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lastly, the blue agent can remedy the impacts of TAP003 through removing the malicious ACLs that TAP003 adds."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()\n",
"\n",
"# Allow TAP003 to add it's malicious rules\n",
"for _ in range(45):\n",
" env.step(0)\n",
"\n",
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-CR\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-DR-1\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.step(44) # ROUTER_ACL_REMOVERULE | st-intra-prv-rt-cr\n",
"env.step(53) # ROUTER_ACL_REMOVERULE | st-intra-prv-rt-dr-1\n",
"env.step(55) # ROUTER_ACL_REMOVERULE | rem-pub-rt-dr"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-CR\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"ST-INTRA-PRV-RT-DR-1\").acl.show()\n",
"env.game.simulation.network.get_node_by_hostname(\"REM-PUB-RT-DR\").acl.show()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"finish_episode_and_print_reward()\n",
"\n",
"for ag in env.game.agents.values():\n",
" print(ag.config.ref, ag.reward_function.total_reward)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.12"
}
},
"nbformat": 4,
"nbformat_minor": 2
}

1141
src/primaite/notebooks/UC7-network_connectivity.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
src/primaite/notebooks/_package_data/uc7/blue_agent_action_and_obs.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.4 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_network.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 159 KiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_network_detailed.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 609 KiB

4
src/primaite/notebooks/_package_data/uc7/uc7_network_detailed_svg.svg Normal file
View File

File diff suppressed because one or more lines are too long
Side by Side

After

Width:  |  Height:  |  Size: 383 KiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_activate.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.0 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_c2.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_download_install.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.0 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_main.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_payload.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_propagate_1.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.0 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_propagate_2.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_propagate_3.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001/uc7_tap001_propagate_4.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap001_training_static_red.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 70 KiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap003/uc7_tap003_main.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

BIN
src/primaite/notebooks/_package_data/uc7/uc7_tap003_training_static_red.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 61 KiB

30
src/primaite/session/environment.py
View File

@@ -26,14 +26,26 @@ except ModuleNotFoundError:
_LOGGER.debug("Torch not available for importing")
def set_random_seed(seed: int) -> Union[None, int]:
def set_random_seed(seed: int, generate_seed_value: bool) -> Union[None, int]:
"""
Set random number generators.
If seed is None or -1 and generate_seed_value is True randomly generate a
seed value.
If seed is > -1 and generate_seed_value is True ignore the latter and use
the provide seed value.
:param seed: int
:param generate_seed_value: bool
:return: None or the int representing the seed used.
"""
if seed is None or seed == -1:
return None
if generate_seed_value:
rng = np.random.default_rng()
# 2**32-1 is highest value for python RNG seed.
seed = int(rng.integers(low=0, high=2**32 - 1))
else:
return None
elif seed < -1:
raise ValueError("Invalid random number seed")
# Seed python RNG
@@ -50,6 +62,13 @@ def set_random_seed(seed: int) -> Union[None, int]:
return seed
def log_seed_value(seed: int):
"""Log the selected seed value to file."""
path = SIM_OUTPUT.path / "seed.log"
with open(path, "w") as file:
file.write(f"Seed value = {seed}")
class PrimaiteGymEnv(gymnasium.Env):
"""
Thin wrapper env to provide agents with a gymnasium API.
@@ -65,7 +84,8 @@ class PrimaiteGymEnv(gymnasium.Env):
"""Object that returns a config corresponding to the current episode."""
self.seed = self.episode_scheduler(0).get("game", {}).get("seed")
"""Get RNG seed from config file. NB: Must be before game instantiation."""
self.seed = set_random_seed(self.seed)
self.generate_seed_value = self.episode_scheduler(0).get("game", {}).get("generate_seed_value")
self.seed = set_random_seed(self.seed, self.generate_seed_value)
self.io = PrimaiteIO.from_config(self.episode_scheduler(0).get("io_settings", {}))
"""Handles IO for the environment. This produces sys logs, agent logs, etc."""
self.game: PrimaiteGame = PrimaiteGame.from_config(self.episode_scheduler(0))
@@ -79,6 +99,8 @@ class PrimaiteGymEnv(gymnasium.Env):
_LOGGER.info(f"PrimaiteGymEnv RNG seed = {self.seed}")
log_seed_value(self.seed)
def action_masks(self) -> np.ndarray:
"""
Return the action mask for the agent.
@@ -146,7 +168,7 @@ class PrimaiteGymEnv(gymnasium.Env):
f"avg. reward: {self.agent.reward_function.total_reward}"
)
if seed is not None:
set_random_seed(seed)
set_random_seed(seed, self.generate_seed_value)
self.total_reward_per_episode[self.episode_counter] = self.agent.reward_function.total_reward
if self.io.settings.save_agent_actions:

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

@@ -864,7 +864,21 @@ class UserManager(Service, discriminator="user-manager"):
"""
rm = super()._init_request_manager()
# todo add doc about requeest schemas
# todo add doc about request schemas
rm.add_request(
"add_user",
RequestType(
func=lambda request, context: RequestResponse.from_bool(
self.add_user(username=request[0], password=request[1], is_admin=request[2])
)
),
)
rm.add_request(
"disable_user",
RequestType(
func=lambda request, context: RequestResponse.from_bool(self.disable_user(username=request[0]))
),
)
rm.add_request(
"change_password",
RequestType(
@@ -1572,7 +1586,7 @@ class Node(SimComponent, ABC):
operating_state: Any = None
users: Any = None # Temporary to appease "extra=forbid"
users: List[Dict] = [] # Temporary to appease "extra=forbid"
config: ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema())
"""Configuration items within Node"""
@@ -1638,6 +1652,8 @@ class Node(SimComponent, ABC):
self._install_system_software()
self.session_manager.node = self
self.session_manager.software_manager = self.software_manager
for user in self.config.users:
self.user_manager.add_user(**user, bypass_can_perform_action=True)
@property
def user_manager(self) -> Optional[UserManager]:
@@ -1769,7 +1785,7 @@ class Node(SimComponent, ABC):
"""
application_name = request[0]
if self.software_manager.software.get(application_name):
self.sys_log.warning(f"Can't install {application_name}. It's already installed.")
self.sys_log.info(f"Can't install {application_name}. It's already installed.")
return RequestResponse(status="success", data={"reason": "already installed"})
application_class = Application._registry[application_name]
self.software_manager.install(application_class)

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

@@ -2,11 +2,12 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Literal, Optional
from typing import Any, ClassVar, Dict, List, Literal, Optional
from pydantic import Field
from primaite import getLogger
from primaite.simulator.file_system.file_type import FileType
from primaite.simulator.network.hardware.base import (
IPWiredNetworkInterface,
Link,
@@ -313,7 +314,7 @@ class HostNode(Node, discriminator="host-node"):
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"HostARP": HostARP,
"host-arp": HostARP,
"icmp": ICMP,
"dns-client": DNSClient,
"ntp-client": NTPClient,
@@ -339,7 +340,7 @@ class HostNode(Node, discriminator="host-node"):
ip_address: IPV4Address
services: Any = None # temporarily unset to appease extra="forbid"
applications: Any = None # temporarily unset to appease extra="forbid"
folders: Any = None # temporarily unset to appease extra="forbid"
folders: List[Dict] = {} # temporarily unset to appease extra="forbid"
network_interfaces: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema())
@@ -348,6 +349,18 @@ class HostNode(Node, discriminator="host-node"):
super().__init__(**kwargs)
self.connect_nic(NIC(ip_address=kwargs["config"].ip_address, subnet_mask=kwargs["config"].subnet_mask))
for folder in self.config.folders:
# handle empty foler defined by just a string
self.file_system.create_folder(folder["folder_name"])
for file in folder.get("files", []):
self.file_system.create_file(
folder_name=folder["folder_name"],
file_name=file["file_name"],
size=file.get("size", 0),
file_type=FileType[file.get("type", "UNKNOWN").upper()],
)
@property
def nmap(self) -> Optional[NMAP]:
"""

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

@@ -49,7 +49,7 @@ class Firewall(Router, discriminator="firewall"):
Example:
>>> from primaite.simulator.network.transmission.network_layer import IPProtocol
>>> from primaite.simulator.network.transmission.transport_layer import Port
>>> from primaite.utils.validation.port import Port
>>> firewall = Firewall(hostname="Firewall1")
>>> firewall.configure_internal_port(ip_address="192.168.1.1", subnet_mask="255.255.255.0")
>>> firewall.configure_external_port(ip_address="10.0.0.1", subnet_mask="255.255.255.0")

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

@@ -467,6 +467,7 @@ class AccessControlList(SimComponent):
"""Check if a packet with the given properties is permitted through the ACL."""
permitted = False
rule: ACLRule = None
for _rule in self._acl:
if not _rule:
continue
@@ -1215,9 +1216,9 @@ class Router(NetworkNode, discriminator="router"):
config: ConfigSchema = Field(default_factory=lambda: Router.ConfigSchema())
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"UserSessionManager": UserSessionManager,
"UserManager": UserManager,
"Terminal": Terminal,
"user-session-manager": UserSessionManager,
"user-manager": UserManager,
"terminal": Terminal,
}
network_interfaces: Dict[str, RouterInterface] = {}
@@ -1385,6 +1386,12 @@ class Router(NetworkNode, discriminator="router"):
return False
def subject_to_acl(self, frame: Frame) -> bool:
"""Check that frame is subject to ACL rules."""
if frame.ip.protocol == "udp" and frame.is_arp:
return False
return True
def receive_frame(self, frame: Frame, from_network_interface: RouterInterface):
"""
Processes an incoming frame received on one of the router's interfaces.
@@ -1398,8 +1405,12 @@ class Router(NetworkNode, discriminator="router"):
if self.operating_state != NodeOperatingState.ON:
return
# Check if it's permitted
permitted, rule = self.acl.is_permitted(frame)
if self.subject_to_acl(frame=frame):
# Check if it's permitted
permitted, rule = self.acl.is_permitted(frame)
else:
permitted = True
rule = None
if not permitted:
at_port = self._get_port_of_nic(from_network_interface)

2
src/primaite/simulator/network/transmission/data_link_layer.py
View File

@@ -163,7 +163,7 @@ class Frame(BaseModel):
"""
Checks if the Frame is an ARP (Address Resolution Protocol) packet.
This is determined by checking if the destination port of the TCP header is equal to the ARP port.
This is determined by checking if the destination and source port of the UDP header is equal to the ARP port.
:return: True if the Frame is an ARP packet, otherwise False.
"""

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

@@ -55,7 +55,7 @@ class ARP(Service, discriminator="arp"):
:param markdown: If True, format the output as Markdown. Otherwise, use plain text.
"""
table = PrettyTable(["IP Address", "MAC Address", "Via"])
table = PrettyTable(["IP Address", "MAC Address", "Via", "Port"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -66,6 +66,7 @@ class ARP(Service, discriminator="arp"):
str(ip),
arp.mac_address,
self.software_manager.node.network_interfaces[arp.network_interface_uuid].mac_address,
self.software_manager.node.network_interfaces[arp.network_interface_uuid].port_num,
]
)
print(table)

78
src/primaite/simulator/system/services/terminal/terminal.py
View File

@@ -142,12 +142,20 @@ class Terminal(Service, discriminator="terminal"):
_client_connection_requests: Dict[str, Optional[Union[str, TerminalClientConnection]]] = {}
"""Dictionary of connect requests made to remote nodes."""
_last_response: Optional[RequestResponse] = None
"""Last response received from RequestManager, for returning remote RequestResponse."""
def __init__(self, **kwargs):
kwargs["name"] = "terminal"
kwargs["port"] = PORT_LOOKUP["SSH"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@property
def last_response(self) -> Optional[RequestResponse]:
"""Public version of _last_response attribute."""
return self._last_response
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -186,7 +194,7 @@ class Terminal(Service, discriminator="terminal"):
return RequestResponse(status="failure", data={})
rm.add_request(
"node-session-remote-login",
"node_session_remote_login",
request_type=RequestType(func=_remote_login),
)
@@ -209,28 +217,45 @@ class Terminal(Service, discriminator="terminal"):
command: str = request[1]["command"]
remote_connection = self._get_connection_from_ip(ip_address=ip_address)
if remote_connection:
outcome = remote_connection.execute(command)
if outcome:
return RequestResponse(
status="success",
data={},
)
else:
return RequestResponse(
status="failure",
data={},
)
remote_connection.execute(command)
return self.last_response if not None else RequestResponse(status="failure", data={})
return RequestResponse(
status="failure",
data={"reason": "Failed to execute command."},
)
rm.add_request(
"send_remote_command",
request_type=RequestType(func=remote_execute_request),
)
def local_execute_request(request: RequestFormat, context: Dict) -> RequestResponse:
"""Executes a command using a local terminal session."""
command: str = request[2]["command"]
local_connection = self._process_local_login(username=request[0], password=request[1])
if local_connection:
outcome = local_connection.execute(command)
if outcome:
return RequestResponse(
status="success",
data={"reason": outcome},
)
return RequestResponse(
status="success",
data={"reason": "Local Terminal failed to resolve command. Potentially invalid credentials?"},
)
rm.add_request(
"send_local_command",
request_type=RequestType(func=local_execute_request),
)
return rm
def execute(self, command: List[Any]) -> Optional[RequestResponse]:
"""Execute a passed ssh command via the request manager."""
return self.parent.apply_request(command)
self._last_response = self.parent.apply_request(command)
return self._last_response
def _get_connection_from_ip(self, ip_address: IPv4Address) -> Optional[RemoteTerminalConnection]:
"""Find Remote Terminal Connection from a given IP."""
@@ -409,6 +434,8 @@ class Terminal(Service, discriminator="terminal"):
"""
source_ip = kwargs["frame"].ip.src_ip_address
self.sys_log.info(f"{self.name}: Received payload: {payload}. Source: {source_ip}")
self._last_response = None # Clear last response
if isinstance(payload, SSHPacket):
if payload.transport_message == SSHTransportMessage.SSH_MSG_USERAUTH_REQUEST:
# validate & add connection
@@ -457,6 +484,9 @@ class Terminal(Service, discriminator="terminal"):
session_id=session_id,
source_ip=source_ip,
)
self._last_response: RequestResponse = RequestResponse(
status="success", data={"reason": "Login Successful"}
)
elif payload.transport_message == SSHTransportMessage.SSH_MSG_SERVICE_REQUEST:
# Requesting a command to be executed
@@ -468,12 +498,32 @@ class Terminal(Service, discriminator="terminal"):
payload.connection_uuid
)
remote_session.last_active_step = self.software_manager.node.user_session_manager.current_timestep
self.execute(command)
self._last_response: RequestResponse = self.execute(command)
if self._last_response.status == "success":
transport_message = SSHTransportMessage.SSH_MSG_SERVICE_SUCCESS
else:
transport_message = SSHTransportMessage.SSH_MSG_SERVICE_FAILED
payload: SSHPacket = SSHPacket(
payload=self._last_response,
transport_message=transport_message,
connection_message=SSHConnectionMessage.SSH_MSG_CHANNEL_DATA,
)
self.software_manager.send_payload_to_session_manager(
payload=payload, dest_port=self.port, session_id=session_id
)
return True
else:
self.sys_log.error(
f"{self.name}: Connection UUID:{payload.connection_uuid} is not valid. Rejecting Command."
)
elif (
payload.transport_message == SSHTransportMessage.SSH_MSG_SERVICE_SUCCESS
or SSHTransportMessage.SSH_MSG_SERVICE_FAILED
):
# Likely receiving command ack from remote.
self._last_response = payload.payload
if isinstance(payload, dict) and payload.get("type"):
if payload["type"] == "disconnect":

55
src/primaite/simulator/system/services/web_server/web_server.py
View File

@@ -117,37 +117,44 @@ class WebServer(Service, discriminator="web-server"):
:type: payload: HttpRequestPacket
"""
response = HttpResponsePacket(status_code=HttpStatusCode.NOT_FOUND, payload=payload)
try:
parsed_url = urlparse(payload.request_url)
path = parsed_url.path.strip("/")
if len(path) < 1:
parsed_url = urlparse(payload.request_url)
path = parsed_url.path.strip("/") if parsed_url and parsed_url.path else ""
if len(path) < 1:
# query succeeded
response.status_code = HttpStatusCode.OK
if path.startswith("users"):
# get data from DatabaseServer
# get all users
if not self._establish_db_connection():
# unable to create a db connection
response.status_code = HttpStatusCode.INTERNAL_SERVER_ERROR
return response
if self.db_connection.query("SELECT"):
# query succeeded
self.set_health_state(SoftwareHealthState.GOOD)
response.status_code = HttpStatusCode.OK
else:
self.set_health_state(SoftwareHealthState.COMPROMISED)
return response
if path.startswith("users"):
# get data from DatabaseServer
# get all users
if not self.db_connection:
self._establish_db_connection()
if self.db_connection.query("SELECT"):
# query succeeded
self.set_health_state(SoftwareHealthState.GOOD)
response.status_code = HttpStatusCode.OK
else:
self.set_health_state(SoftwareHealthState.COMPROMISED)
return response
except Exception: # TODO: refactor this. Likely to cause silent bugs. (ADO ticket #2345 )
# something went wrong on the server
response.status_code = HttpStatusCode.INTERNAL_SERVER_ERROR
return response
def _establish_db_connection(self) -> None:
def _establish_db_connection(self) -> bool:
"""Establish a connection to db."""
# if active db connection, return true
if self.db_connection:
return True
# otherwise, try to create db connection
db_client = self.software_manager.software.get("database-client")
if db_client is None:
return False # database client not installed
self.db_connection: DatabaseClientConnection = db_client.get_new_connection()
return self.db_connection is not None
def send(
self,

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

@@ -25,7 +25,19 @@ game:
- ICMP
- TCP
- UDP
thresholds:
nmne:
high: 100
medium: 25
low: 5
file_access:
high: 10
medium: 5
low: 2
app_executions:
high: 5
medium: 3
low: 2
agents:
- ref: client_2_green_user
team: GREEN
@@ -64,10 +76,16 @@ agents:
options:
hosts:
- hostname: client_1
applications:
- application_name: WebBrowser
folders:
- folder_name: root
files:
- file_name: "test.txt"
- hostname: client_2
- hostname: client_3
num_services: 1
num_applications: 0
num_applications: 1
num_folders: 1
num_files: 1
num_nics: 2
@@ -182,6 +200,10 @@ simulation:
options:
ntp_server_ip: 192.168.1.10
- type: ntp-server
folders:
- folder_name: root
files:
- file_name: test.txt
- hostname: client_2
type: computer
ip_address: 192.168.10.22

226
tests/assets/configs/nodes_with_initial_files.yaml Normal file
View File

@@ -0,0 +1,226 @@
# Basic Switched network
#
# -------------- -------------- --------------
# | client_1 |------| switch_1 |------| client_2 |
# -------------- -------------- --------------
#
io_settings:
save_step_metadata: false
save_pcap_logs: true
save_sys_logs: true
sys_log_level: WARNING
agent_log_level: INFO
save_agent_logs: true
write_agent_log_to_terminal: True
game:
max_episode_length: 256
ports:
- ARP
- DNS
- HTTP
- POSTGRES_SERVER
protocols:
- ICMP
- TCP
- UDP
agents:
- ref: client_2_green_user
team: GREEN
type: periodic-agent
action_space:
action_map:
0:
action: do-nothing
options: {}
1:
action: node-application-execute
options:
node_id: 0
application_id: 0
agent_settings:
possible_start_nodes: [client_2,]
target_application: web-browser
start_step: 5
frequency: 4
variance: 3
- ref: defender
team: BLUE
type: proxy-agent
observation_space:
type: custom
options:
components:
- type: nodes
label: NODES
options:
hosts:
- hostname: client_1
- hostname: client_2
- hostname: client_3
num_services: 1
num_applications: 0
num_folders: 1
num_files: 1
num_nics: 2
include_num_access: false
monitored_traffic:
icmp:
- NONE
tcp:
- DNS
include_nmne: false
routers:
- hostname: router_1
num_ports: 0
ip_list:
- 192.168.10.21
- 192.168.10.22
- 192.168.10.23
wildcard_list:
- 0.0.0.1
port_list:
- 80
- 5432
protocol_list:
- ICMP
- TCP
- UDP
num_rules: 10
- type: links
label: LINKS
options:
link_references:
- switch_1:eth-1<->client_1:eth-1
- switch_1:eth-2<->client_2:eth-1
- type: none
label: ICS
options: {}
action_space:
action_map:
0:
action: do-nothing
options: {}
reward_function:
reward_components:
- type: database-file-integrity
weight: 0.5
options:
node_hostname: database_server
folder_name: database
file_name: database.db
- type: web-server-404-penalty
weight: 0.5
options:
node_hostname: web_server
service_name: web_server_web_service
agent_settings:
flatten_obs: true
simulation:
network:
nodes:
- type: switch
hostname: switch_1
num_ports: 8
- hostname: client_1
type: computer
ip_address: 192.168.10.21
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- type: ransomware-script
- type: web-browser
options:
target_url: http://arcd.com/users/
- type: database-client
options:
db_server_ip: 192.168.1.10
server_password: arcd
- type: data-manipulation-bot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.21
server_password: arcd
- type: dos-bot
options:
target_ip_address: 192.168.10.21
payload: SPOOF DATA
port_scan_p_of_success: 0.8
services:
- type: dns-client
options:
dns_server: 192.168.1.10
- type: dns-server
options:
domain_mapping:
arcd.com: 192.168.1.10
- type: database-service
options:
backup_server_ip: 192.168.1.10
- type: web-server
- type: ftp-server
options:
server_password: arcd
- type: ntp-client
options:
ntp_server_ip: 192.168.1.10
- type: ntp-server
- hostname: client_2
type: computer
ip_address: 192.168.10.22
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
folders:
- folder_name: empty_folder
- folder_name: downloads
files:
- file_name: "test.txt"
- file_name: "another_file.pwtwoti"
- folder_name: root
files:
- file_name: passwords
size: 663
type: TXT
# pre installed services and applications
- hostname: client_3
type: computer
ip_address: 192.168.10.23
subnet_mask: 255.255.255.0
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
start_up_duration: 0
shut_down_duration: 0
operating_state: "OFF"
# pre installed services and applications
links:
- endpoint_a_hostname: switch_1
endpoint_a_port: 1
endpoint_b_hostname: client_1
endpoint_b_port: 1
bandwidth: 200
- endpoint_a_hostname: switch_1
endpoint_a_port: 2
endpoint_b_hostname: client_2
endpoint_b_port: 1
bandwidth: 200

173
tests/e2e_integration_tests/test_uc7_agents.py Normal file
View File

@@ -0,0 +1,173 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG, load
from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon
from primaite.simulator.system.applications.red_applications.ransomware_script import RansomwareScript
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.software import SoftwareHealthState
CONFIG_FILE = _EXAMPLE_CFG / "uc7_config.yaml"
@pytest.fixture(scope="function")
def uc7_environment() -> PrimaiteGymEnv:
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
env = PrimaiteGymEnv(env_config=cfg)
return env
def assert_agent_reward(env: PrimaiteGymEnv, agent_name: str, positive: bool):
"""Asserts that a given agent has a reward that is below/above or equal to 0 dependant on arguments."""
agent_reward = env.game.agents[agent_name].reward_function.total_reward
if agent_name == "defender":
return # ignore blue agent
if positive is True:
assert agent_reward >= 0 # Asserts that no agents are below a total reward of 0
elif positive is False:
assert agent_reward <= 0 # Asserts that no agents are above a total reward of 0
else:
print("Invalid 'positive' argument.")
def test_green_agent_positive_reward(uc7_environment):
"""Confirms that the UC7 Green Agents receive a positive reward (Default Behaviour)."""
env: PrimaiteGymEnv = uc7_environment
# Performing no changes to the environment. Default Behaviour
# Stepping 60 times in the environment
for _ in range(60):
env.step(0)
for agent in env.game.agents:
assert_agent_reward(env=env, agent_name=env.game.agents[agent].config.ref, positive=True)
def test_green_agent_negative_reward(uc7_environment):
"""Confirms that the UC7 Green Agents receive a negative reward. (Disabled web-server and database-service)"""
env: PrimaiteGymEnv = uc7_environment
# Purposefully disabling the following services:
# 1. Disabling the web-server
st_dmz_pub_srv_web: Server = env.game.simulation.network.get_node_by_hostname("ST-DMZ-PUB-SRV-WEB")
st_web_server = st_dmz_pub_srv_web.software_manager.software["web-server"]
st_web_server.operating_state = ServiceOperatingState.DISABLED
assert st_web_server.operating_state == ServiceOperatingState.DISABLED
# 2. Disabling the DatabaseServer
st_data_database_server: Server = env.game.simulation.network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
database_service: DatabaseService = st_data_database_server.software_manager.software["database-service"]
database_service.operating_state = ServiceOperatingState.DISABLED
assert database_service.operating_state == ServiceOperatingState.DISABLED
# Stepping 100 times in the environment
for _ in range(100):
env.step(0)
for agent in env.game.agents:
assert_agent_reward(env=env, agent_name=env.game.agents[agent].config.ref, positive=False)
def test_tap001_default_behaviour(uc7_environment):
"""Confirms that the TAP001 expected simulation impacts works as expected in the UC7 environment."""
env: PrimaiteGymEnv = uc7_environment
env.reset()
network = env.game.simulation.network
# Running for 128 episodes
for _ in range(128):
env.step(0)
some_tech_proj_a_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-1")
# Asserting that the `malware_dropper.ps1` was created.
malware_dropper_file: File = some_tech_proj_a_pc_1.file_system.get_file("downloads", "malware_dropper.ps1")
assert malware_dropper_file.health_status == FileSystemItemHealthStatus.GOOD
# Asserting that the `RansomwareScript` launched successfully.
ransomware_script: RansomwareScript = some_tech_proj_a_pc_1.software_manager.software["ransomware-script"]
assert ransomware_script.health_state_actual == SoftwareHealthState.GOOD
assert ransomware_script.operating_state == ApplicationOperatingState.RUNNING
# Asserting that the `C2Beacon` connected to the `C2Server`.
c2_beacon: C2Beacon = some_tech_proj_a_pc_1.software_manager.software["c2-beacon"]
assert c2_beacon.health_state_actual == SoftwareHealthState.GOOD
assert c2_beacon.operating_state == ApplicationOperatingState.RUNNING
assert c2_beacon.c2_connection_active == True
# Asserting that the target database was successfully corrupted.
some_tech_data_server_database: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
database_file: File = some_tech_data_server_database.file_system.get_file(
folder_name="database", file_name="database.db"
)
assert database_file.health_status == FileSystemItemHealthStatus.CORRUPT
def test_tap003_default_behaviour(uc7_environment):
"""Confirms that the TAP003 expected simulation impacts works as expected in the UC7 environment."""
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.simulator.network.transmission.network_layer import IPPacket, IPProtocol
from primaite.utils.validation.port import PORT_LOOKUP
def uc7_environment_tap003() -> PrimaiteGymEnv:
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["agents"][32]["agent_settings"]["starting_nodes"] = ["ST-PROJ-A-PRV-PC-1"]
cfg["agents"][32]["agent_settings"]["default_starting_node"] = "ST-PROJ-A-PRV-PC-1"
env = PrimaiteGymEnv(env_config=cfg)
return env
env: PrimaiteGymEnv = uc7_environment_tap003()
env.reset()
# Running for 128 episodes
for _ in range(128):
env.step(0)
network = env.game.simulation.network
# Asserting that a malicious ACL has been added to ST-INTRA-PRV-RT-DR-1
st_intra_prv_rt_dr_1: Router = network.get_node_by_hostname(hostname="ST-INTRA-PRV-RT-DR-1")
assert st_intra_prv_rt_dr_1.acl.acl[1].action == ACLAction.DENY
assert st_intra_prv_rt_dr_1.acl.acl[1].protocol == "tcp"
assert st_intra_prv_rt_dr_1.acl.acl[1].src_port == PORT_LOOKUP.get("POSTGRES_SERVER")
assert st_intra_prv_rt_dr_1.acl.acl[1].dst_port == PORT_LOOKUP.get("POSTGRES_SERVER")
# Asserting that a malicious ACL has been added to ST-INTRA-PRV-RT-CR
st_intra_prv_rt_cr: Router = network.get_node_by_hostname(hostname="ST-INTRA-PRV-RT-CR")
assert st_intra_prv_rt_cr.acl.acl[1].action == ACLAction.DENY
assert st_intra_prv_rt_cr.acl.acl[1].protocol == "tcp"
assert st_intra_prv_rt_cr.acl.acl[1].src_port == PORT_LOOKUP.get("HTTP")
assert st_intra_prv_rt_cr.acl.acl[1].dst_port == PORT_LOOKUP.get("HTTP")
# Asserting that a malicious ACL has been added to REM-PUB-RT-DR
rem_pub_rt_dr: Router = network.get_node_by_hostname(hostname="REM-PUB-RT-DR")
assert rem_pub_rt_dr.acl.acl[1].action == ACLAction.DENY
assert rem_pub_rt_dr.acl.acl[1].protocol == "tcp"
assert rem_pub_rt_dr.acl.acl[1].src_port == PORT_LOOKUP.get("DNS")
assert rem_pub_rt_dr.acl.acl[1].dst_port == PORT_LOOKUP.get("DNS")

237
tests/e2e_integration_tests/test_uc7_route_connectivity.py Normal file
View File

@@ -0,0 +1,237 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
CONFIG_FILE = _EXAMPLE_CFG / "uc7_config.yaml"
@pytest.fixture(scope="function")
def uc7_network() -> Network:
with open(file=CONFIG_FILE, mode="r") as f:
cfg = yaml.safe_load(stream=f)
game = PrimaiteGame.from_config(cfg=cfg)
return game.simulation.network
def test_ping_home_office(uc7_network):
"""Asserts that all home_pub_* can ping each-other and the public dns (isp_pub_srv_dns)"""
network = uc7_network
home_pub_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
home_pub_pc_2: Computer = network.get_node_by_hostname("HOME-PUB-PC-2")
home_pub_pc_srv: Server = network.get_node_by_hostname("HOME-PUB-SRV")
home_pub_rt_dr: Router = network.get_node_by_hostname("HOME-PUB-RT-DR")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
assert home_pub_pc_1.ping(isp_pub_srv_dns.network_interface[1].ip_address)
def ping_all_home_office(host):
assert host.ping(home_pub_pc_1.network_interface[1].ip_address)
assert host.ping(home_pub_pc_2.network_interface[1].ip_address)
assert host.ping(home_pub_pc_srv.network_interface[1].ip_address)
assert host.ping(home_pub_rt_dr.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_home_office(home_pub_pc_1)
ping_all_home_office(home_pub_pc_2)
ping_all_home_office(home_pub_pc_srv)
ping_all_home_office(isp_pub_srv_dns)
def test_ping_remote_site(uc7_network):
"""Asserts that all remote_pub_* hosts can ping each-other and the public dns server (isp_pub_srv_dns)"""
network = uc7_network
rem_pub_fw: Firewall = network.get_node_by_hostname(hostname="REM-PUB-FW")
rem_pub_rt_dr: Router = network.get_node_by_hostname(hostname="REM-PUB-RT-DR")
rem_pub_pc_1: Computer = network.get_node_by_hostname(hostname="REM-PUB-PC-1")
rem_pub_pc_2: Computer = network.get_node_by_hostname(hostname="REM-PUB-PC-2")
rem_pub_srv: Computer = network.get_node_by_hostname(hostname="REM-PUB-SRV")
def ping_all_remote_site(host):
assert host.ping(rem_pub_fw.network_interface[1].ip_address)
assert host.ping(rem_pub_rt_dr.network_interface[1].ip_address)
assert host.ping(rem_pub_pc_1.network_interface[1].ip_address)
assert host.ping(rem_pub_pc_2.network_interface[1].ip_address)
assert host.ping(rem_pub_srv.network_interface[1].ip_address)
ping_all_remote_site(host=rem_pub_fw)
ping_all_remote_site(host=rem_pub_rt_dr)
ping_all_remote_site(host=rem_pub_pc_1)
ping_all_remote_site(host=rem_pub_pc_2)
ping_all_remote_site(host=rem_pub_srv)
def test_ping_some_tech_dmz(uc7_network):
"""Asserts that the st_dmz_pub_srv_web and the st_public_firewall can ping each other and remote site and home office."""
network = uc7_network
st_pub_fw: Firewall = network.get_node_by_hostname(hostname="ST-PUB-FW")
st_dmz_pub_srv_web: Server = network.get_node_by_hostname(hostname="ST-DMZ-PUB-SRV-WEB")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
home_pub_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
def ping_all_some_tech_dmz(host):
assert host.ping(st_dmz_pub_srv_web.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_dmz(host=st_pub_fw)
ping_all_some_tech_dmz(host=isp_pub_srv_dns)
ping_all_some_tech_dmz(host=home_pub_pc_1)
def test_ping_some_tech_head_office(uc7_network):
"""Asserts that all the some_tech_* PCs can ping each other and the public dns"""
network = uc7_network
st_home_office_private_pc_1: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-1")
st_home_office_private_pc_2: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-2")
st_home_office_private_pc_3: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-3")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_head_office(host):
assert host.ping(st_home_office_private_pc_1.network_interface[1].ip_address)
assert host.ping(st_home_office_private_pc_2.network_interface[1].ip_address)
assert host.ping(st_home_office_private_pc_3.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_head_office(host=st_home_office_private_pc_1)
ping_all_some_tech_head_office(host=st_home_office_private_pc_2)
ping_all_some_tech_head_office(host=st_home_office_private_pc_3)
def test_ping_some_tech_hr(uc7_network):
"""Assert that all some_tech_hr_* PCs can ping each other and the public dns"""
network = uc7_network
some_tech_hr_pc_1: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-1")
some_tech_hr_pc_2: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-2")
some_tech_hr_pc_3: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-3")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_hr(host):
assert host.ping(some_tech_hr_pc_1.network_interface[1].ip_address)
assert host.ping(some_tech_hr_pc_2.network_interface[1].ip_address)
assert host.ping(some_tech_hr_pc_3.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_hr(some_tech_hr_pc_1)
ping_all_some_tech_hr(some_tech_hr_pc_2)
ping_all_some_tech_hr(some_tech_hr_pc_3)
def test_some_tech_data_hr(uc7_network):
"""Assert that all some_tech_data_* servers can ping each other and the public dns."""
network = uc7_network
some_tech_data_server_storage: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-STORAGE")
some_tech_data_server_database: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_hr(host):
assert host.ping(some_tech_data_server_storage.network_interface[1].ip_address)
assert host.ping(some_tech_data_server_database.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_hr(some_tech_data_server_storage)
ping_all_some_tech_hr(some_tech_data_server_database)
def test_some_tech_project_a(uc7_network):
"""Asserts that all some_tech project A's PCs can ping each other and the public dns."""
network = uc7_network
some_tech_proj_a_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-1")
some_tech_proj_a_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-2")
some_tech_proj_a_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-3")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_proj_a(host):
assert host.ping(some_tech_proj_a_pc_1.network_interface[1].ip_address)
assert host.ping(some_tech_proj_a_pc_2.network_interface[1].ip_address)
assert host.ping(some_tech_proj_a_pc_3.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_proj_a(some_tech_proj_a_pc_1)
ping_all_some_tech_proj_a(some_tech_proj_a_pc_2)
ping_all_some_tech_proj_a(some_tech_proj_a_pc_3)
def test_some_tech_project_b(uc7_network):
"""Asserts that all some_tech_project_b PC's can ping each other and the public dps."""
network = uc7_network
some_tech_proj_b_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-1")
some_tech_proj_b_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-2")
some_tech_proj_b_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-3")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_proj_b(host):
assert host.ping(some_tech_proj_b_pc_1.network_interface[1].ip_address)
assert host.ping(some_tech_proj_b_pc_2.network_interface[1].ip_address)
assert host.ping(some_tech_proj_b_pc_3.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_proj_b(some_tech_proj_b_pc_1)
ping_all_some_tech_proj_b(some_tech_proj_b_pc_2)
ping_all_some_tech_proj_b(some_tech_proj_b_pc_3)
def test_some_tech_project_a(uc7_network):
"""Asserts that all some_tech_project_c PC's can ping each other and the public dps."""
network = uc7_network
some_tech_proj_c_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-1")
some_tech_proj_c_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-2")
some_tech_proj_c_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-3")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
def ping_all_some_tech_proj_c(host):
assert host.ping(some_tech_proj_c_pc_1.network_interface[1].ip_address)
assert host.ping(some_tech_proj_c_pc_2.network_interface[1].ip_address)
assert host.ping(some_tech_proj_c_pc_3.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
ping_all_some_tech_proj_c(some_tech_proj_c_pc_1)
ping_all_some_tech_proj_c(some_tech_proj_c_pc_2)
ping_all_some_tech_proj_c(some_tech_proj_c_pc_3)
def test_ping_all_networks(uc7_network):
"""Asserts that one machine from each network is able to ping all others."""
network = uc7_network
home_office_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
remote_office_pc_1: Computer = network.get_node_by_hostname("REM-PUB-PC-1")
st_head_office_pc_1: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-1")
st_human_resources_pc_1: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-1")
st_data_storage_server: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-STORAGE")
st_data_database_server: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
st_proj_a_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-1")
st_proj_b_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-1")
st_proj_c_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-1")
def ping_network_wide(host):
assert host.ping(home_office_pc_1.network_interface[1].ip_address)
assert host.ping(isp_pub_srv_dns.network_interface[1].ip_address)
assert host.ping(remote_office_pc_1.network_interface[1].ip_address)
assert host.ping(st_head_office_pc_1.network_interface[1].ip_address)
assert host.ping(st_human_resources_pc_1.network_interface[1].ip_address)
assert host.ping(st_data_storage_server.network_interface[1].ip_address)
assert host.ping(st_data_database_server.network_interface[1].ip_address)
assert host.ping(st_proj_a_pc_1.network_interface[1].ip_address)
assert host.ping(st_proj_b_pc_1.network_interface[1].ip_address)
assert host.ping(st_proj_c_pc_1.network_interface[1].ip_address)
ping_network_wide(host=home_office_pc_1)
ping_network_wide(host=isp_pub_srv_dns)
ping_network_wide(host=remote_office_pc_1)
ping_network_wide(host=st_head_office_pc_1)
ping_network_wide(host=st_human_resources_pc_1)
ping_network_wide(host=st_data_storage_server)
ping_network_wide(host=st_data_database_server)
ping_network_wide(host=st_proj_a_pc_1)
ping_network_wide(host=st_proj_b_pc_1)
ping_network_wide(host=st_proj_c_pc_1)

338
tests/e2e_integration_tests/test_uc7_services_and_applications.py Normal file
View File

@@ -0,0 +1,338 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.game import PrimaiteGame
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.host.computer import Computer
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.system.applications.application import ApplicationOperatingState
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.applications.web_browser import WebBrowser
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.dns_server import DNSServer
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
from primaite.simulator.system.services.ftp.ftp_server import FTPServer
from primaite.simulator.system.services.ntp.ntp_client import NTPClient
from primaite.simulator.system.services.ntp.ntp_server import NTPServer
from primaite.simulator.system.services.service import ServiceOperatingState
from primaite.simulator.system.software import SoftwareHealthState
CONFIG_FILE = _EXAMPLE_CFG / "uc7_config.yaml"
@pytest.fixture(scope="function")
def uc7_network() -> Network:
with open(file=CONFIG_FILE, mode="r") as f:
cfg = yaml.safe_load(stream=f)
game = PrimaiteGame.from_config(cfg=cfg)
return game.simulation.network
def assert_ntp_client(host):
"""Confirms that the ntp_client service is present and functioning."""
ntp_client: NTPClient = host.software_manager.software["ntp-client"]
assert ntp_client is not None
assert ntp_client.operating_state == ServiceOperatingState.RUNNING
assert ntp_client.health_state_actual == SoftwareHealthState.GOOD
def assert_dns_client(host):
"""Confirms that the dns_client service is present and functioning."""
dns_client: DNSClient = host.software_manager.software["dns-client"]
assert dns_client is not None
assert dns_client.operating_state == ServiceOperatingState.RUNNING
assert dns_client.health_state_actual == SoftwareHealthState.GOOD
def assert_web_browser(host: Computer):
"""Asserts that the web_browser application is present and functioning."""
web_browser: WebBrowser = host.software_manager.software["web-browser"]
assert web_browser is not None
assert web_browser.operating_state == ApplicationOperatingState.RUNNING
assert web_browser.health_state_actual == SoftwareHealthState.GOOD
def assert_database_client(host: Computer):
"""Asserts that the database_client application is present and functioning."""
database_client = host.software_manager.software["database-client"]
assert database_client is not None
assert database_client.operating_state == ApplicationOperatingState.RUNNING
assert database_client.health_state_actual == SoftwareHealthState.GOOD
def test_home_office_software(uc7_network):
"""Asserts that each host in the home_office network contains the expected software."""
network: Network = uc7_network
home_pub_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
home_pub_pc_2: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
home_pub_srv: Server = network.get_node_by_hostname("HOME-PUB-SRV")
# Home Office PC 1
assert_web_browser(home_pub_pc_1)
assert_database_client(home_pub_pc_1)
assert_dns_client(home_pub_pc_1)
assert_ntp_client(home_pub_pc_1)
# Home Office PC 2
assert_web_browser(home_pub_pc_2)
assert_database_client(home_pub_pc_2)
assert_dns_client(home_pub_pc_2)
assert_ntp_client(home_pub_pc_2)
# Home Office Server
assert_dns_client(home_pub_srv)
assert_ntp_client(home_pub_srv)
def test_internet_dns_server(uc7_network):
"""Asserts that `ISP-PUB-SRV-DNS` host's DNSServer application is operating and functioning as expected."""
network: Network = uc7_network
isp_pub_srv_dns: Server = network.get_node_by_hostname("ISP-PUB-SRV-DNS")
# Confirming that the DNSServer is up and running:
dns_server: DNSServer = isp_pub_srv_dns.software_manager.software["dns-server"]
assert dns_server is not None
assert dns_server.operating_state == ServiceOperatingState.RUNNING
assert dns_server.health_state_actual == SoftwareHealthState.GOOD
# Confirming that the DNSServer is performing as expected by performing a request from a client
home_pub_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
dns_client: DNSClient = home_pub_pc_1.software_manager.software["dns-client"]
assert dns_client.check_domain_exists(target_domain="some_tech.com")
assert dns_client.dns_cache.get("some_tech.com", None) is not None
assert len(dns_client.dns_cache) == 1
def test_remote_office_software(uc7_network):
"""Asserts that each host on the remote_office network has the expected services & applications which are operating as expected."""
network = uc7_network
rem_pub_pc_1: Computer = network.get_node_by_hostname(hostname="REM-PUB-PC-1")
rem_pub_pc_2: Computer = network.get_node_by_hostname(hostname="REM-PUB-PC-2")
rem_pub_srv: Server = network.get_node_by_hostname(hostname="REM-PUB-SRV")
# Remote Site PC 1
assert_web_browser(rem_pub_pc_1)
assert_database_client(rem_pub_pc_1)
assert_dns_client(rem_pub_pc_1)
assert_ntp_client(rem_pub_pc_1)
# Remote Site PC 2
assert_web_browser(rem_pub_pc_2)
assert_database_client(rem_pub_pc_2)
assert_dns_client(rem_pub_pc_2)
assert_ntp_client(rem_pub_pc_2)
# Remote Site Server
assert_dns_client(rem_pub_srv)
assert_ntp_client(rem_pub_srv)
def test_dmz_web_server(uc7_network):
"""Asserts that the DMZ WebServer functions as expected"""
network: Network = uc7_network
st_dmz_pub_srv_web: Server = network.get_node_by_hostname("ST-DMZ-PUB-SRV-WEB")
# Asserting the ST Web Server is working as expected
st_web_server = st_dmz_pub_srv_web.software_manager.software["web-server"]
assert st_web_server is not None
assert st_web_server.operating_state == ServiceOperatingState.RUNNING
assert st_web_server.health_state_actual == SoftwareHealthState.GOOD
# Asserting that WebBrowser can actually connect to the WebServer
# SOME TECH Human Resources --> DMZ Web Server
st_hr_pc_1: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-1")
st_hr_pc_1_web_browser: WebBrowser = st_hr_pc_1.software_manager.software["web-browser"]
assert st_hr_pc_1_web_browser.get_webpage("http://some_tech.com")
# Remote Site --> DMZ Web Server
rem_pub_pc_1: Computer = network.get_node_by_hostname("REM-PUB-PC-1")
rem_pub_pc_1_web_browser: WebBrowser = rem_pub_pc_1.software_manager.software["web-browser"]
assert rem_pub_pc_1_web_browser.get_webpage("http://some_tech.com")
# Home Office --> DMZ Web Server
home_pub_pc_1: Computer = network.get_node_by_hostname("HOME-PUB-PC-1")
home_pub_pc_1_web_browser: WebBrowser = home_pub_pc_1.software_manager.software["web-browser"]
assert home_pub_pc_1_web_browser.get_webpage("http://some_tech.com")
def test_tech_head_office_software(uc7_network):
"""Asserts that each host on the some_tech_head_office network has the expected services & applications which are operating as expected."""
network: Network = uc7_network
st_head_office_private_pc_1: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-1")
st_head_office_private_pc_2: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-2")
st_head_office_private_pc_3: Computer = network.get_node_by_hostname("ST-HO-PRV-PC-3")
# ST Head Office One
assert_web_browser(st_head_office_private_pc_1)
assert_database_client(st_head_office_private_pc_1)
assert_dns_client(st_head_office_private_pc_1)
assert_ntp_client(st_head_office_private_pc_1)
# ST Head Office Two
assert_web_browser(st_head_office_private_pc_2)
assert_database_client(st_head_office_private_pc_2)
assert_dns_client(st_head_office_private_pc_2)
assert_ntp_client(st_head_office_private_pc_2)
# ST Head Office Three
assert_web_browser(st_head_office_private_pc_3)
assert_database_client(st_head_office_private_pc_3)
assert_dns_client(st_head_office_private_pc_3)
assert_ntp_client(st_head_office_private_pc_3)
def test_tech_human_resources_office_software(uc7_network):
"""Asserts that each host on the some_tech human_resources network has the expected services & applications which are operating as expected."""
network: Network = uc7_network
st_hr_pc_1: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-1")
st_hr_pc_2: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-2")
st_hr_pc_3: Computer = network.get_node_by_hostname("ST-HR-PRV-PC-3")
# ST Human Resource PC 1
assert_web_browser(st_hr_pc_1)
assert_database_client(st_hr_pc_1)
assert_dns_client(st_hr_pc_1)
assert_ntp_client(st_hr_pc_1)
# ST Human Resource PC 2
assert_web_browser(st_hr_pc_2)
assert_database_client(st_hr_pc_2)
assert_dns_client(st_hr_pc_2)
assert_ntp_client(st_hr_pc_2)
# ST Human Resource PC 3
assert_web_browser(st_hr_pc_3)
assert_database_client(st_hr_pc_3)
assert_dns_client(st_hr_pc_3)
assert_ntp_client(st_hr_pc_3)
def test_tech_data_software(uc7_network):
"""Asserts the database and database storage servers on the some_tech data network are operating as expected."""
network: Network = uc7_network
st_data_database_server: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
st_data_database_storage: Server = network.get_node_by_hostname("ST-DATA-PRV-SRV-STORAGE")
st_proj_a_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-1")
# Asserting that the database_service is working as expected
database_service: DatabaseService = st_data_database_server.software_manager.software["database-service"]
assert database_service is not None
assert database_service.operating_state == ServiceOperatingState.RUNNING
assert database_service.health_state_actual == SoftwareHealthState.GOOD
# Asserting that the database_client can connect to the database
database_client: DatabaseClient = st_proj_a_pc_1.software_manager.software["database-client"]
assert database_client.server_ip_address is not None
assert database_client.server_ip_address == st_data_database_server.network_interface[1].ip_address
assert database_client.connect()
# Asserting that the database storage works as expected.
assert database_service.backup_server_ip == st_data_database_storage.network_interface[1].ip_address
assert database_service.backup_database()
def test_tech_proj_a_software(uc7_network):
"""Asserts that each host on the some_tech project A network has the expected services & applications which are operating as expected."""
network: Network = uc7_network
st_proj_a_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-1")
st_proj_a_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-2")
st_proj_a_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-A-PRV-PC-3")
# ST Project A - PC 1
assert_web_browser(st_proj_a_pc_1)
assert_database_client(st_proj_a_pc_1)
assert_dns_client(st_proj_a_pc_1)
assert_ntp_client(st_proj_a_pc_1)
# ST Project A - PC 2
assert_web_browser(st_proj_a_pc_2)
assert_database_client(st_proj_a_pc_2)
assert_dns_client(st_proj_a_pc_2)
assert_ntp_client(st_proj_a_pc_2)
# ST Project A - PC 3
assert_web_browser(st_proj_a_pc_3)
assert_database_client(st_proj_a_pc_3)
assert_dns_client(st_proj_a_pc_3)
assert_ntp_client(st_proj_a_pc_3)
def test_tech_proj_b_software(uc7_network):
"""Asserts that each host on the some_tech project A network has the expected services & applications which are operating as expected."""
network: Network = uc7_network
st_proj_b_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-1")
st_proj_b_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-2")
st_proj_b_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-B-PRV-PC-3")
# ST Project B - PC 1
assert_web_browser(st_proj_b_pc_1)
assert_database_client(st_proj_b_pc_1)
assert_dns_client(st_proj_b_pc_1)
assert_ntp_client(st_proj_b_pc_1)
# ST Project B - PC2
assert_web_browser(st_proj_b_pc_2)
assert_database_client(st_proj_b_pc_2)
assert_dns_client(st_proj_b_pc_2)
assert_ntp_client(st_proj_b_pc_2)
# ST Project B - PC3
assert_web_browser(st_proj_b_pc_3)
assert_database_client(st_proj_b_pc_3)
assert_dns_client(st_proj_b_pc_3)
assert_ntp_client(st_proj_b_pc_3)
def test_tech_proj_c_software(uc7_network):
"""Asserts that each host on the some_tech project A network has the expected services & applications which are operating as expected."""
network: Network = uc7_network
st_proj_c_pc_1: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-1")
st_proj_c_pc_2: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-2")
st_proj_c_pc_3: Computer = network.get_node_by_hostname("ST-PROJ-C-PRV-PC-3")
# ST Project C - PC 1
assert_web_browser(st_proj_c_pc_1)
assert_database_client(st_proj_c_pc_1)
assert_dns_client(st_proj_c_pc_1)
assert_ntp_client(st_proj_c_pc_1)
# ST Project C - PC2
assert_web_browser(st_proj_c_pc_2)
assert_database_client(st_proj_c_pc_2)
assert_dns_client(st_proj_c_pc_2)
assert_ntp_client(st_proj_c_pc_2)
# ST Project C - PC3
assert_web_browser(st_proj_c_pc_3)
assert_database_client(st_proj_c_pc_3)
assert_dns_client(st_proj_c_pc_3)
assert_ntp_client(st_proj_c_pc_3)

1
tests/e2e_integration_tests/threat_actor_profiles/__init__.py Normal file
View File

@@ -0,0 +1 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK

143
tests/e2e_integration_tests/threat_actor_profiles/test_abstract_tap.py Normal file
View File

@@ -0,0 +1,143 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
def uc7_tap001_env() -> PrimaiteGymEnv:
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
for a in cfg["agents"]:
if a["ref"] == "attacker":
tap_cfg = a
tap_cfg["agent_settings"]["start_step"] = 1
tap_cfg["agent_settings"]["frequency"] = 5
tap_cfg["agent_settings"]["variance"] = 0
env = PrimaiteGymEnv(env_config=cfg)
return env
def uc7_tap003_env(**kwargs) -> PrimaiteGymEnv:
"""Setups the UC7 TAP003 Game with the following settings:
start_step = Start on Step 1
frequency = Attack Every 5 Steps
Each PyTest will define the rest of the TAP & Kill Chain settings via **Kwargs
"""
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", "r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
for a in cfg["agents"]:
if a["ref"] == "attacker":
tap_cfg = a
tap_cfg["agent_settings"]["start_step"] = 1
tap_cfg["agent_settings"]["frequency"] = 5
tap_cfg["agent_settings"]["variance"] = 0
if "repeat_kill_chain" in kwargs:
tap_cfg["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
if "repeat_kill_chain_stages" in kwargs:
tap_cfg["agent_settings"]["repeat_kill_chain_stages"] = kwargs["repeat_kill_chain_stages"]
if "planning_probability" in kwargs:
tap_cfg["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs["planning_probability"]
if "custom_kill_chain" in kwargs:
tap_cfg["agent_settings"]["kill_chain"] = kwargs["custom_kill_chain"]
if "starting_nodes" in kwargs:
tap_cfg["agent_settings"]["starting_nodes"] = kwargs["starting_nodes"]
if "target_nodes" in kwargs:
tap_cfg["agent_settings"]["target_nodes"] = kwargs["target_nodes"]
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap001_setup():
"""Tests abstract TAP's following methods:
1. _setup_kill_chain
2. _setup_agent_kill_chain
3. _setup_tap_applications
"""
env = uc7_tap001_env() # Using TAP001 for PyTests.
tap: TAP001 = env.game.agents["attacker"]
# check the kill chain loaded correctly
assert tap.selected_kill_chain is MobileMalwareKillChain
assert tap.selected_kill_chain.FAILED == BaseKillChain.FAILED
assert tap.selected_kill_chain.SUCCEEDED == BaseKillChain.SUCCEEDED
assert tap.selected_kill_chain.NOT_STARTED == BaseKillChain.NOT_STARTED
if sn := tap.config.agent_settings.default_starting_node:
assert tap.starting_node == sn
else:
assert tap.starting_node in tap.config.agent_settings.starting_nodes
if ti := tap.config.agent_settings.default_target_ip:
assert tap.target_ip == ti
else:
assert tap.target_ip in tap.config.agent_settings.target_ips
assert tap.next_execution_timestep == tap.config.agent_settings.start_step
assert tap.current_host == tap.starting_node
def test_abstract_tap_select_start_node():
"""Tests that Abstract TAP's _select_start_node"""
env = uc7_tap003_env(repeat_kill_chain=True, repeat_kill_chain_stages=True) # Using TAP003 for PyTests.
tap: TAP003 = env.game.agents["attacker"]
assert tap.starting_node == "ST-PROJ-A-PRV-PC-1"
assert tap.current_host == tap.starting_node
def test_outcome_handler():
"""Tests Abstract Tap's outcome handler concludes the episode when the kill chain fails."""
env = uc7_tap003_env(repeat_kill_chain=False, repeat_kill_chain_stages=False) # Using TAP003 for PyTests.
tap: TAP003 = env.game.agents["attacker"]
tap.current_kill_chain_stage = BaseKillChain.FAILED
for _ in range(6):
env.step(0)
assert tap.actions_concluded == True
def test_abstract_tap_kill_chain_repeat():
"""Tests that the kill chain repeats from the beginning upon failure."""
env = uc7_tap003_env(repeat_kill_chain=True, repeat_kill_chain_stages=False) # Using TAP003 for PyTests.
tap: TAP003 = env.game.agents["attacker"]
for _ in range(15):
env.step(0) # Steps directly to the Access Stage
assert tap.current_kill_chain_stage == InsiderKillChain.ACCESS
tap.current_kill_chain_stage = BaseKillChain.FAILED
for _ in range(5):
env.step(0) # Steps to manipulation - but failure causes the kill chain to restart.
assert tap.actions_concluded == False
assert tap.current_kill_chain_stage == InsiderKillChain.RECONNAISSANCE
"""Tests that kill chain stages repeat when expected"""
env = uc7_tap003_env(
repeat_kill_chain=True, repeat_kill_chain_stages=True, planning_probability=0
) # Using TAP003 for PyTests.
tap: TAP003 = env.game.agents["attacker"]
tap.current_kill_chain_stage = InsiderKillChain.PLANNING
for _ in range(15):
env.step(0) # Attempts to progress past the PLANNING stage multiple times.
assert tap.actions_concluded == False
assert tap.current_kill_chain_stage == InsiderKillChain.PLANNING

69
tests/e2e_integration_tests/threat_actor_profiles/test_kill_chain_methods.py Normal file
View File

@@ -0,0 +1,69 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
def uc7_tap003_env() -> PrimaiteGymEnv:
"""Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
env = PrimaiteGymEnv(env_config=cfg)
return env
def uc7_tap001_env() -> PrimaiteGymEnv:
"""Setup the UC7 TAP001 Game with the start_step & frequency set to 1 & 2 respectively. Probabilities are set to 1"""
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap003_insider_kill_chain_load():
"""Tests that tap003's insider kill chain is successfully loaded into the tap.selected_kill_chain attribute."""
env = uc7_tap003_env() # Using TAP003 for PyTests.
tap: TAP003 = env.game.agents["attacker"]
# Asserting that the Base Kill Chain intEnum stages are loaded
assert BaseKillChain.FAILED in [enums for enums in tap.selected_kill_chain]
assert BaseKillChain.SUCCEEDED in [enums for enums in tap.selected_kill_chain]
assert BaseKillChain.NOT_STARTED in [enums for enums in tap.selected_kill_chain]
# Asserting that the Insider Kill Chain intenum stages are loaded.
assert len(InsiderKillChain.__members__) == len(tap.selected_kill_chain.__members__)
def test_tap001_mobile_malware_kill_chain_load():
"""Tests that tap001's mobile malware is successfully loaded into the tap.selected_kill_chain attribute."""
env = uc7_tap001_env() # Using TAP003 for PyTests.
tap: TAP001 = env.game.agents["attacker"]
# Asserting that the Base Kill Chain intEnum stages are loaded.
assert BaseKillChain.FAILED in [enums for enums in tap.selected_kill_chain]
assert BaseKillChain.SUCCEEDED in [enums for enums in tap.selected_kill_chain]
assert BaseKillChain.NOT_STARTED in [enums for enums in tap.selected_kill_chain]
# Asserting that the Insider Kill Chain intEnum stages are loaded.
assert len(MobileMalwareKillChain.__members__) == len(tap.selected_kill_chain.__members__)

109
tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_repeat.py Normal file
View File

@@ -0,0 +1,109 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
"""Setups the UC7 tap001 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = kwargs["repeat_kill_chain_stages"]
cfg["agents"][32]["agent_settings"]["kill_chain"]["PROPAGATE"]["probability"] = kwargs["propagate_probability"]
cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["probability"] = kwargs["payload_probability"]
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap001_repeating_kill_chain():
"""Tests to check that tap001 repeats it's kill chain after success"""
env = uc7_tap001_env(
repeat_kill_chain=True,
repeat_kill_chain_stages=True,
payload_probability=1,
propagate_probability=1,
)
tap001: TAP001 = env.game.agents["attacker"]
# Looping for 50 timesteps - As the agent is set to execute an action every 2 timesteps
# This is the equivalent of the agent taking 20 actions.
for _ in range(50): # This for loop should never actually fully complete.
if tap001.current_kill_chain_stage == BaseKillChain.SUCCEEDED:
break
env.step(0)
# Catches if the above for loop fully completes.
# This test uses a probability of 1 for all stages and a variance of 2 timesteps
# Thus the for loop above should never fail.
# If this occurs then there is an error somewhere in either:
# 1. The TAP Logic
# 2. Failing Agent Actions are causing the TAP to fail. (See tap_return_handler).
if tap001.current_kill_chain_stage != BaseKillChain.SUCCEEDED:
pytest.fail("Attacker Never Reached SUCCEEDED - Please evaluate current TAP Logic.")
# Stepping twice for the succeeded logic to kick in:
env.step(0)
env.step(0)
env.step(0)
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.DOWNLOAD.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.INSTALL.name
def test_tap001_repeating_kill_chain_stages():
"""Tests to check that tap001 repeats it's kill chain after failing a kill chain stage."""
env = uc7_tap001_env(
repeat_kill_chain=True,
repeat_kill_chain_stages=True,
payload_probability=1,
propagate_probability=0,
# Probability 0 = Will never be able to perform the access stage and progress to Manipulation.
)
tap001: TAP001 = env.game.agents["attacker"]
env.step(0) # Skipping not started
env.step(0) # Successful on the first stage
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.DOWNLOAD.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.INSTALL.name
env.step(0) # Successful progression to the second stage
env.step(0)
env.step(0)
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.INSTALL.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.ACTIVATE.name
env.step(0) # Successful progression to the third stage
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.ACTIVATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
env.step(0) # Successful progression to the Fourth stage
env.step(0)
env.step(0)
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
env.step(0) # FAILURE -- Unsuccessful progression to the Fourth stage
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
assert tap001.current_stage_progress == KillChainStageProgress.PENDING

215
tests/e2e_integration_tests/threat_actor_profiles/test_tap001_kill_chain_stages.py Normal file
View File

@@ -0,0 +1,215 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon
from primaite.simulator.system.services.database.database_service import DatabaseService
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
REPEAT_KILL_CHAIN = False # Should the TAP repeat the kill chain after success/failure?
REPEAT_KILL_CHAIN_STAGES = False # Should the TAP restart from it's previous stage on failure?
KILL_CHAIN_PROBABILITY = 1 # Blank probability for agent 'success'
DATA_EXFIL = True # Data exfiltration on the payload stage is enabled.
def uc7_tap001_env() -> PrimaiteGymEnv:
"""Setups the UC7 tap001 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["probability"] = KILL_CHAIN_PROBABILITY
cfg["agents"][32]["agent_settings"]["kill_chain"]["PROPAGATE"]["probability"] = KILL_CHAIN_PROBABILITY
cfg["agents"][32]["agent_settings"]["kill_chain"]["PAYLOAD"]["exfiltrate"] = DATA_EXFIL
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap001_kill_chain_stage_DOWNLOAD():
"""Tests that the DOWNLOAD Mobile Malware step works as expected and the expected impacts are made in the simulation."""
# Instantiating the relevant simulation/game objects:
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
starting_host = env.game.simulation.network.get_node_by_hostname(tap001.starting_node)
assert tap001.current_kill_chain_stage == BaseKillChain.NOT_STARTED
# Frequency is set to two steps
env.step(0)
env.step(0)
env.step(0)
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.DOWNLOAD.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.INSTALL.name
assert tap001.current_stage_progress == KillChainStageProgress.IN_PROGRESS
# Creating the "downloads" folder
env.step(0)
env.step(0)
assert starting_host.software_manager.file_system.get_folder(folder_name="downloads")
assert starting_host.software_manager.file_system.get_file(folder_name="downloads", file_name="malware_dropper.ps1")
# Testing that the num_file_increase works
assert starting_host.file_system.num_file_creations == 1
def test_tap001_kill_chain_stage_INSTALL():
"""Tests that the INSTALL Mobile Malware step works as expected and the expected impacts are made in the simulation."""
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
# The tap001's Starting Client:
starting_host = env.game.simulation.network.get_node_by_hostname(tap001.starting_node)
# Skipping directly to the activate stage
for _ in range(6):
env.step(0)
# Testing that tap001 Enters into the expected kill chain stages
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.INSTALL.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.ACTIVATE.name
env.step(0) # Allows the agent action to resolve.
env.step(0)
ransomware_dropper_file = starting_host.software_manager.file_system.get_file(
folder_name="downloads", file_name="malware_dropper.ps1"
)
assert ransomware_dropper_file.num_access == 1
def test_tap001_kill_chain_stage_ACTIVATE():
"""Tests that the ACTIVATE Mobile Malware step works as expected and the current impacts are made in the simulation."""
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
# The tap001's Starting Client:
starting_host = env.game.simulation.network.get_node_by_hostname(tap001.starting_node)
# Skipping directly to the activate stage
for _ in range(8):
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.ACTIVATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
# Installing ransomware-script Application
env.step(0)
env.step(0)
# Installing NMAP Application
env.step(0)
env.step(0)
# These asserts will fail if the applications are not present in the software_manager
assert starting_host.software_manager.software["ransomware-script"]
assert starting_host.software_manager.software["nmap"]
def test_tap001_kill_chain_stage_PROPAGATE():
"""Tests that the ACTIVATE Mobile Malware step works as expected and the current impacts are made in the simulation."""
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
for _ in range(12):
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
# Specific Stage by Stage Propagate Testing is done in test_tap001_propagate.
fail_safe_var = 0
while tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE:
env.step(0)
assert tap001.current_stage_progress == KillChainStageProgress.IN_PROGRESS
fail_safe_var += 1
if fail_safe_var == 100:
pytest.fail("Fail Safe Variable was hit! -- Propagate step is running indefinitely")
def test_tap001_kill_chain_stage_COMMAND_AND_CONTROL():
"""Tests that the Command And Control Mobile Malware step works as expected and the current impacts are made in the simulation."""
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
fail_safe_var = 0
for _ in range(28):
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.PAYLOAD.name
while tap001.current_kill_chain_stage == MobileMalwareKillChain.COMMAND_AND_CONTROL:
env.step(0)
fail_safe_var += 1
env.game.simulation.network.airspace.show()
if fail_safe_var == 100:
pytest.fail(reason="Fail Safe Variable was hit! -- Propagate step is running indefinitely")
starting_host = env.game.simulation.network.get_node_by_hostname(tap001.starting_node)
c2_beacon: C2Beacon = starting_host.software_manager.software["c2-beacon"]
assert c2_beacon.c2_connection_active is True
def test_tap001_kill_chain_stage_PAYLOAD():
"""Tests that the PAYLOAD Mobile Malware step works as expected and the current impacts are made in the simulation."""
env = uc7_tap001_env()
tap001: TAP001 = env.game.agents["attacker"]
# The tap001's Target Database
target_host = env.game.simulation.network.get_node_by_hostname("ST-DATA-PRV-SRV-DB")
db_server_service: DatabaseService = target_host.software_manager.software.get("database-service")
# Green agent status requests are tested within the ransomware application tests.
# See test_ransomware_disrupts_green_agent_connection for further reference.
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.GOOD
fail_safe_var = 0
while tap001.current_kill_chain_stage != MobileMalwareKillChain.PAYLOAD:
env.step(0)
fail_safe_var += 1
if fail_safe_var == 100:
pytest.fail(reason="Fail Safe Variable was hit! -- a step is running indefinitely")
for _ in range(12):
env.step(0)
assert db_server_service.db_file.health_status is FileSystemItemHealthStatus.CORRUPT
# Asserting we've managed to the database.db file onto the starting node & server
starting_host = env.game.simulation.network.get_node_by_hostname(tap001.starting_node)
c2_host = env.game.simulation.network.get_node_by_hostname(tap001.c2_settings["c2_server"])
assert starting_host.file_system.access_file(folder_name="exfiltration_folder", file_name="database.db")
assert c2_host.file_system.access_file(folder_name="exfiltration_folder", file_name="database.db")

140
tests/e2e_integration_tests/threat_actor_profiles/test_tap001_propagate_stage.py Normal file
View File

@@ -0,0 +1,140 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.system.applications.red_applications.c2.c2_beacon import C2Beacon
from primaite.simulator.system.services.database.database_service import DatabaseService
# Defining generic tap constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
REPEAT_KILL_CHAIN = False # Should the TAP repeat the kill chain after success/failure?
REPEAT_KILL_CHAIN_STAGES = False # Should the TAP restart from it's previous stage on failure?
KILL_CHAIN_PROBABILITY = 1 # Blank probability for agent 'success'
def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
"""Setups the UC7 tap001 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
agent_cfg = cfg["agents"][32]["agent_settings"]
agent_cfg["start_step"] = START_STEP
agent_cfg["frequency"] = FREQUENCY
agent_cfg["variance"] = VARIANCE
agent_cfg["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
agent_cfg["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
agent_cfg["kill_chain"]["PAYLOAD"]["probability"] = KILL_CHAIN_PROBABILITY
agent_cfg["kill_chain"]["PROPAGATE"]["probability"] = KILL_CHAIN_PROBABILITY
agent_cfg["kill_chain"]["PROPAGATE"]["scan_attempts"] = kwargs["scan_attempts"]
agent_cfg["kill_chain"]["PAYLOAD"]["payload"] = kwargs["payload"]
agent_cfg["kill_chain"]["PROPAGATE"]["network_addresses"] = kwargs["network_addresses"]
if "repeat_scan" in kwargs:
agent_cfg["kill_chain"]["PROPAGATE"]["repeat_scan"] = kwargs["repeat_scan"]
if "starting_nodes" in kwargs:
agent_cfg["starting_nodes"] = kwargs["starting_nodes"]
agent_cfg["default_starting_node"] = kwargs["starting_nodes"][0]
if "target_ips" in kwargs:
agent_cfg["target_ips"] = kwargs["target_ips"]
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap001_kill_chain_stage_PROPAGATE_default():
"""Tests that the PROPAGATE Mobile Malware step works as expected and the current impacts are made in the simulation."""
payload = "ENCRYPT"
scan_attempts = 10
network_addresses = [
"192.168.230.0/29",
"192.168.10.0/26",
"192.168.20.0/30",
"192.168.240.0/29",
"192.168.220.0/29",
]
env = uc7_tap001_env(payload=payload, scan_attempts=scan_attempts, network_addresses=network_addresses)
tap001: TAP001 = env.game.agents["attacker"]
# First Kill Chain Stages
for _ in range(12):
env.step(0)
# Assert that we're about to enter into the propagate stage.
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
# Move into the propagate stage.
while tap001.current_kill_chain_stage == MobileMalwareKillChain.PROPAGATE:
env.step(0)
# Assert that we've successfully moved into the command and control stage.
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.PAYLOAD.name
def test_tap001_kill_chain_stage_PROPAGATE_different_starting_node():
"""Tests that the PROPAGATE Mobile Malware step works as expected and the current impacts are made in the simulation from a different starting node."""
payload = "ENCRYPT"
scan_attempts = 10
network_addresses = [
"192.168.230.0/29",
"192.168.10.0/26",
"192.168.20.0/30",
"192.168.240.0/29",
"192.168.220.0/29",
]
starting_nodes = ["ST-PROJ-B-PRV-PC-2", "ST-PROJ-C-PRV-PC-3"]
env = uc7_tap001_env(
payload=payload, scan_attempts=scan_attempts, network_addresses=network_addresses, starting_nodes=starting_nodes
)
tap001: TAP001 = env.game.agents["attacker"]
for _ in range(12):
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.PROPAGATE.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
# Specific Stage by Stage Propagate Testing is done in test_tap001_propagate.
while tap001.current_kill_chain_stage == MobileMalwareKillChain.PROPAGATE:
env.step(0)
assert tap001.current_kill_chain_stage.name == MobileMalwareKillChain.COMMAND_AND_CONTROL.name
assert tap001.next_kill_chain_stage.name == MobileMalwareKillChain.PAYLOAD.name
def test_tap001_kill_chain_stage_PROPAGATE_repeat_scan():
"""Tests that the PROPAGATE Mobile Malware step will fail when the target is unable to be located."""
payload = "ENCRYPT"
scan_attempts = 20
repeat_scan = True
network_addresses = ["192.168.1.0/24", "192.168.0.0/28", "100.64.0.0/30", "172.168.0.0/28"]
env = uc7_tap001_env(
payload=payload, scan_attempts=scan_attempts, network_addresses=network_addresses, repeat_scan=repeat_scan
)
for _ in range(12):
env.step(0)
tap001: TAP001 = env.game.agents["attacker"]
while tap001.current_kill_chain_stage == MobileMalwareKillChain.PROPAGATE:
env.step(0)
# As the given network_address does not contain the target, we should failed because the maximum amount of scan attempts has been reached
assert tap001.scans_complete == 20
assert tap001.current_kill_chain_stage == MobileMalwareKillChain.FAILED

101
tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_repeat.py Normal file
View File

@@ -0,0 +1,101 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
def uc7_tap003_env(**kwargs) -> PrimaiteGymEnv:
"""Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = kwargs["repeat_kill_chain_stages"]
cfg["agents"][32]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = kwargs[
"manipulation_probability"
]
cfg["agents"][32]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = kwargs["access_probability"]
cfg["agents"][32]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs["planning_probability"]
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap003_repeating_kill_chain():
"""Tests to check that TAP003 repeats it's kill chain after success"""
env = uc7_tap003_env(
repeat_kill_chain=True,
repeat_kill_chain_stages=True,
manipulation_probability=1,
access_probability=1,
planning_probability=1,
)
tap003: TAP003 = env.game.agents["attacker"]
for _ in range(40): # This for loop should never actually fully complete.
if tap003.current_kill_chain_stage == BaseKillChain.SUCCEEDED:
break
env.step(0)
# Catches if the above for loop fully completes.
# This test uses a probability of 1 for all stages and a variance of 2 timesteps
# Thus the for loop above should never fail.
# If this occurs then there is an error somewhere in either:
# 1. The TAP Logic
# 2. Failing Agent Actions are causing the TAP to fail. (See tap_return_handler).
if tap003.current_kill_chain_stage != BaseKillChain.SUCCEEDED:
pytest.fail("Attacker Never Reached SUCCEEDED - Please evaluate current TAP Logic.")
# Stepping twice for the succeeded logic to kick in:
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
def test_tap003_repeating_kill_chain_stages():
"""Tests to check that TAP003 repeats it's kill chain after failing a kill chain stage."""
env = uc7_tap003_env(
repeat_kill_chain=True,
repeat_kill_chain_stages=True,
manipulation_probability=1,
# Probability 0 = Will never be able to perform the access stage and progress to Manipulation.
access_probability=0,
planning_probability=1,
)
tap003: TAP003 = env.game.agents["attacker"]
env.step(0) # Skipping not started
env.step(0) # Successful on the first stage
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
env.step(0) # Successful progression to the second stage
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.PLANNING.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.ACCESS.name
env.step(0) # Successfully moved onto access.
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.ACCESS.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
env.step(0) # Failure to progress past the third stage.
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.ACCESS.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name

232
tests/e2e_integration_tests/threat_actor_profiles/test_tap003_kill_chain_stages.py Normal file
View File

@@ -0,0 +1,232 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.hardware.nodes.network.router import ACLAction
# Defining constants.
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
REPEAT_KILL_CHAIN = False # Should the TAP repeat the kill chain after success/failure?
REPEAT_KILL_CHAIN_STAGES = False # Should the TAP restart from it's previous stage on failure?
KILL_CHAIN_PROBABILITY = 1 # Blank probability for agent 'success'
def uc7_tap003_env() -> PrimaiteGymEnv:
"""Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][32]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][32]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][32]["agent_settings"]["variance"] = VARIANCE
cfg["agents"][32]["agent_settings"]["repeat_kill_chain"] = REPEAT_KILL_CHAIN_STAGES
cfg["agents"][32]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
cfg["agents"][32]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = KILL_CHAIN_PROBABILITY
cfg["agents"][32]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = KILL_CHAIN_PROBABILITY
cfg["agents"][32]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = KILL_CHAIN_PROBABILITY
cfg["agents"][32]["agent_settings"]["kill_chain"]["EXPLOIT"]["probability"] = KILL_CHAIN_PROBABILITY
env = PrimaiteGymEnv(env_config=cfg)
return env
def environment_step(i: int, env: PrimaiteGymEnv) -> PrimaiteGymEnv:
"""Carries out i (given parameter) steps in the environment.."""
for x in range(i):
env.step(0)
return env
def test_tap003_kill_chain_stage_reconnaissance():
"""Tests the successful/failed handlers in the reconnaissance stage in the Insider Kill Chain InsiderKillChain"""
# Instantiating the relevant simulation/game objects:
env = uc7_tap003_env()
tap003: TAP003 = env.game.agents["attacker"]
assert tap003.current_kill_chain_stage == BaseKillChain.NOT_STARTED
# Frequency is set to two steps
env = environment_step(i=2, env=env)
# Testing that TAP003 Enters into the expected kill chain stages
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
def test_tap003_kill_chain_stage_planning():
"""Tests the successful/failed handlers in the planning stage in the Insider Kill Chain (TAP003)"""
env = uc7_tap003_env()
tap003: TAP003 = env.game.agents["attacker"]
assert tap003.current_kill_chain_stage == BaseKillChain.NOT_STARTED
env = environment_step(i=2, env=env)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
env = environment_step(i=2, env=env)
# Testing that TAP003 Enters into the expected kill chain stages
assert tap003.current_kill_chain_stage.name == InsiderKillChain.PLANNING.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.ACCESS.name
env = environment_step(i=2, env=env)
# Testing that the stage successfully impacted the simulation - User is logged in
# TODO: Add an assert for this.
def test_tap003_kill_chain_stage_access():
"""Tests the successful/failed handlers in the access stage in the InsiderKillChain"""
env = uc7_tap003_env()
tap003: TAP003 = env.game.agents["attacker"]
assert tap003.current_kill_chain_stage == BaseKillChain.NOT_STARTED
env = environment_step(i=2, env=env)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
env = environment_step(i=2, env=env)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.PLANNING.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.ACCESS.name
env = environment_step(i=2, env=env)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.ACCESS.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
env = environment_step(i=2, env=env)
def test_tap003_kill_chain_stage_manipulation():
"""Tests the successful/failed handlers in the manipulation stage in the InsiderKillChain"""
env = uc7_tap003_env()
env.reset()
tap003: TAP003 = env.game.agents["attacker"]
assert tap003.current_kill_chain_stage == BaseKillChain.NOT_STARTED
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.PLANNING.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.ACCESS.name
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.ACCESS.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
# Testing that the stage successfully impacted the simulation - Accounts Altered
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
st_intra_prv_rt_dr_1: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-DR-1")
assert st_intra_prv_rt_dr_1.user_manager.admins["admin"].password == "red_pass"
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
st_intra_prv_rt_cr: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-CR")
assert st_intra_prv_rt_cr.user_manager.admins["admin"].password == "red_pass"
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
rem_pub_rt_dr: Router = env.game.simulation.network.get_node_by_hostname("REM-PUB-RT-DR")
assert rem_pub_rt_dr.user_manager.admins["admin"].password == "red_pass"
def test_tap003_kill_chain_stage_exploit():
"""Tests the successful/failed handlers in the exploit stage in the InsiderKillChain"""
env = uc7_tap003_env()
tap003: TAP003 = env.game.agents["attacker"]
# The TAP003's Target Router/Firewall
st_intra_prv_rt_dr_1: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-DR-1")
st_intra_prv_rt_cr: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-CR")
rem_pub_rt_dr: Router = env.game.simulation.network.get_node_by_hostname("REM-PUB-RT-DR")
assert tap003.current_kill_chain_stage == BaseKillChain.NOT_STARTED
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.RECONNAISSANCE.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.PLANNING.name
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.PLANNING.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.ACCESS.name
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.ACCESS.name
assert tap003.next_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
env.step(0)
assert tap003.current_kill_chain_stage.name == InsiderKillChain.EXPLOIT.name
# Testing that the stage successfully impacted the simulation - Malicious ACL Added:
for _ in range(32):
env.step(0)
# Tests that the ACL has been added and that the action is deny.
st_intra_prv_rt_dr_1_acl_list = st_intra_prv_rt_dr_1.acl
assert st_intra_prv_rt_dr_1_acl_list.acl[1].action != None
assert st_intra_prv_rt_dr_1_acl_list.acl[1].action == ACLAction.DENY
st_intra_prv_rt_cr_acl_list = st_intra_prv_rt_cr.acl
assert st_intra_prv_rt_cr_acl_list.acl[1].action != None
assert st_intra_prv_rt_cr_acl_list.acl[1].action == ACLAction.DENY
rem_pub_rt_dr_acl_list = rem_pub_rt_dr.acl
assert rem_pub_rt_dr_acl_list.acl[1].action != None
assert rem_pub_rt_dr_acl_list.acl[1].action == ACLAction.DENY

201
tests/e2e_integration_tests/threat_actor_profiles/test_tap003_multiple_rules.py Normal file
View File

@@ -0,0 +1,201 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Protocol
import pytest
import yaml
from primaite.config.load import _EXAMPLE_CFG
from primaite.game.agent.scripted_agents.abstract_tap import (
AbstractTAP,
BaseKillChain,
KillChainOptions,
KillChainStageOptions,
KillChainStageProgress,
)
from primaite.game.agent.scripted_agents.TAP001 import MobileMalwareKillChain
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain
from primaite.session.environment import PrimaiteGymEnv
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
# Defining constants.
ATTACK_AGENT_INDEX = 32
START_STEP = 1 # The starting step of the agent.
FREQUENCY = 2 # The frequency of kill chain stage progression (E.g it's next attempt at "attacking").
VARIANCE = 0 # The timestep variance between kill chain progression (E.g Next timestep = Frequency +/- variance)
REPEAT_KILL_CHAIN = False # Should the TAP repeat the kill chain after success/failure?
REPEAT_KILL_CHAIN_STAGES = False # Should the TAP restart from it's previous stage on failure?
KILL_CHAIN_PROBABILITY = 1 # Blank probability for agent 'success'
RULES = [
{
"target_router": "ST-INTRA-PRV-RT-DR-1",
"position": 1,
"permission": "DENY",
"src_ip": "192.168.220.3",
"src_wildcard": "NONE",
"dst_ip": "192.168.220.3",
"dst_wildcard": "NONE",
"src_port": "ALL",
"dst_port": "ALL",
"protocol_name": "ALL",
},
{
"target_router": "ST-INTRA-PRV-RT-DR-2",
"position": 5,
"permission": "DENY",
"src_ip": "192.168.220.3",
"src_wildcard": "NONE",
"dst_ip": "ALL",
"dst_wildcard": "NONE",
"src_port": "ALL",
"dst_port": "ALL",
"protocol_name": "ALL",
},
{
"target_router": "ST-INTRA-PRV-RT-CR",
"position": 6,
"permission": "PERMIT",
"src_ip": "192.168.220.3",
"src_wildcard": "NONE",
"dst_ip": "ALL",
"dst_wildcard": "NONE",
"src_port": "ALL",
"dst_port": "ALL",
"protocol_name": "ALL",
},
{
"target_router": "REM-PUB-RT-DR",
"position": 3,
"permission": "PERMIT",
"src_ip": "192.168.220.3",
"src_wildcard": "0.0.0.1",
"dst_ip": "192.168.220.3",
"dst_wildcard": "0.0.0.1",
"src_port": "FTP",
"dst_port": "FTP",
"protocol_name": "TCP",
},
#
]
def uc7_tap003_env(**kwargs) -> PrimaiteGymEnv:
"""Setups the UC7 TAP003 Game with the start_step & frequency set to 1 with probabilities set to 1 as well"""
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
cfg = yaml.safe_load(uc7_config)
cfg["io_settings"]["save_sys_logs"] = False
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["start_step"] = START_STEP
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["frequency"] = FREQUENCY
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["variance"] = VARIANCE
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = kwargs[
"repeat_kill_chain_stages"
]
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = kwargs[
"manipulation_probability"
]
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = kwargs[
"access_probability"
]
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs[
"planning_probability"
]
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["EXPLOIT"]["malicious_acls"] = RULES
# Adding the new test target to TAP003's starting knowledge:
new_target_dict = {
"ST-INTRA-PRV-RT-DR-2": {
"ip_address": "192.168.170.2",
"username": "admin",
"password": "admin",
}
}
new_target_manipulation = {
"host": "ST-INTRA-PRV-RT-DR-2",
"ip_address": "192.168.170.2",
"action": "change_password",
"username": "admin",
"new_password": "red_pass",
}
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["starting_network_knowledge"][
"credentials"
].update(new_target_dict)
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"]["account_changes"].append(
new_target_manipulation
)
env = PrimaiteGymEnv(env_config=cfg)
return env
def test_tap003_cycling_rules():
"""Tests to check that TAP003 repeats it's kill chain after success"""
env = uc7_tap003_env(
repeat_kill_chain=True,
repeat_kill_chain_stages=True,
manipulation_probability=1,
access_probability=1,
planning_probability=1,
)
tap003: TAP003 = env.game.agents["attacker"]
def wait_until_attack():
for _ in range(120):
# check if the agent has executed and therefore moved onto the next rule index
env.step(0)
if tap003.history[-1].action == "node-send-remote-command":
if tap003.history[-1].parameters["command"][0] == "acl":
return
pytest.fail("While testing the cycling of TAP003 rules, the agent unexpectedly didn't execute its attack.")
wait_until_attack()
target_node: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-DR-1")
assert (rule_0 := target_node.acl.acl[1]) is not None
assert rule_0.action == ACLAction.DENY
assert rule_0.protocol == None
assert rule_0.src_ip_address == IPV4Address("192.168.220.3")
assert rule_0.src_wildcard_mask == None
assert rule_0.dst_ip_address == IPV4Address("192.168.220.3")
assert rule_0.dst_wildcard_mask == None
assert rule_0.src_port == None
assert rule_0.dst_port == None
target_node: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-DR-2")
wait_until_attack()
assert (rule_1 := target_node.acl.acl[5]) is not None
assert rule_1.action == ACLAction.DENY
assert rule_1.protocol == None
assert rule_1.src_ip_address == IPV4Address("192.168.220.3")
assert rule_1.src_wildcard_mask == None
assert rule_1.dst_ip_address == None
assert rule_1.dst_wildcard_mask == None
assert rule_1.src_port == None
assert rule_1.dst_port == None
wait_until_attack()
target_node: Router = env.game.simulation.network.get_node_by_hostname("ST-INTRA-PRV-RT-CR")
assert (rule_2 := target_node.acl.acl[6]) is not None
assert rule_2.action == ACLAction.PERMIT
assert rule_2.protocol == None
assert rule_2.src_ip_address == IPV4Address("192.168.220.3")
assert rule_2.src_wildcard_mask == None # default
assert rule_2.dst_ip_address == None
assert rule_2.dst_wildcard_mask == None # default
assert rule_2.src_port == None
assert rule_2.dst_port == None
wait_until_attack()
target_node: Router = env.game.simulation.network.get_node_by_hostname("REM-PUB-RT-DR")
assert (rule_3 := target_node.acl.acl[3]) is not None
assert rule_3.action == ACLAction.PERMIT
assert rule_3.protocol == PROTOCOL_LOOKUP["TCP"]
assert rule_3.src_ip_address == IPV4Address("192.168.220.3")
assert rule_3.src_wildcard_mask == IPV4Address("0.0.0.1")
assert rule_3.dst_ip_address == IPV4Address("192.168.220.3")
assert rule_3.dst_wildcard_mask == IPV4Address("0.0.0.1")
assert rule_3.src_port == PORT_LOOKUP["FTP"]
assert rule_3.dst_port == PORT_LOOKUP["FTP"]
# If we've gotten this fair then we can pass the test :)
pass

41
tests/integration_tests/configuration_file_parsing/test_game_options_config.py
View File

@@ -8,7 +8,7 @@ from primaite.config.load import data_manipulation_config_path
from primaite.game.game import PrimaiteGame
from tests import TEST_ASSETS_ROOT
BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
BASIC_SWITCHED_NETWORK_CONFIG = TEST_ASSETS_ROOT / "configs/basic_switched_network.yaml"
def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
@@ -24,3 +24,42 @@ def test_thresholds():
game = load_config(data_manipulation_config_path())
assert game.options.thresholds is not None
def test_nmne_threshold():
"""Test that the NMNE thresholds are properly loaded in by observation."""
game = load_config(BASIC_SWITCHED_NETWORK_CONFIG)
assert game.options.thresholds["nmne"] is not None
# get NIC observation
nic_obs = game.agents["defender"].observation_manager.obs.components["NODES"].hosts[0].nics[0]
assert nic_obs.low_nmne_threshold == 5
assert nic_obs.med_nmne_threshold == 25
assert nic_obs.high_nmne_threshold == 100
def test_file_access_threshold():
"""Test that the NMNE thresholds are properly loaded in by observation."""
game = load_config(BASIC_SWITCHED_NETWORK_CONFIG)
assert game.options.thresholds["file_access"] is not None
# get file observation
file_obs = game.agents["defender"].observation_manager.obs.components["NODES"].hosts[0].folders[0].files[0]
assert file_obs.low_file_access_threshold == 2
assert file_obs.med_file_access_threshold == 5
assert file_obs.high_file_access_threshold == 10
def test_app_executions_threshold():
"""Test that the NMNE thresholds are properly loaded in by observation."""
game = load_config(BASIC_SWITCHED_NETWORK_CONFIG)
assert game.options.thresholds["app_executions"] is not None
# get application observation
app_obs = game.agents["defender"].observation_manager.obs.components["NODES"].hosts[0].applications[0]
assert app_obs.low_app_execution_threshold == 2
assert app_obs.med_app_execution_threshold == 3
assert app_obs.high_app_execution_threshold == 5

64
tests/integration_tests/configuration_file_parsing/test_node_file_system_config.py Normal file
View File

@@ -0,0 +1,64 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from pathlib import Path
from typing import Union
import yaml
from primaite.game.game import PrimaiteGame
from primaite.simulator.file_system.file_type import FileType
from tests import TEST_ASSETS_ROOT
BASIC_CONFIG = TEST_ASSETS_ROOT / "configs/nodes_with_initial_files.yaml"
def load_config(config_path: Union[str, Path]) -> PrimaiteGame:
"""Returns a PrimaiteGame object which loads the contents of a given yaml path."""
with open(config_path, "r") as f:
cfg = yaml.safe_load(f)
return PrimaiteGame.from_config(cfg)
def test_node_file_system_from_config():
"""Test that the appropriate files are instantiated in nodes when loaded from config."""
game = load_config(BASIC_CONFIG)
client_1 = game.simulation.network.get_node_by_hostname("client_1")
assert client_1.software_manager.software.get("database-service") # database service should be installed
assert client_1.file_system.get_file(folder_name="database", file_name="database.db") # database files should exist
assert client_1.software_manager.software.get("web-server") # web server should be installed
assert client_1.file_system.get_file(folder_name="primaite", file_name="index.html") # web files should exist
client_2 = game.simulation.network.get_node_by_hostname("client_2")
# database service should not be installed
assert client_2.software_manager.software.get("database-service") is None
# database files should not exist
assert client_2.file_system.get_file(folder_name="database", file_name="database.db") is None
# web server should not be installed
assert client_2.software_manager.software.get("web-server") is None
# web files should not exist
assert client_2.file_system.get_file(folder_name="primaite", file_name="index.html") is None
empty_folder = client_2.file_system.get_folder(folder_name="empty_folder")
assert empty_folder
assert len(empty_folder.files) == 0 # should have no files
password_file = client_2.file_system.get_file(folder_name="root", file_name="passwords.txt")
assert password_file # should exist
assert password_file.file_type is FileType.TXT
assert password_file.size == 663
downloads_folder = client_2.file_system.get_folder(folder_name="downloads")
assert downloads_folder # downloads folder should exist
test_txt = downloads_folder.get_file(file_name="test.txt")
assert test_txt # test.txt should exist
assert test_txt.file_type is FileType.TXT
unknown_file_type = downloads_folder.get_file(file_name="another_file.pwtwoti")
assert unknown_file_type # unknown_file_type should exist
assert unknown_file_type.file_type is FileType.UNKNOWN

Some files were not shown because too many files have changed in this diff Show More