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

Merged PR 612: #2925 Use Case 7 Scenario Modelling

Browse Source 
## Summary
This PR contains the entirety of the #2925 UC7 implementation and updated to work with 4.0.0.

Specifically, this PR contains the following:

- New UC7 Scenario config (#2483)
- New UC7 TAP001 Config #2909
- New UC7 TAP003 Config #2910
- New UC7 default blue agent #3070
- New UC7 green agent POL #3067
- Multiple UC7 detailed diagrams #3068
- Multiple new UC7 notebooks #3069

## Test process

Pre-existing tests have been re-modelled to use UC7 as well as a few more new UC7 specific tests that ensure all of the expected default behaviour is working.

Additionally, multiple notebooks exist which utilise a large amount of the UC7 scenario and thus also act as tests.

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

Related work items: #2483, #2909, #2910, #3067, #3068, #3069, #3070, #3071, #3086
This commit is contained in:
Archer Bowen
2025-03-03 09:13:08 +00:00
committed by Marek Wolan
parent 75eb1050c9 83c9a90c0a
commit dd5571359a
53 changed files with 17424 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

File diff suppressed because it is too large Load Diff

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

File diff suppressed because it is too large Load Diff

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

@@ -0,0 +1,40 @@
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:
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.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

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

@@ -0,0 +1,40 @@
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:
kill_chain:
ACTIVATE:
probability: 1
PROPAGATE:
probability: 1
scan_attempts: 20
repeat_scan: false
network_addresses:
- 192.168.240.0/29 # ST Project B
- 192.168.10.0/26 # Remote Site
- 192.168.20.0/30 # Remote DMZ
- 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

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

@@ -0,0 +1,40 @@
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:
kill_chain:
ACTIVATE:
probability: 1
PROPAGATE:
probability: 1
scan_attempts: 20
repeat_scan: false
network_addresses:
- 192.168.250.0/29 # ST Project C
- 192.168.10.0/26 # Remote Site
- 192.168.20.0/30 # Remote DMZ
- 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

2635
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

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

@@ -37,15 +37,18 @@ class NodeAbstractAction(AbstractAction, ABC):
return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, discriminator="node-os-scan"):
class NodeOSScanAction(AbstractAction, discriminator="node-os-scan"):
"""Action which scans a node's OS."""
config: "NodeOSScanAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for Node actions."""
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration schema for NodeOSScanAction."""
node_name: str
verb: ClassVar[str] = "scan"
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", config.node_name, "os", "scan"]
class NodeShutdownAction(NodeAbstractAction, discriminator="node-shutdown"):

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

@@ -124,8 +124,8 @@ class AbstractAgent(BaseModel, ABC):
pass
else:
# format dict by putting each key-value entry on a separate line and putting a blank line on the end.
param_string = "\n".join([*[f"{k}: {v:.30}" for k, v in item.parameters.items()], ""])
data_string = "\n".join([*[f"{k}: {v:.30}" for k, v in item.response.data], ""])
param_string = "\n".join([*[f"{k}: {str(v):.80}" for k, v in item.parameters.items()], ""])
data_string = "\n".join([*[f"{k}: {str(v):.80}" for k, v in item.response.data.items()], ""])
table.add_row([item.timestep, item.action, param_string, item.response.status, data_string])
print(table)

23
src/primaite/game/agent/scripted_agents/TAP003.py
View File

@@ -257,7 +257,7 @@ class TAP003(AbstractTAP, discriminator="tap-003"):
self.logger.debug(f"Updating network knowledge. Changed {username}'s password to {password} on {hostname}.")
self._change_password_target_host = ""
# local password change
elif last_hist_item.action == "node-accounts-change-password" and last_hist_item.response.status == "success":
elif last_hist_item.action == "node-account-change-password" and last_hist_item.response.status == "success":
self.network_knowledge["current_session"] = {}
username = last_hist_item.request[6]
password = last_hist_item.request[8]
@@ -338,15 +338,17 @@ class TAP003(AbstractTAP, discriminator="tap-003"):
"""
if self.current_kill_chain_stage == self.selected_kill_chain.MANIPULATION:
if self._agent_trial_handler(self.config.agent_settings.kill_chain.MANIPULATION.probability):
self.logger.info(f"TAP003 reached the {self.current_kill_chain_stage.name}")
if self.current_stage_progress == KillChainStageProgress.PENDING:
self.logger.info(f"TAP003 reached the {self.current_kill_chain_stage.name}.")
self.current_stage_progress = KillChainStageProgress.IN_PROGRESS
self.current_host = self.starting_node
account_changes = self.config.agent_settings.kill_chain.MANIPULATION.account_changes
if len(account_changes) > 0:
if len(account_changes) > 0 or self._next_account_change:
if not self._next_account_change:
self._next_account_change = account_changes.pop(0)
if self._next_account_change["host"] == self.current_host:
# do a local password change
self.chosen_action = "node-accounts-change-password", {
self.chosen_action = "node-account-change-password", {
"node_name": self.current_host,
"username": self._next_account_change["username"],
"current_password": self.network_knowledge["credentials"][self.current_host]["password"],
@@ -382,14 +384,15 @@ class TAP003(AbstractTAP, discriminator="tap-003"):
],
}
self.logger.info(f"Changing password on remote node {hostname}")
self._next_account_change = account_changes.pop(0)
if len(account_changes) == 0:
self.logger.info("No further account changes required.")
self._next_account_change = None
else:
self._next_account_change = account_changes.pop(0)
self._change_password_target_host = hostname
if len(account_changes) == 0:
self._next_account_change = None
self.logger.info("Finished changing passwords.")
if not self._next_account_change:
self.logger.info("Manipulation complete. Progressing to exploit...")
self._progress_kill_chain()
self.current_stage_progress = KillChainStageProgress.PENDING
else:
if self.config.agent_settings.repeat_kill_chain_stages == False:
self.current_kill_chain_stage = self.selected_kill_chain.FAILED

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

@@ -14,6 +14,7 @@ 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

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

@@ -1041,7 +1041,7 @@
"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",
@@ -1129,7 +1129,7 @@
"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",
@@ -1254,7 +1254,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)"
]
},
@@ -1670,6 +1670,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": {

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

@@ -380,18 +380,6 @@
"!primaite setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -712,7 +700,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": ".venv",
"language": "python",
"name": "python3"
},
@@ -726,7 +714,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.11"
"version": "3.10.12"
}
},
"nbformat": 4,

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

@@ -16,6 +16,13 @@
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simulation Layer Implementation."
]
},
{
"cell_type": "markdown",
"metadata": {},

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

File diff suppressed because it is too large Load Diff

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

File diff suppressed because it is too large Load Diff

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

File diff suppressed because it is too large Load Diff

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

@@ -0,0 +1,147 @@
{
"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": [
"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 `do-nothing` 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-nmap\"):\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 do-nothing 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
}

1126
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.2 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_tap003/uc7_tap003_main.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

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

@@ -1534,6 +1534,12 @@ class Node(SimComponent, ABC):
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
red_scan_countdown: int = 0
"Time steps until reveal to red scan is complete."
node_scan_countdown: int = 0
"Time steps until scan is complete"
# TODO: this should not be set for abstract classes.
_discriminator: ClassVar[str]
"""discriminator for this particular class, used for printing and logging. Each subclass redefines this."""
@@ -1570,12 +1576,6 @@ class Node(SimComponent, ABC):
node_scan_duration: int = 10
"How many timesteps until the whole node is scanned. Default 10 time steps."
node_scan_countdown: int = 0
"Time steps until scan is complete"
red_scan_countdown: int = 0
"Time steps until reveal to red scan is complete."
dns_server: Optional[IPv4Address] = None
"List of IP addresses of DNS servers used for name resolution."
@@ -2019,10 +2019,10 @@ class Node(SimComponent, ABC):
# time steps which require the node to be on
if self.operating_state == NodeOperatingState.ON:
# node scanning
if self.config.node_scan_countdown > 0:
self.config.node_scan_countdown -= 1
if self.node_scan_countdown > 0:
self.node_scan_countdown -= 1
if self.config.node_scan_countdown == 0:
if self.node_scan_countdown == 0:
# scan everything!
for process_id in self.processes:
self.processes[process_id].scan()
@@ -2038,10 +2038,10 @@ class Node(SimComponent, ABC):
# scan file system
self.file_system.scan(instant_scan=True)
if self.config.red_scan_countdown > 0:
self.config.red_scan_countdown -= 1
if self.red_scan_countdown > 0:
self.red_scan_countdown -= 1
if self.config.red_scan_countdown == 0:
if self.red_scan_countdown == 0:
# scan processes
for process_id in self.processes:
self.processes[process_id].reveal_to_red()
@@ -2098,7 +2098,7 @@ class Node(SimComponent, ABC):
to the red agent.
"""
self.config.node_scan_countdown = self.config.node_scan_duration
self.node_scan_countdown = self.config.node_scan_duration
return True
def reveal_to_red(self) -> bool:
@@ -2114,7 +2114,7 @@ class Node(SimComponent, ABC):
`revealed_to_red` to `True`.
"""
self.config.red_scan_countdown = self.config.node_scan_duration
self.red_scan_countdown = self.config.node_scan_duration
return True
def power_on(self) -> bool: