Merge remote-tracking branch 'origin/dev' into feature/Updated-How-To-Guides
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
2677
src/primaite/config/_package_data/uc7_config.yaml
Normal file
2730
src/primaite/config/_package_data/uc7_config_tap003.yaml
Normal 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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -65,8 +65,7 @@ class ACLObservation(AbstractObservation, discriminator="acl"):
|
||||
self.port_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(port_list)}
|
||||
self.protocol_to_id: Dict[str, int] = {p: i + 2 for i, p in enumerate(protocol_list)}
|
||||
self.default_observation: Dict = {
|
||||
i
|
||||
+ 1: {
|
||||
i: {
|
||||
"position": i,
|
||||
"permission": 0,
|
||||
"source_ip_id": 0,
|
||||
@@ -94,12 +93,11 @@ class ACLObservation(AbstractObservation, discriminator="acl"):
|
||||
return self.default_observation
|
||||
obs = {}
|
||||
acl_items = dict(acl_state.items())
|
||||
i = 1 # don't show rule 0 for compatibility reasons.
|
||||
while i < self.num_rules + 1:
|
||||
for i in range(self.num_rules):
|
||||
rule_state = acl_items[i]
|
||||
if rule_state is None:
|
||||
obs[i] = {
|
||||
"position": i - 1,
|
||||
"position": i,
|
||||
"permission": 0,
|
||||
"source_ip_id": 0,
|
||||
"source_wildcard_id": 0,
|
||||
@@ -125,7 +123,7 @@ class ACLObservation(AbstractObservation, discriminator="acl"):
|
||||
protocol = rule_state["protocol"]
|
||||
protocol_id = self.protocol_to_id.get(protocol, 1)
|
||||
obs[i] = {
|
||||
"position": i - 1,
|
||||
"position": i,
|
||||
"permission": rule_state["action"],
|
||||
"source_ip_id": src_node_id,
|
||||
"source_wildcard_id": src_wildcard_id,
|
||||
@@ -135,7 +133,6 @@ class ACLObservation(AbstractObservation, discriminator="acl"):
|
||||
"dest_port_id": dst_port_id,
|
||||
"protocol_id": protocol_id,
|
||||
}
|
||||
i += 1
|
||||
return obs
|
||||
|
||||
@property
|
||||
@@ -148,8 +145,7 @@ class ACLObservation(AbstractObservation, discriminator="acl"):
|
||||
"""
|
||||
return spaces.Dict(
|
||||
{
|
||||
i
|
||||
+ 1: spaces.Dict(
|
||||
i: spaces.Dict(
|
||||
{
|
||||
"position": spaces.Discrete(self.num_rules),
|
||||
"permission": spaces.Discrete(3),
|
||||
|
||||
1080
src/primaite/game/agent/scripted_agents/TAP001.py
Normal file
468
src/primaite/game/agent/scripted_agents/TAP003.py
Normal file
@@ -0,0 +1,468 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from enum import IntEnum
|
||||
from typing import Dict, List, Literal, Optional, Tuple, Type
|
||||
|
||||
from gymnasium.core import ObsType
|
||||
from pydantic import Field
|
||||
|
||||
from primaite.game.agent.actions.acl import RouterACLAddRuleAction
|
||||
from primaite.game.agent.scripted_agents.abstract_tap import (
|
||||
AbstractTAP,
|
||||
KillChainOptions,
|
||||
KillChainStageOptions,
|
||||
KillChainStageProgress,
|
||||
)
|
||||
|
||||
|
||||
class InsiderKillChainOptions(KillChainOptions):
|
||||
"""Model validation for TAP003's Kill Chain."""
|
||||
|
||||
class _PlanningOptions(KillChainStageOptions):
|
||||
"""Valid options for the `PLANNING` InsiderKillChain stage."""
|
||||
|
||||
starting_network_knowledge: Dict # TODO: more specific schema here?
|
||||
|
||||
class _AccessOptions(KillChainStageOptions):
|
||||
"""Valid options for the `ACCESS` InsiderKillChain stage."""
|
||||
|
||||
pass
|
||||
|
||||
class _ManipulationOptions(KillChainStageOptions):
|
||||
"""Valid options for the `MANIPULATION` InsiderKillChain stage."""
|
||||
|
||||
account_changes: List[Dict] = [] # TODO: More specific schema here?
|
||||
|
||||
class _ExploitOptions(KillChainStageOptions):
|
||||
"""Valid options for the `EXPLOIT` InsiderKillChain stage."""
|
||||
|
||||
malicious_acls: List[RouterACLAddRuleAction.ConfigSchema] = []
|
||||
|
||||
PLANNING: _PlanningOptions = Field(default_factory=lambda: InsiderKillChainOptions._PlanningOptions())
|
||||
ACCESS: _AccessOptions = Field(default_factory=lambda: InsiderKillChainOptions._AccessOptions())
|
||||
MANIPULATION: _ManipulationOptions = Field(default_factory=lambda: InsiderKillChainOptions._ManipulationOptions())
|
||||
EXPLOIT: _ExploitOptions = Field(default_factory=lambda: InsiderKillChainOptions._ExploitOptions())
|
||||
|
||||
|
||||
class InsiderKillChain(IntEnum):
|
||||
"""
|
||||
Enumeration representing different attack stages of the vulnerability and backdoor creation kill chain.
|
||||
|
||||
This kill chain is designed around the TAP003 - Malicious Insider Corporal Pearson.
|
||||
Each stage represents a specific phase in the kill chain.
|
||||
Please refer to the TAP003 notebook for the current version's implementation of this kill chain.
|
||||
"""
|
||||
|
||||
RECONNAISSANCE = 1
|
||||
"Represents TAP003 identifying sensitive systems, data and access control mechanisms"
|
||||
PLANNING = 2
|
||||
"Represents TAP003 devising a plan to exploit their elevated privileges."
|
||||
ACCESS = 3
|
||||
"Represents TAP003's using legitimate credentials to access the access control settings."
|
||||
MANIPULATION = 4
|
||||
"Represents TAP003 altering ACLs, User & Group Attributes & other control mechanisms to grant unauthorised access"
|
||||
EXPLOIT = 5
|
||||
"Represents TAP003 exploiting their insider knowledge and privilege to implement changes for sabotage."
|
||||
EMBED = 6
|
||||
"Represents TAP003's additional changes to ensure continued access"
|
||||
CONCEAL = 7
|
||||
"Represents TAP003's efforts in hiding their traces of malicious activities"
|
||||
EXTRACT = 8
|
||||
"Represents TAP003 removing sensitive data from the organisation, either for personal gain or to inflict harm."
|
||||
ERASE = 9
|
||||
"Represents TAP003 covering their tracks by removing any tools, reverting temporary changes and logging out"
|
||||
|
||||
# These Enums must be included in all kill chains.
|
||||
# Due to limitations in Python and Enums, it is not possible to inherit these Enums from an base class.
|
||||
|
||||
NOT_STARTED = 100
|
||||
"Indicates that the Kill Chain has not started."
|
||||
SUCCEEDED = 200
|
||||
"Indicates that the kill chain has succeeded."
|
||||
FAILED = 300
|
||||
"Indicates that the attack has failed."
|
||||
|
||||
def initial_stage(self) -> "InsiderKillChain":
|
||||
"""Returns the first stage in the kill chain. Used by Abstract TAP for TAP Agent Setup."""
|
||||
return self.RECONNAISSANCE
|
||||
|
||||
|
||||
class TAP003(AbstractTAP, discriminator="tap-003"):
|
||||
"""
|
||||
TAP003 | Malicious Insider Corporal Pearson.
|
||||
|
||||
Currently implements one kill chain: Backdoor & Vulnerability Creation.
|
||||
This Threat Actor Profile (TAP) aims to introduce subtle cyber attack.
|
||||
For example, the Backdoor & Vulnerability creation kill chain
|
||||
creates DENY firewall rules which do not trigger NMNE.
|
||||
Please see the TAP003-Kill-Chain-E2E.ipynb for more information.
|
||||
"""
|
||||
|
||||
class AgentSettingsSchema(AbstractTAP.AgentSettingsSchema):
|
||||
"""Agent Settings Schema that enforces TAP003's `kill_chain` config to use the InsiderKillChainOptions."""
|
||||
|
||||
kill_chain: InsiderKillChainOptions # = Field(default_factory=lambda: MobileMalwareKillChainOptions())
|
||||
|
||||
class ConfigSchema(AbstractTAP.ConfigSchema):
|
||||
"""Config Schema for the TAP001 agent."""
|
||||
|
||||
type: Literal["tap-003"] = "tap-003"
|
||||
agent_settings: "TAP003.AgentSettingsSchema" = Field(default_factory=lambda: TAP003.AgentSettingsSchema())
|
||||
|
||||
config: ConfigSchema
|
||||
selected_kill_chain: Type[InsiderKillChain] = InsiderKillChain
|
||||
_current_acl: int = 0
|
||||
network_knowledge: Dict = {} # TODO: more specific typing
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._change_password_target_host: str = ""
|
||||
"""If we have just sent a change password request over SSH, this variable keeps track of the hostname."""
|
||||
self._ssh_target_host: str = ""
|
||||
"""If we have just send a SSH_LOGIN request, keeps track of the hostname to which we are attempting to SSH."""
|
||||
self._next_account_change: Optional[Dict] = None
|
||||
self._num_acls = len(self.config.agent_settings.kill_chain.EXPLOIT.malicious_acls)
|
||||
|
||||
self.network_knowledge: dict = {"credentials": {}, "current_session": {}}
|
||||
"""Keep track of current network state based on responses after sending actions. Populated during PLANNING."""
|
||||
self.setup_agent()
|
||||
|
||||
def _progress_kill_chain(self) -> None:
|
||||
"""Private Method used to progress the kill chain to the next stage."""
|
||||
if self.next_kill_chain_stage == self.selected_kill_chain.EXPLOIT: # Covering final stage edge case.
|
||||
self.current_kill_chain_stage = self.selected_kill_chain(self.current_kill_chain_stage + 1)
|
||||
self.next_kill_chain_stage = self.selected_kill_chain.SUCCEEDED
|
||||
else:
|
||||
# Otherwise, set the current stage to the previous next and increment the next kill chain stage.
|
||||
self.current_kill_chain_stage = self.next_kill_chain_stage
|
||||
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.SUCCEEDED:
|
||||
self.next_kill_chain_stage = self.selected_kill_chain.NOT_STARTED
|
||||
else:
|
||||
self.next_kill_chain_stage = self.selected_kill_chain(self.current_kill_chain_stage + 1)
|
||||
|
||||
self.current_stage_progress = KillChainStageProgress.PENDING
|
||||
|
||||
def setup_agent(self) -> None:
|
||||
"""Responsible for agent setup upon episode reset.
|
||||
|
||||
Explicitly this method performs the following:
|
||||
1. Loads the inherited attribute 'selected_kill_chain' with the InsiderKillChain
|
||||
2. Selects the starting node from the given user tap settings
|
||||
3. Selects the target node from the given user tap settings
|
||||
4. Sets the next execution timestep to the given user tap settings - start step
|
||||
5. Sets TAP's current host as the selected starting node.
|
||||
"""
|
||||
# TAP Boilerplate Setup
|
||||
self._setup_agent_kill_chain(InsiderKillChain)
|
||||
|
||||
# TAP003 Specific Setup
|
||||
self._select_start_node()
|
||||
self._set_next_execution_timestep(self.config.agent_settings.start_step)
|
||||
self.current_host = self.starting_node
|
||||
|
||||
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
|
||||
"""Follows the TAP003 Backdoor Vulnerability Kill Chain.
|
||||
|
||||
Calls the next TAP003 Action Stage. Uses private methods to schedule kill chain stages.
|
||||
See TAP003-Kill-Chain-E2E.ipynb for further information on the TAP003 agent.
|
||||
|
||||
:param obs: Current observation for this agent.
|
||||
:type obs: ObsType
|
||||
:param timestep: The current simulation timestep, used for scheduling actions
|
||||
:type timestep: int
|
||||
:return: Action formatted in CAOS format
|
||||
:rtype: Tuple[str, Dict]
|
||||
"""
|
||||
self._handle_login_response()
|
||||
self._handle_change_password_response()
|
||||
if timestep < self.next_execution_timestep or self.actions_concluded:
|
||||
return "do-nothing", {} # bypasses self.chosen_action
|
||||
|
||||
# self.current_timestep is currently the previous execution timestep
|
||||
# So it can be used to index action history.
|
||||
if not self._tap_return_handler(self.current_timestep):
|
||||
# If the application is already installed, don't keep retrying - this is an acceptable fail
|
||||
if self.current_kill_chain_stage == InsiderKillChain.PLANNING:
|
||||
last_action = self.history[self.current_timestep].action
|
||||
fail_reason = self.history[self.current_timestep].response.data["reason"]
|
||||
if last_action == "node-application-install" and fail_reason == "already installed":
|
||||
pass
|
||||
else:
|
||||
self.update_current_timestep(new_timestep=timestep)
|
||||
self._set_next_execution_timestep(timestep + self.config.agent_settings.frequency)
|
||||
self._tap_outcome_handler(InsiderKillChain)
|
||||
return self.chosen_action
|
||||
|
||||
self.update_current_timestep(new_timestep=timestep)
|
||||
self._set_next_execution_timestep(timestep + self.config.agent_settings.frequency)
|
||||
self._tap_outcome_handler(InsiderKillChain) # Handles successes and failures
|
||||
|
||||
# The kill chain is called in reverse order
|
||||
# The kill chain sequence must be called in reverse order to ensure proper execution.
|
||||
|
||||
self._exploit()
|
||||
self._manipulation()
|
||||
self._access()
|
||||
self._planning()
|
||||
self._reconnaissance()
|
||||
self._tap_start(InsiderKillChain)
|
||||
|
||||
return self.chosen_action
|
||||
|
||||
def _handle_login_response(self) -> None:
|
||||
"""If the last request was an SSH login attempt, update the current session in network knowledge."""
|
||||
if not self.history:
|
||||
return
|
||||
last_hist_item = self.history[-1]
|
||||
if not last_hist_item.action == "node-session-remote-login" or last_hist_item.response.status != "success":
|
||||
return
|
||||
|
||||
self.network_knowledge["current_session"] = {
|
||||
"hostname": self._ssh_target_host,
|
||||
"ip_address": last_hist_item.response.data["ip_address"],
|
||||
"username": last_hist_item.response.data["username"],
|
||||
}
|
||||
self.logger.debug(
|
||||
f"Updating network knowledge. Logged in as {last_hist_item.response.data['username']} on "
|
||||
f"{self._ssh_target_host}"
|
||||
)
|
||||
|
||||
def _handle_change_password_response(self) -> None:
|
||||
if not self.history:
|
||||
return
|
||||
last_hist_item = self.history[-1]
|
||||
|
||||
# when sending remote change password command, this must get populated
|
||||
if not self._change_password_target_host:
|
||||
return
|
||||
|
||||
if (
|
||||
last_hist_item.action == "node-send-remote-command"
|
||||
and last_hist_item.parameters["command"][2] == "change_password"
|
||||
and last_hist_item.response.status == "success"
|
||||
):
|
||||
# changing password logs us out, so our current session needs to be cleared
|
||||
self.network_knowledge["current_session"] = {}
|
||||
|
||||
# update internal knowledge with the new password
|
||||
ip = last_hist_item.parameters["remote_ip"]
|
||||
username = last_hist_item.parameters["command"][3]
|
||||
password = last_hist_item.parameters["command"][5]
|
||||
hostname = self._change_password_target_host
|
||||
self.network_knowledge["credentials"][hostname] = {
|
||||
"ip_address": ip,
|
||||
"username": username,
|
||||
"password": password,
|
||||
}
|
||||
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-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]
|
||||
hostname = last_hist_item.request[2]
|
||||
self.network_knowledge["credentials"][hostname] = {"username": username, "password": password}
|
||||
self.logger.debug(f"Updating network knowledge. Changed {username}'s password to {password} on {hostname}.")
|
||||
self._change_password_target_host = ""
|
||||
|
||||
def _reconnaissance(self) -> None:
|
||||
"""Insider Kill Chain | Reconnaissance Stage.
|
||||
|
||||
First stage in the Insider kill chain.
|
||||
|
||||
Sets the self.chosen attribute to the "do-nothing" CAOS action
|
||||
and then calls the self._progress_kill_chain() method.
|
||||
"""
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.RECONNAISSANCE:
|
||||
self.chosen_action = "do-nothing", {}
|
||||
self._progress_kill_chain()
|
||||
|
||||
def _planning(self) -> None:
|
||||
"""Insider Kill Chain | Planning Stage.
|
||||
|
||||
Second stage in the Insider kill chain.
|
||||
Performs a trial using the given user PLANNING stage probability.
|
||||
|
||||
If the trial is successful then the agent populates its knowledge base with information from the config.
|
||||
|
||||
Otherwise, the stage is not progressed. Additionally, the agent's kill chain is set
|
||||
to failure if the repeat_kill_chain_stages parameter is set to FALSE.
|
||||
"""
|
||||
if not self.current_kill_chain_stage == self.selected_kill_chain.PLANNING:
|
||||
return
|
||||
|
||||
if not self._agent_trial_handler(self.config.agent_settings.kill_chain.PLANNING.probability):
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
self.chosen_action = "do-nothing", {}
|
||||
|
||||
else:
|
||||
self.network_knowledge[
|
||||
"credentials"
|
||||
] = self.config.agent_settings.kill_chain.PLANNING.starting_network_knowledge["credentials"]
|
||||
self.current_host = self.starting_node
|
||||
self.logger.info("Resolving starting knowledge.")
|
||||
self._progress_kill_chain()
|
||||
if self.current_stage_progress == KillChainStageProgress.PENDING:
|
||||
self.logger.info(f"TAP003 reached the {self.current_kill_chain_stage.name}")
|
||||
|
||||
def _access(self) -> None:
|
||||
"""Insider Kill Chain | Planning Stage.
|
||||
|
||||
Third stage in the Insider kill chain.
|
||||
Performs a trial using the given user ACCESS stage probability.
|
||||
|
||||
This currently does nothing.
|
||||
"""
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.ACCESS:
|
||||
if self._agent_trial_handler(self.config.agent_settings.kill_chain.ACCESS.probability):
|
||||
self._progress_kill_chain()
|
||||
self.chosen_action = "do-nothing", {}
|
||||
else:
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
self.chosen_action = "do-nothing", {}
|
||||
|
||||
def _manipulation(self) -> None:
|
||||
"""Insider Kill Chain | Manipulation Stage.
|
||||
|
||||
Fourth stage in the Insider kill chain.
|
||||
Performs a trial using the given user MANIPULATION stage probability.
|
||||
|
||||
If the trial is successful, the agent will change passwords for accounts that will later be used to execute
|
||||
malicious commands
|
||||
|
||||
Otherwise if the stage is not progressed. Additionally, the agent's kill chain is set
|
||||
to failure if the repeat_kill_chain_stages parameter is set to FALSE.
|
||||
"""
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.MANIPULATION:
|
||||
if self._agent_trial_handler(self.config.agent_settings.kill_chain.MANIPULATION.probability):
|
||||
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 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-account-change-password", {
|
||||
"node_name": self.current_host,
|
||||
"username": self._next_account_change["username"],
|
||||
"current_password": self.network_knowledge["credentials"][self.current_host]["password"],
|
||||
"new_password": self._next_account_change["new_password"],
|
||||
}
|
||||
self.logger.info("Changing local password.")
|
||||
self._next_account_change = account_changes.pop(0)
|
||||
self._change_password_target_host = self.current_host
|
||||
else:
|
||||
# make sure we are logged in via ssh to remote node
|
||||
hostname = self._next_account_change["host"]
|
||||
if self.network_knowledge.get("current_session", {}).get("hostname") != hostname:
|
||||
self._ssh_target_host = hostname
|
||||
self.chosen_action = "node-session-remote-login", {
|
||||
"node_name": self.starting_node,
|
||||
"username": self.network_knowledge["credentials"][hostname]["username"],
|
||||
"password": self.network_knowledge["credentials"][hostname]["password"],
|
||||
"remote_ip": self.network_knowledge["credentials"][hostname]["ip_address"],
|
||||
}
|
||||
self.logger.info(f"Logging into {hostname} in order to change password.")
|
||||
# once we know we are logged in, send a command to change password
|
||||
else:
|
||||
self.chosen_action = "node-send-remote-command", {
|
||||
"node_name": self.starting_node,
|
||||
"remote_ip": self.network_knowledge["credentials"][hostname]["ip_address"],
|
||||
"command": [
|
||||
"service",
|
||||
"user-manager",
|
||||
"change_password",
|
||||
self._next_account_change["username"],
|
||||
self.network_knowledge["credentials"][hostname]["password"],
|
||||
self._next_account_change["new_password"],
|
||||
],
|
||||
}
|
||||
self.logger.info(f"Changing password on remote node {hostname}")
|
||||
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 not self._next_account_change:
|
||||
self.logger.info("Manipulation complete. Progressing to exploit...")
|
||||
self._progress_kill_chain()
|
||||
else:
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
self.chosen_action = "do-nothing", {}
|
||||
|
||||
def _exploit(self) -> None:
|
||||
"""Insider Kill Chain | Exploit Stage.
|
||||
|
||||
Fifth stage in the Insider kill chain.
|
||||
Performs a trial using the given user EXPLOIT stage probability.
|
||||
|
||||
If the trial is successful then self.chosen_action attribute is set to the
|
||||
"node-send-remote-command" CAOS action with the "ROUTER_ACL_ADDRULE" as it's chosen command.
|
||||
|
||||
The impact of the ROUTER_ACL_ADDRULE is dependant on user given parameters. At current
|
||||
the default impact of this stage is to block green agent traffic. An example of TAP003's
|
||||
manipulation stage in action can be found in the TAP003 notebook.
|
||||
|
||||
Otherwise if the stage is not progressed. Additionally, the agent's kill chain is set
|
||||
to failure if the repeat_kill_chain_stages parameter is set to FALSE.
|
||||
"""
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.EXPLOIT:
|
||||
if self.current_kill_chain_stage == KillChainStageProgress.PENDING:
|
||||
# Perform the probability of success once upon entering the stage.
|
||||
if not self._agent_trial_handler(self.config.agent_settings.kill_chain.EXPLOIT.probability):
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
self.chosen_action = "do-nothing", {}
|
||||
return
|
||||
self.current_kill_chain_stage = KillChainStageProgress.IN_PROGRESS
|
||||
|
||||
self.config.agent_settings.kill_chain.EXPLOIT.malicious_acls = (
|
||||
self.config.agent_settings.kill_chain.EXPLOIT.malicious_acls
|
||||
)
|
||||
self._num_acls = len(self.config.agent_settings.kill_chain.EXPLOIT.malicious_acls)
|
||||
malicious_acl = self.config.agent_settings.kill_chain.EXPLOIT.malicious_acls[self._current_acl]
|
||||
hostname = malicious_acl.target_router
|
||||
|
||||
if self.network_knowledge.get("current_session", {}).get("hostname") != hostname:
|
||||
self._ssh_target_host = hostname
|
||||
self.chosen_action = "node-session-remote-login", {
|
||||
"node_name": self.starting_node,
|
||||
"username": self.network_knowledge["credentials"][hostname]["username"],
|
||||
"password": self.network_knowledge["credentials"][hostname]["password"],
|
||||
"remote_ip": self.network_knowledge["credentials"][hostname]["ip_address"],
|
||||
}
|
||||
self.logger.info(f"Logging into {hostname} in order to add ACL rules.")
|
||||
# once we know we are logged in, send a command to change password
|
||||
else:
|
||||
self.chosen_action = "node-send-remote-command", {
|
||||
"node_name": self.starting_node,
|
||||
"remote_ip": self.network_knowledge["credentials"][hostname]["ip_address"],
|
||||
"command": [
|
||||
"acl",
|
||||
"add_rule",
|
||||
malicious_acl.permission,
|
||||
malicious_acl.protocol_name,
|
||||
str(malicious_acl.src_ip),
|
||||
str(malicious_acl.src_wildcard),
|
||||
malicious_acl.src_port,
|
||||
str(malicious_acl.dst_ip),
|
||||
str(malicious_acl.dst_wildcard),
|
||||
malicious_acl.dst_port,
|
||||
malicious_acl.position,
|
||||
],
|
||||
}
|
||||
self.logger.info(f"Adding ACL rule to {hostname}")
|
||||
self._current_acl = self._current_acl + 1
|
||||
|
||||
if self._current_acl == self._num_acls:
|
||||
self._current_acl = 0
|
||||
self.logger.info("Finished adding ACL rules.")
|
||||
self._progress_kill_chain()
|
||||
@@ -1,6 +1,21 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
|
||||
from primaite.game.agent import interface
|
||||
from primaite.game.agent.scripted_agents import abstract_tap, data_manipulation_bot, probabilistic_agent, random_agent
|
||||
from primaite.game.agent.scripted_agents import (
|
||||
abstract_tap,
|
||||
data_manipulation_bot,
|
||||
probabilistic_agent,
|
||||
random_agent,
|
||||
TAP001,
|
||||
TAP003,
|
||||
)
|
||||
|
||||
__all__ = ("abstract_tap", "data_manipulation_bot", "interface", "probabilistic_agent", "random_agent")
|
||||
__all__ = (
|
||||
"abstract_tap",
|
||||
"data_manipulation_bot",
|
||||
"interface",
|
||||
"probabilistic_agent",
|
||||
"random_agent",
|
||||
"TAP001",
|
||||
"TAP003",
|
||||
)
|
||||
|
||||
@@ -1,49 +1,202 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from abc import abstractmethod
|
||||
from enum import Enum, IntEnum
|
||||
from typing import Dict, List, Optional, Tuple, Type
|
||||
|
||||
from gymnasium.core import ObsType
|
||||
from pydantic import Field
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
|
||||
|
||||
__all__ = "AbstractTAPAgent"
|
||||
from primaite.game.agent.interface import AbstractScriptedAgent
|
||||
from primaite.game.science import simulate_trial
|
||||
|
||||
|
||||
class AbstractTAPAgent(PeriodicAgent, ABC):
|
||||
"""Base class for TAP agents to inherit from."""
|
||||
# This class is required for abstract tap. The IntEnums in this class are repeated in other kill chains.
|
||||
class BaseKillChain(IntEnum):
|
||||
"""A generic kill chain for abstract tap initialisation.
|
||||
|
||||
config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema())
|
||||
next_execution_timestep: int = 0
|
||||
The IntEnums in this class are repeated in other kill chains
|
||||
As IntEnums cannot be directly extended by a inheritance.
|
||||
"""
|
||||
|
||||
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema, ABC):
|
||||
"""Schema for the `agent_settings` part of the agent config."""
|
||||
NOT_STARTED = 100
|
||||
"Indicates that the Kill Chain has not started."
|
||||
SUCCEEDED = 200
|
||||
"Indicates that the kill chain has succeeded."
|
||||
FAILED = 300
|
||||
"Indicates that the attack has failed."
|
||||
|
||||
possible_starting_nodes: List[str] = Field(default_factory=list)
|
||||
# The original approach is to extend the base class during runtime via class methods.
|
||||
# However, this approach drastically impacted the readability and complexity of the code
|
||||
# So the decision was made to ignore the DRY Principle for kill chains.
|
||||
|
||||
class ConfigSchema(PeriodicAgent.ConfigSchema, ABC):
|
||||
"""Configuration schema for Abstract TAP agents."""
|
||||
@abstractmethod
|
||||
def initial_stage(self) -> "BaseKillChain":
|
||||
"""Returns the first stage in the kill chain. Used for Abstract TAP Setup."""
|
||||
return self.NOT_STARTED
|
||||
|
||||
type: str = "abstract-tap"
|
||||
agent_settings: AbstractTAPAgent.AgentSettingsSchema = Field(
|
||||
default_factory=lambda: AbstractTAPAgent.AgentSettingsSchema()
|
||||
|
||||
class KillChainStageProgress(Enum):
|
||||
"""Generic Progress Enums. Used by TAP Agents to keep track of kill chain stages that required multiple actions."""
|
||||
|
||||
PENDING = 0
|
||||
"""Indicates that the current kill chain stage is yet to start."""
|
||||
IN_PROGRESS = 1
|
||||
"""Indicates that the current kill chain stage is not yet completed."""
|
||||
FINISHED = 2
|
||||
"""Indicates that the current kill chain stage stage has been completed."""
|
||||
|
||||
|
||||
class KillChainOptions(BaseModel):
|
||||
"""Base Class for Kill Chain Options. Inherited by all TAP Type Agents."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class KillChainStageOptions(BaseModel):
|
||||
"""Shared options for generic Kill Chain Stages."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
probability: float = 1
|
||||
|
||||
|
||||
class AbstractTAP(AbstractScriptedAgent):
|
||||
"""Abstract class for Threat Actor Persona (TAP) Type Agents must inherit from.
|
||||
|
||||
This abstract base class provides TAP agents an interface which provides
|
||||
TAP type agents the necessary methods to execute kill chain(s) with
|
||||
configurable parameters.
|
||||
|
||||
TAP Actions are returned to the Request Manager as a Tuple
|
||||
in CAOS format via the get_action method in line with other agents.
|
||||
|
||||
Abstract TAP Class intends to provide each TAP the following:
|
||||
|
||||
1. Kill Chain Progression
|
||||
|
||||
Kill Chains are IntEnums which define the different stages within a kill chain.
|
||||
These stages are intended to be used across multiple ARCD environments.
|
||||
|
||||
2. Abstract Methods For Kill Chain Control Flow
|
||||
|
||||
Abstract methods _progress_kill_chain & _setup_kill_chain
|
||||
are intended to provide TAP type agent additional control
|
||||
over execution flow in comparison to AbstractScriptedAgent.
|
||||
|
||||
Usually these methods handle kill chain progression & success criteria.
|
||||
|
||||
For more information about Abstract TAPs please refer
|
||||
to the methods & attributes documentation directly.
|
||||
|
||||
Additionally, Refer to a specific TAP for a more specific example.
|
||||
"""
|
||||
|
||||
class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema):
|
||||
"""Agent Settings Schema. Default settings applied for all threat actor profiles."""
|
||||
|
||||
start_step: int = 5
|
||||
frequency: int = 5
|
||||
variance: int = 0
|
||||
repeat_kill_chain: bool = False
|
||||
repeat_kill_chain_stages: bool = True
|
||||
starting_nodes: Optional[List[str]] = []
|
||||
default_starting_node: str
|
||||
kill_chain: KillChainOptions
|
||||
|
||||
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
|
||||
"""Configuration schema applicable to all TAP agents."""
|
||||
|
||||
agent_settings: "AbstractTAP.AgentSettingsSchema" = Field(
|
||||
default_factory=lambda: AbstractTAP.AgentSettingsSchema()
|
||||
)
|
||||
|
||||
starting_node: Optional[str] = None
|
||||
config: ConfigSchema = Field(default_factory=lambda: AbstractTAP.ConfigSchema())
|
||||
|
||||
selected_kill_chain: Type[BaseKillChain]
|
||||
"""A combination of TAP's base & default kill chain. Loaded dynamically during agent setup."""
|
||||
next_execution_timestep: int = 0
|
||||
"""The next timestep in which the agent will attempt to progress the kill chain."""
|
||||
starting_node: str = ""
|
||||
"""The name (string) of TAP agent's starting node. This attribute is initialised via _self_select_starting_node."""
|
||||
|
||||
actions_concluded: bool = False
|
||||
"""Boolean value which indicates if a TAP Agent has completed it's attack for the episode."""
|
||||
|
||||
next_kill_chain_stage: BaseKillChain = BaseKillChain.NOT_STARTED
|
||||
"""The IntEnum of the next kill chain stage to be executed.
|
||||
|
||||
This attribute is initialised via _tap_start.
|
||||
Afterwards, this attribute is loaded dynamically via _progress_kill_chain.
|
||||
"""
|
||||
current_kill_chain_stage: BaseKillChain = BaseKillChain.NOT_STARTED
|
||||
"""The TAP agent's current kill chain.
|
||||
|
||||
This attribute is used as a state to indicate the current progress in a kill chain.
|
||||
"""
|
||||
current_stage_progress: KillChainStageProgress = KillChainStageProgress.PENDING
|
||||
"""The TAP agent's current progress in a stage within a kill chain.
|
||||
|
||||
This attribute is used as a state to indicate the current progress in a individual kill chain stage.
|
||||
|
||||
Some TAP's require multiple actions to take place before moving onto the next stage in a kill chain.
|
||||
This attribute is used to keep track of the current progress within an individual stage.
|
||||
"""
|
||||
chosen_action: Tuple[str, Dict] = "do-nothing", {}
|
||||
"""The next agent's chosen action. Returned in CAOS format at the end of each timestep."""
|
||||
|
||||
current_host: str = ""
|
||||
"""The name (str) of a TAP agent's currently selected host.
|
||||
|
||||
This attribute is set dynamically during tap execution via _set_current_host.
|
||||
"""
|
||||
current_timestep: int = 0
|
||||
"""The current timestep (int) of the game.
|
||||
|
||||
This attribute is set to the "timestep" argument passed to get_action.
|
||||
Mainly used to by kill chain stages for complex execution flow that is dependant on the simulation.
|
||||
|
||||
Specifically, this attribute is used for indexing previous actions in the self.history inherited attribute.
|
||||
|
||||
For more information please refer to AbstractAgent's "self.history" attribute
|
||||
And for action responses see 'request.py' and the .data attribute.
|
||||
|
||||
Lastly, a demonstration of the above capability can be found in the PROPAGATE step in the tap001-e2e notebook.
|
||||
"""
|
||||
|
||||
def update_current_timestep(self, new_timestep: int):
|
||||
"""Updates the current time_step attribute to the given timestep argument."""
|
||||
self.current_timestep = new_timestep
|
||||
|
||||
@abstractmethod
|
||||
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
|
||||
"""Return an action to be taken in the environment."""
|
||||
return super().get_action(obs=obs, timestep=timestep)
|
||||
def _progress_kill_chain(self):
|
||||
"""Private Abstract method which defines the default kill chain progression.
|
||||
|
||||
@abstractmethod
|
||||
def setup_agent(self) -> None:
|
||||
"""Set up agent."""
|
||||
This abstract method intend to allow TAPs to control the logic flow of their kill chain.
|
||||
In a majority of cases this method handles the success criteria and incrementing the current kill chain intenum.
|
||||
|
||||
This method is abstract so TAPs can configure this behaviour for tap specific implementations.
|
||||
"""
|
||||
pass
|
||||
|
||||
def _select_start_node(self) -> None:
|
||||
"""
|
||||
Handles setting the starting node behaviour of TAP type agents.
|
||||
|
||||
If the user given tap_settings provides a starting_node list then the starting node
|
||||
is set to a random node given in the starting_node list.
|
||||
Otherwise, the starting node is set to the 'default_starting_node' option.
|
||||
"""
|
||||
# Catches empty starting nodes.
|
||||
if not self.config.agent_settings.starting_nodes:
|
||||
self.starting_node = self.config.agent_settings.default_starting_node
|
||||
else:
|
||||
self.starting_node = random.choice(self.config.agent_settings.starting_nodes)
|
||||
|
||||
def _setup_agent_kill_chain(self, given_kill_chain: BaseKillChain) -> None:
|
||||
"""Sets the 'next_kill_chain_stage' TAP attribute via the public kill chain method 'initial_stage'."""
|
||||
self.selected_kill_chain = given_kill_chain
|
||||
self.next_kill_chain_stage = self.selected_kill_chain.initial_stage(given_kill_chain)
|
||||
|
||||
def _set_next_execution_timestep(self, timestep: int) -> None:
|
||||
"""Set the next execution timestep with a configured random variance.
|
||||
|
||||
@@ -54,8 +207,96 @@ class AbstractTAPAgent(PeriodicAgent, ABC):
|
||||
)
|
||||
self.next_execution_timestep = timestep + random_timestep_increment
|
||||
|
||||
def _select_start_node(self) -> None:
|
||||
"""Set the starting starting node of the agent to be a random node from this agent's action manager."""
|
||||
# we are assuming that every node in the node manager has a data manipulation application at idx 0
|
||||
self.starting_node = random.choice(self.config.agent_settings.possible_starting_nodes)
|
||||
self.logger.debug(f"Selected starting node: {self.starting_node}")
|
||||
def _agent_trial_handler(self, agent_probability_of_success: int) -> bool:
|
||||
"""Acts as a wrapper around simulate trial - Sets kill chain stage to failed if the relevant setting is set.
|
||||
|
||||
:param agent_probability_of_success: The probability of the action success to be passed to simulate_trial.
|
||||
:type agent_probability_of_success: int.
|
||||
:rtype: Bool.
|
||||
"""
|
||||
if simulate_trial(agent_probability_of_success):
|
||||
return True
|
||||
else:
|
||||
self.logger.info(
|
||||
f"failed to reach kill chain stage {self.next_kill_chain_stage.name} due to probability of failure."
|
||||
)
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.logger.info(f"Thus {self.config.ref} has failed the kill chain")
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
return False
|
||||
else:
|
||||
self.logger.info(f"Retrying from stage {self.current_kill_chain_stage.name}.")
|
||||
return False
|
||||
|
||||
def _tap_outcome_handler(self, selected_kill_chain_class: BaseKillChain) -> None:
|
||||
"""
|
||||
Default TAP behaviour for base kill chain stages.
|
||||
|
||||
Upon Success and failure:
|
||||
TAPs will either repeat or re-attack dependant on the user given settings.
|
||||
|
||||
:param tap_kill_chain: The TAP agent's currently selected kill chain.
|
||||
:type tap_kill_chain: BaseKillChain
|
||||
"""
|
||||
if (
|
||||
self.current_kill_chain_stage == self.selected_kill_chain.SUCCEEDED
|
||||
or self.current_kill_chain_stage == self.selected_kill_chain.FAILED
|
||||
):
|
||||
if self.actions_concluded == True: # Prevents Further logging via a guard clause boolean
|
||||
self.chosen_action = "do-nothing", {}
|
||||
return
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.SUCCEEDED:
|
||||
self.logger.info(f"{self.config.ref} has successfully carried out the kill chain.")
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.FAILED:
|
||||
self.logger.info(f"{self.config.ref} has failed the Kill Chain.")
|
||||
if self.config.agent_settings.repeat_kill_chain == True:
|
||||
self.logger.info(f"{self.config.ref} has opted to re-attack!")
|
||||
self.current_kill_chain_stage = BaseKillChain.NOT_STARTED
|
||||
self.next_kill_chain_stage = selected_kill_chain_class.initial_stage(selected_kill_chain_class)
|
||||
else:
|
||||
self.logger.info(f"{self.config.ref} has opted to forgo any further attacks.")
|
||||
self.actions_concluded = True # Guard Clause Bool
|
||||
self.chosen_action = "do-nothing", {}
|
||||
|
||||
def _tap_return_handler(self, timestep: int) -> bool:
|
||||
# Intelligence | Use the request_response system to enable different behaviour
|
||||
"""
|
||||
Handles the request_manager's response query. Sets Kill Chain to false if failed.
|
||||
|
||||
If the previous action failed due to the simulation state,
|
||||
the kill chain is considered to have failed.
|
||||
|
||||
Returns True if the previous action was successful.
|
||||
Returns False if the previous action was any other state.
|
||||
(Including Pending and Failure)
|
||||
|
||||
:param timestep: The current primAITE game layer timestep.
|
||||
:type timestep: int
|
||||
:rtype bool
|
||||
"""
|
||||
if self.history[timestep].response.status != "success":
|
||||
self.logger.info(
|
||||
f"{self.config.ref} has failed to successfully carry out {self.current_kill_chain_stage.name}"
|
||||
)
|
||||
self.logger.info(f"due to the simulation state: {self.history[timestep].response.data}")
|
||||
if self.config.agent_settings.repeat_kill_chain_stages == False:
|
||||
self.logger.info(
|
||||
f"Thus {self.config.ref} has failed this kill chain attempt on {self.current_kill_chain_stage.name}"
|
||||
)
|
||||
self.current_kill_chain_stage = self.selected_kill_chain.FAILED
|
||||
else:
|
||||
self.logger.info(f"Retrying from stage {self.current_kill_chain_stage.name}!")
|
||||
return False
|
||||
return True
|
||||
|
||||
def _tap_start(self, tap_kill_chain: Type[BaseKillChain]) -> None:
|
||||
"""
|
||||
Sets the TAP Agent's beginning current/next kill chain stages.
|
||||
|
||||
:param IntEnum tap_kill_chain: A currently selected kill chain.
|
||||
"""
|
||||
if self.current_kill_chain_stage == self.selected_kill_chain.NOT_STARTED:
|
||||
self.current_kill_chain_stage = tap_kill_chain.initial_stage(tap_kill_chain)
|
||||
self.next_kill_chain_stage = self.selected_kill_chain(self.current_kill_chain_stage + 1)
|
||||
self.logger.info(f"{self.config.ref} has begun it's attack!")
|
||||
self.chosen_action = "do-nothing", {}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -700,7 +700,7 @@
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "venv2",
|
||||
"display_name": ".venv",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
@@ -714,7 +714,7 @@
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.11.10"
|
||||
"version": "3.10.12"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
|
||||
@@ -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
1844
src/primaite/notebooks/UC7-TAP001-Kill-Chain-E2E.ipynb
Normal file
1714
src/primaite/notebooks/UC7-TAP003-Kill-Chain-E2E.ipynb
Normal file
147
src/primaite/notebooks/UC7-Training.ipynb
Normal 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
@@ -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
|
After Width: | Height: | Size: 1.4 MiB |
BIN
src/primaite/notebooks/_package_data/uc7/uc7_network.png
Normal file
|
After Width: | Height: | Size: 159 KiB |
|
After Width: | Height: | Size: 609 KiB |
|
After Width: | Height: | Size: 383 KiB |
|
After Width: | Height: | Size: 1.0 MiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 1.0 MiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.0 MiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.3 MiB |
@@ -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:
|
||||
|
||||
172
tests/e2e_integration_tests/test_uc7_agents.py
Normal file
@@ -0,0 +1,172 @@
|
||||
# © 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"
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
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"][ATTACK_AGENT_INDEX]["agent_settings"]["starting_nodes"] = ["ST_PROJ-A-PRV-PC-1"]
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["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
@@ -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)
|
||||
@@ -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)
|
||||
@@ -0,0 +1 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
@@ -0,0 +1,142 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
from primaite.session.environment import PrimaiteGymEnv
|
||||
|
||||
START_STEP = 1 # The starting step of the agent.
|
||||
FREQUENCY = 5 # 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)
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def uc7_tap001_env() -> PrimaiteGymEnv:
|
||||
with open(_EXAMPLE_CFG / "uc7_config.yaml", mode="r") as uc7_config:
|
||||
cfg = yaml.safe_load(uc7_config)
|
||||
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
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
if "repeat_kill_chain" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain"] = kwargs["repeat_kill_chain"]
|
||||
if "repeat_kill_chain_stages" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = kwargs[
|
||||
"repeat_kill_chain_stages"
|
||||
]
|
||||
if "planning_probability" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = kwargs[
|
||||
"planning_probability"
|
||||
]
|
||||
if "custom_kill_chain" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"] = kwargs["custom_kill_chain"]
|
||||
if "starting_nodes" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["starting_nodes"] = kwargs["starting_nodes"]
|
||||
if "target_nodes" in kwargs:
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["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
|
||||
# 6 Iterations as the attacker frequency is set to 5. Therefore, after 6 timesteps we expect the agent to realise it's failed the kill chain.
|
||||
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
|
||||
@@ -0,0 +1,72 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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)
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def uc7_tap003_env() -> PrimaiteGymEnv:
|
||||
"""Setups the UC7 TAP003 Game with a 1 timestep start_step, frequency of 2 and 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
|
||||
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"][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
|
||||
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
|
||||
kill_chain_enums = [enums for enums in tap.selected_kill_chain]
|
||||
assert BaseKillChain.FAILED in kill_chain_enums
|
||||
assert BaseKillChain.SUCCEEDED in kill_chain_enums
|
||||
assert BaseKillChain.NOT_STARTED in kill_chain_enums
|
||||
# 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.
|
||||
kill_chain_enums = [enums for enums in tap.selected_kill_chain]
|
||||
assert BaseKillChain.FAILED in kill_chain_enums
|
||||
assert BaseKillChain.SUCCEEDED in kill_chain_enums
|
||||
assert BaseKillChain.NOT_STARTED in kill_chain_enums
|
||||
# Asserting that the Insider Kill Chain intEnum stages are loaded.
|
||||
assert len(MobileMalwareKillChain.__members__) == len(tap.selected_kill_chain.__members__)
|
||||
@@ -0,0 +1,116 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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)
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def uc7_tap001_env(**kwargs) -> PrimaiteGymEnv:
|
||||
"""Setups the UC7 tap001 game with a 1 timestep start_step, frequency of 2 and 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"][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"]["PROPAGATE"]["probability"] = kwargs[
|
||||
"propagate_probability"
|
||||
]
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["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 25 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 frequency of 2 timesteps
|
||||
# Thus the for loop above should never complete it's full 50 iterations.
|
||||
# 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 four timesteps (2 TAP001 steps) 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
|
||||
@@ -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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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)
|
||||
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.
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
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"][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"] = REPEAT_KILL_CHAIN_STAGES
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PAYLOAD"][
|
||||
"probability"
|
||||
] = KILL_CHAIN_PROBABILITY
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PROPAGATE"][
|
||||
"probability"
|
||||
] = KILL_CHAIN_PROBABILITY
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["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 so we need to step through the environment a couple of times
|
||||
# In order for TAP001 to move onto the next kill chain stage.
|
||||
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")
|
||||
@@ -0,0 +1,141 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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's.
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
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"][ATTACK_AGENT_INDEX]["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 fail because the maximum amount of scan attempts has been reached
|
||||
assert tap001.scans_complete == scan_attempts
|
||||
assert tap001.current_kill_chain_stage == MobileMalwareKillChain.FAILED
|
||||
@@ -0,0 +1,108 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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)
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def uc7_tap003_env(**kwargs) -> PrimaiteGymEnv:
|
||||
"""Setups the UC7 TAP003 Game with a 1 timestep start_step, frequency of 2 and 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"
|
||||
]
|
||||
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,
|
||||
# access_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
|
||||
@@ -0,0 +1,218 @@
|
||||
# © 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.TAP003 import InsiderKillChain, TAP003
|
||||
from primaite.session.environment import PrimaiteGymEnv
|
||||
from primaite.simulator.network.hardware.nodes.network.firewall import Firewall
|
||||
from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
|
||||
|
||||
# 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'
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def uc7_tap003_env() -> PrimaiteGymEnv:
|
||||
"""Setups the UC7 TAP003 Game with a 1 timestep start_step, frequency of 2 and 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["io_settings"]["save_agent_logs"] = True
|
||||
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"] = REPEAT_KILL_CHAIN_STAGES
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["repeat_kill_chain_stages"] = REPEAT_KILL_CHAIN_STAGES
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"][
|
||||
"probability"
|
||||
] = KILL_CHAIN_PROBABILITY
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"][
|
||||
"probability"
|
||||
] = KILL_CHAIN_PROBABILITY
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"][
|
||||
"probability"
|
||||
] = KILL_CHAIN_PROBABILITY
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["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 successful - TAP003 has loaded it's starting network knowledge into it's network knowledge.
|
||||
|
||||
# At this point TAP003 will parse it's starting network knowledge config into it's a private attribute (`network_knowledge`)
|
||||
assert (
|
||||
tap003.network_knowledge["credentials"]
|
||||
== tap003.config.agent_settings.kill_chain.PLANNING.starting_network_knowledge["credentials"]
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def test_tap003_kill_chain_stage_manipulation():
|
||||
"""Tests the successful/failed handlers in the manipulation 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)
|
||||
|
||||
assert tap003.current_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
|
||||
|
||||
# Testing that the stage successfully impacted the simulation - Accounts Altered
|
||||
env = environment_step(i=5, env=env)
|
||||
st_intra_prv_rt_dr_1: Router = env.game.simulation.network.get_node_by_hostname("ST_INTRA-PRV-RT-DR-1")
|
||||
assert tap003.current_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
|
||||
assert st_intra_prv_rt_dr_1.user_manager.admins["admin"].password == "red_pass"
|
||||
|
||||
env = environment_step(i=5, env=env)
|
||||
st_intra_prv_rt_cr: Router = env.game.simulation.network.get_node_by_hostname("ST_INTRA-PRV-RT-CR")
|
||||
assert tap003.current_kill_chain_stage.name == InsiderKillChain.MANIPULATION.name
|
||||
assert st_intra_prv_rt_cr.user_manager.admins["admin"].password == "red_pass"
|
||||
|
||||
env = environment_step(i=5, env=env)
|
||||
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 = 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=16, env=env)
|
||||
|
||||
assert tap003.current_kill_chain_stage.name == InsiderKillChain.EXPLOIT.name
|
||||
|
||||
# Testing that the stage successfully impacted the simulation - Malicious ACL Added:
|
||||
for _ in range(14):
|
||||
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
|
||||
@@ -0,0 +1,203 @@
|
||||
# © 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, TAP001
|
||||
from primaite.game.agent.scripted_agents.TAP003 import InsiderKillChain, TAP003
|
||||
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 a 1 timestep start_step, frequency of 2 and 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():
|
||||
# 120 environment steps to ensure that TAP003 reaches manipulate.
|
||||
# If this loop finishes 120 iterations before the test finishes then TAP003 is struggling to
|
||||
# reach or finish the manipulation kill chain stage correctly.
|
||||
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 :)
|
||||
@@ -0,0 +1 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
@@ -0,0 +1,25 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import yaml
|
||||
|
||||
from primaite.config.load import _EXAMPLE_CFG
|
||||
from primaite.game.agent.scripted_agents.TAP003 import TAP003
|
||||
from primaite.game.game import PrimaiteGame
|
||||
|
||||
ATTACK_AGENT_INDEX = 32
|
||||
|
||||
|
||||
def test_tap003_kill_chain_settings_load_config():
|
||||
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
|
||||
cfg = yaml.safe_load(uc7_config)
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["MANIPULATION"]["probability"] = 0.5
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["ACCESS"]["probability"] = 0.5
|
||||
cfg["agents"][ATTACK_AGENT_INDEX]["agent_settings"]["kill_chain"]["PLANNING"]["probability"] = 0.5
|
||||
game = PrimaiteGame.from_config(cfg)
|
||||
tap: TAP003 = game.agents["attacker"]
|
||||
kill_chain = tap.config.agent_settings.kill_chain
|
||||
assert kill_chain.MANIPULATION.probability == 0.5
|
||||
assert kill_chain.ACCESS.probability == 0.5
|
||||
assert kill_chain.PLANNING.probability == 0.5
|
||||
@@ -0,0 +1,35 @@
|
||||
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
|
||||
from pathlib import Path
|
||||
from typing import Union
|
||||
|
||||
import yaml
|
||||
|
||||
from primaite.config.load import _EXAMPLE_CFG
|
||||
from primaite.game.agent.scripted_agents.TAP003 import TAP003
|
||||
from primaite.game.game import PrimaiteGame
|
||||
|
||||
|
||||
def test_threat_actor_profile_load_config():
|
||||
"""Test to check that threat actor profiles are able to be loaded."""
|
||||
with open(_EXAMPLE_CFG / "uc7_config_tap003.yaml", mode="r") as uc7_config:
|
||||
cfg = yaml.safe_load(uc7_config)
|
||||
|
||||
game = PrimaiteGame.from_config(cfg)
|
||||
# tap003 is found and loaded TODO: Once tuple digestion is implemented, change to hardcoded 'tap003' test.
|
||||
assert "attacker" in game.agents
|
||||
assert isinstance(game.agents["attacker"], TAP003)
|
||||
agent: TAP003 = game.agents["attacker"]
|
||||
assert agent.config.agent_settings.start_step == 1
|
||||
assert agent.config.agent_settings.frequency == 3
|
||||
assert agent.config.agent_settings.variance == 0
|
||||
assert not agent.config.agent_settings.repeat_kill_chain
|
||||
assert agent.config.agent_settings.repeat_kill_chain_stages
|
||||
assert agent.config.agent_settings.default_starting_node == "ST_PROJ-A-PRV-PC-1"
|
||||
assert not agent.config.agent_settings.starting_nodes
|
||||
assert agent.config.agent_settings.kill_chain.PLANNING.probability == 1
|
||||
assert len(agent.config.agent_settings.kill_chain.PLANNING.starting_network_knowledge["credentials"]) == 6
|
||||
assert agent.config.agent_settings.kill_chain.ACCESS.probability == 1
|
||||
assert agent.config.agent_settings.kill_chain.MANIPULATION.probability == 1
|
||||
assert len(agent.config.agent_settings.kill_chain.MANIPULATION.account_changes) == 3
|
||||
assert agent.config.agent_settings.kill_chain.EXPLOIT.probability == 1
|
||||
assert len(agent.config.agent_settings.kill_chain.EXPLOIT.malicious_acls) == 3
|
||||
@@ -47,7 +47,7 @@ def test_acl_observations(simulation):
|
||||
observation_space = acl_obs.observe(simulation.describe_state())
|
||||
assert observation_space.get(1) is not None
|
||||
rule_obs = observation_space.get(1) # this is the ACL Rule added to allow NTP
|
||||
assert rule_obs.get("position") == 0 # rule was put at position 1 (0 because counting from 1 instead of 1)
|
||||
assert rule_obs.get("position") == 1 # rule was put at position 1
|
||||
assert rule_obs.get("permission") == 1 # permit = 1 deny = 2
|
||||
assert rule_obs.get("source_ip_id") == 1 # applies to all source nodes
|
||||
assert rule_obs.get("dest_ip_id") == 1 # applies to all destination nodes
|
||||
@@ -60,7 +60,7 @@ def test_acl_observations(simulation):
|
||||
observation_space = acl_obs.observe(simulation.describe_state())
|
||||
assert observation_space.get(1) is not None
|
||||
rule_obs = observation_space.get(1) # this is the ACL Rule added to allow NTP
|
||||
assert rule_obs.get("position") == 0
|
||||
assert rule_obs.get("position") == 1
|
||||
assert rule_obs.get("permission") == 0
|
||||
assert rule_obs.get("source_ip_id") == 0
|
||||
assert rule_obs.get("dest_ip_id") == 0
|
||||
|
||||
@@ -11,15 +11,15 @@ from primaite.utils.validation.port import PORT_LOOKUP
|
||||
|
||||
def check_default_rules(acl_obs):
|
||||
assert len(acl_obs) == 7
|
||||
assert all(acl_obs[i]["position"] == i - 1 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["permission"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_ip_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_wildcard_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["source_port_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_ip_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_wildcard_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["dest_port_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["protocol_id"] == 0 for i in range(1, 8))
|
||||
assert all(acl_obs[i]["position"] == i for i in range(7))
|
||||
assert all(acl_obs[i]["permission"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["source_ip_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["source_wildcard_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["source_port_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["dest_ip_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["dest_wildcard_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["dest_port_id"] == 0 for i in range(7))
|
||||
assert all(acl_obs[i]["protocol_id"] == 0 for i in range(7))
|
||||
|
||||
|
||||
def test_firewall_observation():
|
||||
@@ -75,7 +75,7 @@ def test_firewall_observation():
|
||||
|
||||
observation = firewall_observation.observe(firewall.describe_state())
|
||||
observed_rule = observation["ACL"]["INTERNAL"]["INBOUND"][5]
|
||||
assert observed_rule["position"] == 4
|
||||
assert observed_rule["position"] == 5
|
||||
assert observed_rule["permission"] == 2
|
||||
assert observed_rule["source_ip_id"] == 2
|
||||
assert observed_rule["source_wildcard_id"] == 3
|
||||
|
||||
@@ -53,7 +53,7 @@ def test_router_observation():
|
||||
# Observe the state using the RouterObservation instance
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
observed_rule = observed_output["ACL"][5]
|
||||
assert observed_rule["position"] == 4
|
||||
assert observed_rule["position"] == 5
|
||||
assert observed_rule["permission"] == 2
|
||||
assert observed_rule["source_ip_id"] == 2
|
||||
assert observed_rule["source_wildcard_id"] == 3
|
||||
@@ -77,7 +77,7 @@ def test_router_observation():
|
||||
)
|
||||
observed_output = router_observation.observe(router.describe_state())
|
||||
observed_rule = observed_output["ACL"][2]
|
||||
assert observed_rule["position"] == 1
|
||||
assert observed_rule["position"] == 2
|
||||
assert observed_rule["permission"] == 1
|
||||
assert observed_rule["source_ip_id"] == 1
|
||||
assert observed_rule["source_wildcard_id"] == 1
|
||||
|
||||