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

Merge 'origin/dev' into feature/2319-database-admin

Browse Source
This commit is contained in:
Marek Wolan
2024-02-29 13:21:44 +00:00
parent af8ca82fcb 312b5c35bd
commit f7c535b557
62 changed files with 865 additions and 1163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/primaite/config/_package_data/example_config.yaml
View File

@@ -14,7 +14,7 @@ io_settings:
save_checkpoints: true
checkpoint_interval: 5
save_step_metadata: false
save_pcap_logs: true
save_pcap_logs: false
save_sys_logs: true
@@ -627,6 +627,10 @@ agents:
simulation:
network:
nmne_config:
capture_nmne: true
nmne_capture_keywords:
- DELETE
nodes:
- ref: router_1
@@ -696,12 +700,13 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server

458
src/primaite/config/_package_data/example_config_2_rl_agents.yaml
View File

@@ -1,15 +1,23 @@
training_config:
rl_framework: RLLIB_multi_agent
# rl_framework: SB3
rl_algorithm: PPO
seed: 333
n_learn_episodes: 1
n_eval_episodes: 5
max_steps_per_episode: 256
deterministic_eval: false
n_agents: 2
agent_references:
- defender_1
- defender_2
io_settings:
save_checkpoints: true
checkpoint_interval: 5
save_step_metadata: false
save_pcap_logs: false
save_sys_logs: true
game:
@@ -36,9 +44,9 @@ agents:
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_ref: client_2
- node_name: client_2
applications:
- application_ref: client_2_web_browser
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
@@ -54,6 +62,31 @@ agents:
frequency: 4
variance: 3
- ref: client_1_green_user
team: GREEN
type: GreenWebBrowsingAgent
observation_space:
type: UC2GreenObservation
action_space:
action_list:
- type: DONOTHING
- type: NODE_APPLICATION_EXECUTE
options:
nodes:
- node_name: client_1
applications:
- application_name: WebBrowser
max_folders_per_node: 1
max_files_per_folder: 1
max_services_per_node: 1
max_applications_per_node: 1
reward_function:
reward_components:
- type: DUMMY
- ref: data_manipulation_attacker
team: RED
type: RedDatabaseCorruptingAgent
@@ -72,7 +105,7 @@ agents:
- type: NODE_OS_SCAN
options:
nodes:
- node_ref: client_1
- node_name: client_1
applications:
- application_name: DataManipulationBot
- node_name: client_2
@@ -104,25 +137,21 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: DNSServer
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
services:
- service_ref: database_service
- service_name: WebServer
- node_hostname: database_server
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: backup_server
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -137,23 +166,23 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
- node_hostname: domain_controller
nic_num: 1
- node_ref: web_server
- node_hostname: web_server
nic_num: 1
- node_ref: database_server
- node_hostname: database_server
nic_num: 1
- node_ref: backup_server
- node_hostname: backup_server
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 1
- node_ref: client_1
- node_hostname: client_1
nic_num: 1
- node_ref: client_2
- node_hostname: client_2
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 2
ics: null
@@ -184,10 +213,10 @@ agents:
- type: NODE_RESET
- type: NETWORK_ACL_ADDRULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_ACL_REMOVERULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_NIC_ENABLE
- type: NETWORK_NIC_DISABLE
@@ -242,25 +271,25 @@ agents:
action: "NODE_FILE_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
10:
action: "NODE_FILE_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
11:
action: "NODE_FILE_DELETE"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
12:
action: "NODE_FILE_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
13:
action: "NODE_SERVICE_PATCH"
@@ -271,22 +300,22 @@ agents:
action: "NODE_FOLDER_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
15:
action: "NODE_FOLDER_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
16:
action: "NODE_FOLDER_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
17:
action: "NODE_FOLDER_RESTORE"
options:
node_id: 2
folder_id: 1
folder_id: 0
18:
action: "NODE_OS_SCAN"
options:
@@ -303,63 +332,63 @@ agents:
action: "NODE_RESET"
options:
node_id: 5
22:
22: # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
permission: 2
source_ip_id: 7
dest_ip_id: 1
source_ip_id: 7 # client 1
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
23:
23: # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 2
permission: 2
source_ip_id: 8
dest_ip_id: 1
source_ip_id: 8 # client 2
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
24:
24: # block tcp traffic from client 1 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 3
permission: 2
source_ip_id: 7
dest_ip_id: 3
source_ip_id: 7 # client 1
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
25:
25: # block tcp traffic from client 2 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 4
permission: 2
source_ip_id: 8
dest_ip_id: 3
source_ip_id: 8 # client 2
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
26:
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 5
permission: 2
source_ip_id: 7
dest_ip_id: 4
source_ip_id: 7 # client 1
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
27:
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 6
permission: 2
source_ip_id: 8
dest_ip_id: 4
source_ip_id: 8 # client 2
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
@@ -407,123 +436,148 @@ agents:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
39:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
40:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
41:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
42:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
43:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
44:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
45:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
46:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
47:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
48:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
49:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
50:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
51:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
52:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
53:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
options:
nodes:
- node_ref: domain_controller
- node_ref: web_server
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_ref: web_server_web_service
- node_ref: database_server
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_ref: database_service
- node_ref: backup_server
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_address_order:
- node_name: domain_controller
nic_num: 1
- node_name: web_server
nic_num: 1
- node_name: database_server
nic_num: 1
- node_name: backup_server
nic_num: 1
- node_name: security_suite
nic_num: 1
- node_name: client_1
nic_num: 1
- node_name: client_2
nic_num: 1
- node_name: security_suite
nic_num: 2
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
weight: 0.34
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
- type: WEB_SERVER_404_PENALTY
weight: 0.5
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: client_1
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_hostname: client_2
agent_settings:
# ...
flatten_obs: true
- ref: defender_2
team: BLUE
@@ -537,25 +591,21 @@ agents:
num_files_per_folder: 1
num_nics_per_node: 2
nodes:
- node_ref: domain_controller
- node_hostname: domain_controller
services:
- service_ref: domain_controller_dns_server
- node_ref: web_server
- service_name: DNSServer
- node_hostname: web_server
services:
- service_ref: web_server_database_client
- node_ref: database_server
services:
- service_ref: database_service
- service_name: WebServer
- node_hostname: database_server
folders:
- folder_name: database
files:
- file_name: database.db
- node_ref: backup_server
# services:
# - service_ref: backup_service
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- node_hostname: backup_server
- node_hostname: security_suite
- node_hostname: client_1
- node_hostname: client_2
links:
- link_ref: router_1___switch_1
- link_ref: router_1___switch_2
@@ -570,23 +620,23 @@ agents:
acl:
options:
max_acl_rules: 10
router_node_ref: router_1
router_hostname: router_1
ip_address_order:
- node_ref: domain_controller
- node_hostname: domain_controller
nic_num: 1
- node_ref: web_server
- node_hostname: web_server
nic_num: 1
- node_ref: database_server
- node_hostname: database_server
nic_num: 1
- node_ref: backup_server
- node_hostname: backup_server
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 1
- node_ref: client_1
- node_hostname: client_1
nic_num: 1
- node_ref: client_2
- node_hostname: client_2
nic_num: 1
- node_ref: security_suite
- node_hostname: security_suite
nic_num: 2
ics: null
@@ -617,10 +667,10 @@ agents:
- type: NODE_RESET
- type: NETWORK_ACL_ADDRULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_ACL_REMOVERULE
options:
target_router_ref: router_1
target_router_hostname: router_1
- type: NETWORK_NIC_ENABLE
- type: NETWORK_NIC_DISABLE
@@ -675,25 +725,25 @@ agents:
action: "NODE_FILE_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
10:
action: "NODE_FILE_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
11:
action: "NODE_FILE_DELETE"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
12:
action: "NODE_FILE_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
file_id: 0
13:
action: "NODE_SERVICE_PATCH"
@@ -704,22 +754,22 @@ agents:
action: "NODE_FOLDER_SCAN"
options:
node_id: 2
folder_id: 1
folder_id: 0
15:
action: "NODE_FOLDER_CHECKHASH"
options:
node_id: 2
folder_id: 1
folder_id: 0
16:
action: "NODE_FOLDER_REPAIR"
options:
node_id: 2
folder_id: 1
folder_id: 0
17:
action: "NODE_FOLDER_RESTORE"
options:
node_id: 2
folder_id: 1
folder_id: 0
18:
action: "NODE_OS_SCAN"
options:
@@ -736,63 +786,63 @@ agents:
action: "NODE_RESET"
options:
node_id: 5
22:
22: # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
permission: 2
source_ip_id: 7
dest_ip_id: 1
source_ip_id: 7 # client 1
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
23:
23: # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 2
permission: 2
source_ip_id: 8
dest_ip_id: 1
source_ip_id: 8 # client 2
dest_ip_id: 1 # ALL
source_port_id: 1
dest_port_id: 1
protocol_id: 1
24:
24: # block tcp traffic from client 1 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 3
permission: 2
source_ip_id: 7
dest_ip_id: 3
source_ip_id: 7 # client 1
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
25:
25: # block tcp traffic from client 2 to web app
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 4
permission: 2
source_ip_id: 8
dest_ip_id: 3
source_ip_id: 8 # client 2
dest_ip_id: 3 # web server
source_port_id: 1
dest_port_id: 1
protocol_id: 3
26:
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 5
permission: 2
source_ip_id: 7
dest_ip_id: 4
source_ip_id: 7 # client 1
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
27:
action: "NETWORK_ACL_ADDRULE"
options:
position: 1
position: 6
permission: 2
source_ip_id: 8
dest_ip_id: 4
source_ip_id: 8 # client 2
dest_ip_id: 4 # database
source_port_id: 1
dest_port_id: 1
protocol_id: 3
@@ -840,122 +890,148 @@ agents:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
39:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 0
nic_id: 1
nic_id: 0
40:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
41:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 1
nic_id: 1
nic_id: 0
42:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
43:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 2
nic_id: 1
nic_id: 0
44:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
45:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 3
nic_id: 1
nic_id: 0
46:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
47:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 1
nic_id: 0
48:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
49:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 4
nic_id: 2
nic_id: 1
50:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
51:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 5
nic_id: 1
nic_id: 0
52:
action: "NETWORK_NIC_DISABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
53:
action: "NETWORK_NIC_ENABLE"
options:
node_id: 6
nic_id: 1
nic_id: 0
options:
nodes:
- node_ref: domain_controller
- node_ref: web_server
- node_name: domain_controller
- node_name: web_server
applications:
- application_name: DatabaseClient
services:
- service_ref: web_server_web_service
- node_ref: database_server
- service_name: WebServer
- node_name: database_server
folders:
- folder_name: database
files:
- file_name: database.db
services:
- service_ref: database_service
- node_ref: backup_server
- node_ref: security_suite
- node_ref: client_1
- node_ref: client_2
- service_name: DatabaseService
- node_name: backup_server
- node_name: security_suite
- node_name: client_1
- node_name: client_2
max_folders_per_node: 2
max_files_per_folder: 2
max_services_per_node: 2
max_nics_per_node: 8
max_acl_rules: 10
ip_address_order:
- node_name: domain_controller
nic_num: 1
- node_name: web_server
nic_num: 1
- node_name: database_server
nic_num: 1
- node_name: backup_server
nic_num: 1
- node_name: security_suite
nic_num: 1
- node_name: client_1
nic_num: 1
- node_name: client_2
nic_num: 1
- node_name: security_suite
nic_num: 2
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
weight: 0.5
weight: 0.34
options:
node_ref: database_server
node_hostname: database_server
folder_name: database
file_name: database.db
- type: WEB_SERVER_404_PENALTY
weight: 0.5
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_ref: web_server
service_ref: web_server_web_service
node_hostname: client_1
- type: WEBPAGE_UNAVAILABLE_PENALTY
weight: 0.33
options:
node_hostname: client_2
agent_settings:
# ...
flatten_obs: true
@@ -963,6 +1039,10 @@ agents:
simulation:
network:
nmne_config:
capture_nmne: true
nmne_capture_keywords:
- DELETE
nodes:
- ref: router_1
@@ -1032,12 +1112,13 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- ref: web_server_web_service
type: WebServer
applications:
- ref: web_server_database_client
type: DatabaseClient
options:
db_server_ip: 192.168.1.14
- ref: web_server_web_service
type: WebServer
- ref: database_server
@@ -1089,10 +1170,14 @@ simulation:
- ref: data_manipulation_bot
type: DataManipulationBot
options:
port_scan_p_of_success: 0.1
data_manipulation_p_of_success: 0.1
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- ref: client_1_web_browser
type: WebBrowser
options:
target_url: http://arcd.com/users/
services:
- ref: client_1_dns_client
type: DNSClient
@@ -1109,6 +1194,13 @@ simulation:
type: WebBrowser
options:
target_url: http://arcd.com/users/
- ref: data_manipulation_bot
type: DataManipulationBot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
services:
- ref: client_2_dns_client
type: DNSClient

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

@@ -8,6 +8,7 @@ from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.nmne import CAPTURE_NMNE
_LOGGER = getLogger(__name__)
@@ -346,7 +347,14 @@ class FolderObservation(AbstractObservation):
class NicObservation(AbstractObservation):
"""Observation of a Network Interface Card (NIC) in the network."""
default_observation: spaces.Space = {"nic_status": 0}
@property
def default_observation(self) -> Dict:
"""The default NIC observation dict."""
data = {"nic_status": 0}
if CAPTURE_NMNE:
data.update({"nmne": {"inbound": 0, "outbound": 0}})
return data
def __init__(self, where: Optional[Tuple[str]] = None) -> None:
"""Initialise NIC observation.
@@ -360,6 +368,29 @@ class NicObservation(AbstractObservation):
super().__init__()
self.where: Optional[Tuple[str]] = where
def _categorise_mne_count(self, nmne_count: int) -> int:
"""
Categorise the number of Malicious Network Events (NMNEs) into discrete bins.
This helps in classifying the severity or volume of MNEs into manageable levels for the agent.
Bins are defined as follows:
- 0: No MNEs detected (0 events).
- 1: Low number of MNEs (1-5 events).
- 2: Moderate number of MNEs (6-10 events).
- 3: High number of MNEs (more than 10 events).
:param nmne_count: Number of MNEs detected.
:return: Bin number corresponding to the number of MNEs. Returns 0, 1, 2, or 3 based on the detected MNE count.
"""
if nmne_count > 10:
return 3
elif nmne_count > 5:
return 2
elif nmne_count > 0:
return 1
return 0
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
@@ -371,15 +402,31 @@ class NicObservation(AbstractObservation):
if self.where is None:
return self.default_observation
nic_state = access_from_nested_dict(state, self.where)
if nic_state is NOT_PRESENT_IN_STATE:
return self.default_observation
else:
return {"nic_status": 1 if nic_state["enabled"] else 2}
obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2}
if CAPTURE_NMNE:
obs_dict.update({"nmne": {}})
direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
inbound_count = inbound_keywords.get("*", 0)
outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {})
outbound_count = outbound_keywords.get("*", 0)
obs_dict["nmne"]["inbound"] = self._categorise_mne_count(inbound_count)
obs_dict["nmne"]["outbound"] = self._categorise_mne_count(outbound_count)
return obs_dict
@property
def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape."""
return spaces.Dict({"nic_status": spaces.Discrete(3)})
return spaces.Dict(
{
"nic_status": spaces.Discrete(3),
"nmne": spaces.Dict({"inbound": spaces.Discrete(6), "outbound": spaces.Discrete(6)}),
}
)
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]]) -> "NicObservation":

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

@@ -18,6 +18,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import NIC
from primaite.simulator.network.hardware.nodes.host.server import Server
from primaite.simulator.network.hardware.nodes.network.router import Router
from primaite.simulator.network.hardware.nodes.network.switch import Switch
from primaite.simulator.network.nmne import set_nmne_config
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
@@ -80,18 +81,15 @@ class PrimaiteGame:
self.simulation: Simulation = Simulation()
"""Simulation object with which the agents will interact."""
self.agents: List[AbstractAgent] = []
"""List of agents."""
self.agents: Dict[str, AbstractAgent] = {}
"""Mapping from agent name to agent object."""
self.rl_agents: List[ProxyAgent] = []
"""Subset of agent list including only the reinforcement learning agents."""
self.rl_agents: Dict[str, ProxyAgent] = {}
"""Subset of agents which are intended for reinforcement learning."""
self.step_counter: int = 0
"""Current timestep within the episode."""
self.episode_counter: int = 0
"""Current episode number."""
self.options: PrimaiteGameOptions
"""Special options that apply for the entire game."""
@@ -137,7 +135,7 @@ class PrimaiteGame:
self.update_agents(sim_state)
# Apply all actions to simulation as requests
agent_actions = self.apply_agent_actions() # noqa
self.apply_agent_actions()
# Advance timestep
self.advance_timestep()
@@ -148,7 +146,7 @@ class PrimaiteGame:
def update_agents(self, state: Dict) -> None:
"""Update agents' observations and rewards based on the current state."""
for agent in self.agents:
for _, agent in self.agents.items():
agent.update_observation(state)
agent.update_reward(state)
agent.reward_function.total_reward += agent.reward_function.current_reward
@@ -162,7 +160,7 @@ class PrimaiteGame:
"""
agent_actions = {}
for agent in self.agents:
for _, agent in self.agents.items():
obs = agent.observation_manager.current_observation
rew = agent.reward_function.current_reward
action_choice, options = agent.get_action(obs, rew, timestep=self.step_counter)
@@ -185,20 +183,14 @@ class PrimaiteGame:
return True
return False
def reset(self) -> None:
"""Reset the game, this will reset the simulation."""
self.episode_counter += 1
self.step_counter = 0
_LOGGER.debug(f"Resetting primaite game, episode = {self.episode_counter}")
self.simulation.reset_component_for_episode(episode=self.episode_counter)
for agent in self.agents:
agent.reward_function.total_reward = 0.0
agent.reset_agent_for_episode()
def close(self) -> None:
"""Close the game, this will close the simulation."""
return NotImplemented
def setup_for_episode(self, episode: int) -> None:
"""Perform any final configuration of components to make them ready for the game to start."""
self.simulation.setup_for_episode(episode=episode)
@classmethod
def from_config(cls, cfg: Dict) -> "PrimaiteGame":
"""Create a PrimaiteGame object from a config dictionary.
@@ -272,7 +264,9 @@ class PrimaiteGame:
new_service = new_node.software_manager.software[service_type]
game.ref_map_services[service_ref] = new_service.uuid
else:
_LOGGER.warning(f"service type not found {service_type}")
msg = f"Configuration contains an invalid service type: {service_type}"
_LOGGER.error(msg)
raise ValueError(msg)
# service-dependent options
if service_type == "DNSClient":
if "options" in service_cfg:
@@ -311,7 +305,9 @@ class PrimaiteGame:
new_application = new_node.software_manager.software[application_type]
game.ref_map_applications[application_ref] = new_application.uuid
else:
_LOGGER.warning(f"application type not found {application_type}")
msg = f"Configuration contains an invalid application type: {application_type}"
_LOGGER.error(msg)
raise ValueError(msg)
if application_type == "DataManipulationBot":
if "options" in application_cfg:
@@ -403,7 +399,6 @@ class PrimaiteGame:
reward_function=reward_function,
settings=settings,
)
game.agents.append(new_agent)
elif agent_type == "ProxyAgent":
new_agent = ProxyAgent(
agent_name=agent_cfg["ref"],
@@ -412,8 +407,7 @@ class PrimaiteGame:
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
game.rl_agents.append(new_agent)
game.rl_agents[agent_cfg["ref"]] = new_agent
elif agent_type == "RedDatabaseCorruptingAgent":
new_agent = DataManipulationAgent(
agent_name=agent_cfg["ref"],
@@ -422,10 +416,13 @@ class PrimaiteGame:
reward_function=reward_function,
agent_settings=agent_settings,
)
game.agents.append(new_agent)
else:
_LOGGER.warning(f"agent type {agent_type} not found")
msg(f"Configuration error: {agent_type} is not a valid agent type.")
_LOGGER.error(msg)
raise ValueError(msg)
game.agents[agent_cfg["ref"]] = new_agent
game.simulation.set_original_state()
# Set the NMNE capture config
set_nmne_config(cfg["simulation"]["network"].get("nmne_config", {}))
return game

9
src/primaite/notebooks/training_example_ray_multi_agent.ipynb
View File

@@ -60,7 +60,7 @@
" policies={'defender_1','defender_2'}, # These names are the same as the agents defined in the example config.\n",
" policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,\n",
" )\n",
" .environment(env=PrimaiteRayMARLEnv, env_config={\"cfg\":cfg})#, disable_env_checking=True)\n",
" .environment(env=PrimaiteRayMARLEnv, env_config=cfg)#, disable_env_checking=True)\n",
" .rollouts(num_rollout_workers=0)\n",
" .training(train_batch_size=128)\n",
" )\n"
@@ -88,6 +88,13 @@
" param_space=config\n",
").fit()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {

2
src/primaite/notebooks/training_example_ray_single_agent.ipynb
View File

@@ -54,7 +54,7 @@
"metadata": {},
"outputs": [],
"source": [
"env_config = {\"cfg\":cfg}\n",
"env_config = cfg\n",
"\n",
"config = (\n",
" PPOConfig()\n",

15
src/primaite/notebooks/training_example_sb3.ipynb
View File

@@ -27,9 +27,7 @@
"outputs": [],
"source": [
"with open(example_config_path(), 'r') as f:\n",
" cfg = yaml.safe_load(f)\n",
"\n",
"game = PrimaiteGame.from_config(cfg)"
" cfg = yaml.safe_load(f)\n"
]
},
{
@@ -38,7 +36,7 @@
"metadata": {},
"outputs": [],
"source": [
"gym = PrimaiteGymEnv(game=game)"
"gym = PrimaiteGymEnv(game_config=cfg)"
]
},
{
@@ -65,7 +63,7 @@
"metadata": {},
"outputs": [],
"source": [
"model.learn(total_timesteps=1000)\n"
"model.learn(total_timesteps=10)\n"
]
},
{
@@ -76,6 +74,13 @@
"source": [
"model.save(\"deleteme\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {

256
src/primaite/notebooks/uc2_demo.ipynb
View File

@@ -130,6 +130,9 @@
" - NETWORK_INTERFACES\n",
" - <nic_id 1-2>\n",
" - nic_status\n",
" - nmne\n",
" - inbound\n",
" - outbound\n",
" - operating_status\n",
"- LINKS\n",
" - <link_id 1-10>\n",
@@ -220,6 +223,14 @@
"|1|ENABLED|\n",
"|2|DISABLED|\n",
"\n",
"NMNE (number of malicious network events) means, for inbound or outbound traffic, means:\n",
"|value|NMNEs|\n",
"|--|--|\n",
"|0|None|\n",
"|1|1 - 5|\n",
"|2|6 - 10|\n",
"|3|More than 10|\n",
"\n",
"Link load has the following meaning:\n",
"|load|percent utilisation|\n",
"|--|--|\n",
@@ -371,150 +382,9 @@
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2024-02-27 09:43:39,312::WARNING::primaite.game.game::275::service type not found DatabaseClient\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Resetting environment, episode 0, avg. reward: 0.0\n",
"env created successfully\n",
"{'ACL': {1: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 0,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 2: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 1,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 3: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 2,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 4: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 3,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 5: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 4,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 6: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 5,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 7: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 6,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 8: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 7,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 9: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 8,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0},\n",
" 10: {'dest_node_id': 0,\n",
" 'dest_port': 0,\n",
" 'permission': 0,\n",
" 'position': 9,\n",
" 'protocol': 0,\n",
" 'source_node_id': 0,\n",
" 'source_port': 0}},\n",
" 'ICS': 0,\n",
" 'LINKS': {1: {'PROTOCOLS': {'ALL': 1}},\n",
" 2: {'PROTOCOLS': {'ALL': 1}},\n",
" 3: {'PROTOCOLS': {'ALL': 1}},\n",
" 4: {'PROTOCOLS': {'ALL': 1}},\n",
" 5: {'PROTOCOLS': {'ALL': 1}},\n",
" 6: {'PROTOCOLS': {'ALL': 1}},\n",
" 7: {'PROTOCOLS': {'ALL': 1}},\n",
" 8: {'PROTOCOLS': {'ALL': 1}},\n",
" 9: {'PROTOCOLS': {'ALL': 1}},\n",
" 10: {'PROTOCOLS': {'ALL': 0}}},\n",
" 'NODES': {1: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 2: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 1}},\n",
" 'operating_status': 1},\n",
" 3: {'FOLDERS': {1: {'FILES': {1: {'health_status': 1}},\n",
" 'health_status': 1}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 4: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 5: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 6: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1},\n",
" 7: {'FOLDERS': {1: {'FILES': {1: {'health_status': 0}},\n",
" 'health_status': 0}},\n",
" 'NETWORK_INTERFACES': {1: {'nic_status': 1},\n",
" 2: {'nic_status': 0}},\n",
" 'SERVICES': {1: {'health_status': 0, 'operating_status': 0}},\n",
" 'operating_status': 1}}}\n"
]
}
],
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# create the env\n",
"with open(example_config_path(), 'r') as f:\n",
@@ -524,10 +394,10 @@
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['data_manipulation_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][8]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
" cfg['simulation']['network']['nodes'][9]['applications'][0]['options']['port_scan_p_of_success'] = 1.0\n",
"game = PrimaiteGame.from_config(cfg)\n",
"env = PrimaiteGymEnv(game = game)\n",
"# Don't flatten obs as we are not training an agent and we wish to see the dict-formatted observations\n",
"env.agent.flatten_obs = False\n",
" # don't flatten observations so that we can see what is going on\n",
" cfg['agents'][3]['agent_settings']['flatten_obs'] = False\n",
"\n",
"env = PrimaiteGymEnv(game_config = cfg)\n",
"obs, info = env.reset()\n",
"print('env created successfully')\n",
"pprint(obs)"
@@ -560,53 +430,9 @@
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"step: 211, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 212, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 213, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 214, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 215, Red action: DO NOTHING, Blue reward:-0.42\n",
"step: 216, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 217, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 218, Red action: DO NOTHING, Blue reward:-0.42\n",
"step: 219, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 220, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 221, Red action: ATTACK from client 2, Blue reward:-0.32\n",
"step: 222, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 223, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 224, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 225, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 226, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 227, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 228, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 229, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 230, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 231, Red action: DO NOTHING, Blue reward:-0.42\n",
"step: 232, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 233, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 234, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 235, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 236, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 237, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 238, Red action: ATTACK from client 2, Blue reward:-0.32\n",
"step: 239, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 240, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 241, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 242, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 243, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 244, Red action: DO NOTHING, Blue reward:-0.32\n",
"step: 245, Red action: DO NOTHING, Blue reward:-0.32\n"
]
}
],
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for step in range(35):\n",
" obs, reward, terminated, truncated, info = env.step(0)\n",
@@ -623,9 +449,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"pprint(obs['NODES'])"
@@ -641,9 +465,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(9) # scan database file\n",
@@ -669,9 +491,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(13) # patch the database\n",
@@ -696,9 +516,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"obs, reward, terminated, truncated, info = env.step(0) # patch the database\n",
@@ -721,9 +539,7 @@
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"tags": []
},
"metadata": {},
"outputs": [],
"source": [
"env.step(13) # Patch the database\n",
@@ -763,6 +579,22 @@
"obs['ACL']"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Reset the environment, you can rerun the other cells to verify that the attack works the same every episode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"env.reset()"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -800,7 +632,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"display_name": "venv",
"language": "python",
"name": "python3"
},
@@ -818,5 +650,5 @@
}
},
"nbformat": 4,
"nbformat_minor": 4
"nbformat_minor": 2
}

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

@@ -1,5 +1,5 @@
import json
from typing import Any, Dict, Final, Optional, SupportsFloat, Tuple
from typing import Any, Dict, Optional, SupportsFloat, Tuple
import gymnasium
from gymnasium.core import ActType, ObsType
@@ -18,11 +18,23 @@ class PrimaiteGymEnv(gymnasium.Env):
assumptions about the agent list always having a list of length 1.
"""
def __init__(self, game: PrimaiteGame):
def __init__(self, game_config: Dict):
"""Initialise the environment."""
super().__init__()
self.game: "PrimaiteGame" = game
self.agent: ProxyAgent = self.game.rl_agents[0]
self.game_config: Dict = game_config
"""PrimaiteGame definition. This can be changed between episodes to enable curriculum learning."""
self.game: PrimaiteGame = PrimaiteGame.from_config(self.game_config)
"""Current game."""
self._agent_name = next(iter(self.game.rl_agents))
"""Name of the RL agent. Since there should only be one RL agent we can just pull the first and only key."""
self.episode_counter: int = 0
"""Current episode number."""
@property
def agent(self) -> ProxyAgent:
"""Grab a fresh reference to the agent object because it will be reinstantiated each episode."""
return self.game.rl_agents[self._agent_name]
def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
"""Perform a step in the environment."""
@@ -45,13 +57,13 @@ class PrimaiteGymEnv(gymnasium.Env):
return next_obs, reward, terminated, truncated, info
def _write_step_metadata_json(self, action: int, state: Dict, reward: int):
output_dir = SIM_OUTPUT.path / f"episode_{self.game.episode_counter}" / "step_metadata"
output_dir = SIM_OUTPUT.path / f"episode_{self.episode_counter}" / "step_metadata"
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"step_{self.game.step_counter}.json"
data = {
"episode": self.game.episode_counter,
"episode": self.episode_counter,
"step": self.game.step_counter,
"action": int(action),
"reward": int(reward),
@@ -63,10 +75,12 @@ class PrimaiteGymEnv(gymnasium.Env):
def reset(self, seed: Optional[int] = None) -> Tuple[ObsType, Dict[str, Any]]:
"""Reset the environment."""
print(
f"Resetting environment, episode {self.game.episode_counter}, "
f"avg. reward: {self.game.rl_agents[0].reward_function.total_reward}"
f"Resetting environment, episode {self.episode_counter}, "
f"avg. reward: {self.agent.reward_function.total_reward}"
)
self.game.reset()
self.game: PrimaiteGame = PrimaiteGame.from_config(cfg=self.game_config)
self.game.setup_for_episode(episode=self.episode_counter)
self.episode_counter += 1
state = self.game.get_sim_state()
self.game.update_agents(state)
next_obs = self._get_obs()
@@ -88,12 +102,12 @@ class PrimaiteGymEnv(gymnasium.Env):
def _get_obs(self) -> ObsType:
"""Return the current observation."""
if not self.agent.flatten_obs:
return self.agent.observation_manager.current_observation
else:
if self.agent.flatten_obs:
unflat_space = self.agent.observation_manager.space
unflat_obs = self.agent.observation_manager.current_observation
return gymnasium.spaces.flatten(unflat_space, unflat_obs)
else:
return self.agent.observation_manager.current_observation
class PrimaiteRayEnv(gymnasium.Env):
@@ -102,12 +116,11 @@ class PrimaiteRayEnv(gymnasium.Env):
def __init__(self, env_config: Dict) -> None:
"""Initialise the environment.
:param env_config: A dictionary containing the environment configuration. It must contain a single key, `game`
which is the PrimaiteGame instance.
:type env_config: Dict[str, PrimaiteGame]
:param env_config: A dictionary containing the environment configuration.
:type env_config: Dict
"""
self.env = PrimaiteGymEnv(game=PrimaiteGame.from_config(env_config["cfg"]))
self.env.game.episode_counter -= 1
self.env = PrimaiteGymEnv(game_config=env_config)
self.env.episode_counter -= 1
self.action_space = self.env.action_space
self.observation_space = self.env.observation_space
@@ -128,13 +141,16 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
:param env_config: A dictionary containing the environment configuration. It must contain a single key, `game`
which is the PrimaiteGame instance.
:type env_config: Dict[str, PrimaiteGame]
:type env_config: Dict
"""
self.game: PrimaiteGame = PrimaiteGame.from_config(env_config["cfg"])
self.game_config: Dict = env_config
"""PrimaiteGame definition. This can be changed between episodes to enable curriculum learning."""
self.game: PrimaiteGame = PrimaiteGame.from_config(self.game_config)
"""Reference to the primaite game"""
self.agents: Final[Dict[str, ProxyAgent]] = {agent.agent_name: agent for agent in self.game.rl_agents}
"""List of all possible agents in the environment. This list should not change!"""
self._agent_ids = list(self.agents.keys())
self._agent_ids = list(self.game.rl_agents.keys())
"""Agent ids. This is a list of strings of agent names."""
self.episode_counter: int = 0
"""Current episode number."""
self.terminateds = set()
self.truncateds = set()
@@ -149,9 +165,16 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
)
super().__init__()
@property
def agents(self) -> Dict[str, ProxyAgent]:
"""Grab a fresh reference to the agents from this episode's game object."""
return {name: self.game.rl_agents[name] for name in self._agent_ids}
def reset(self, *, seed: int = None, options: dict = None) -> Tuple[ObsType, Dict]:
"""Reset the environment."""
self.game.reset()
self.game: PrimaiteGame = PrimaiteGame.from_config(cfg=self.game_config)
self.game.setup_for_episode(episode=self.episode_counter)
self.episode_counter += 1
state = self.game.get_sim_state()
self.game.update_agents(state)
next_obs = self._get_obs()
@@ -172,7 +195,7 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
# 1. Perform actions
for agent_name, action in actions.items():
self.agents[agent_name].store_action(action)
agent_actions = self.game.apply_agent_actions()
self.game.apply_agent_actions()
# 2. Advance timestep
self.game.advance_timestep()
@@ -186,7 +209,7 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
rewards = {name: agent.reward_function.current_reward for name, agent in self.agents.items()}
terminateds = {name: False for name, _ in self.agents.items()}
truncateds = {name: self.game.calculate_truncated() for name, _ in self.agents.items()}
infos = {"agent_actions": agent_actions}
infos = {name: {} for name, _ in self.agents.items()}
terminateds["__all__"] = len(self.terminateds) == len(self.agents)
truncateds["__all__"] = self.game.calculate_truncated()
if self.game.save_step_metadata:
@@ -194,13 +217,13 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
return next_obs, rewards, terminateds, truncateds, infos
def _write_step_metadata_json(self, actions: Dict, state: Dict, rewards: Dict):
output_dir = SIM_OUTPUT.path / f"episode_{self.game.episode_counter}" / "step_metadata"
output_dir = SIM_OUTPUT.path / f"episode_{self.episode_counter}" / "step_metadata"
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"step_{self.game.step_counter}.json"
data = {
"episode": self.game.episode_counter,
"episode": self.episode_counter,
"step": self.game.step_counter,
"actions": {agent_name: int(action) for agent_name, action in actions.items()},
"reward": rewards,
@@ -212,8 +235,9 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
def _get_obs(self) -> Dict[str, ObsType]:
"""Return the current observation."""
obs = {}
for name, agent in self.agents.items():
for agent_name in self._agent_ids:
agent = self.game.rl_agents[agent_name]
unflat_space = agent.observation_manager.space
unflat_obs = agent.observation_manager.current_observation
obs[name] = gymnasium.spaces.flatten(unflat_space, unflat_obs)
obs[agent_name] = gymnasium.spaces.flatten(unflat_space, unflat_obs)
return obs

6
src/primaite/session/session.py
View File

@@ -101,11 +101,11 @@ class PrimaiteSession:
# CREATE ENVIRONMENT
if sess.training_options.rl_framework == "RLLIB_single_agent":
sess.env = PrimaiteRayEnv(env_config={"cfg": cfg})
sess.env = PrimaiteRayEnv(env_config=cfg)
elif sess.training_options.rl_framework == "RLLIB_multi_agent":
sess.env = PrimaiteRayMARLEnv(env_config={"cfg": cfg})
sess.env = PrimaiteRayMARLEnv(env_config=cfg)
elif sess.training_options.rl_framework == "SB3":
sess.env = PrimaiteGymEnv(game=game)
sess.env = PrimaiteGymEnv(game_config=cfg)
sess.policy = PolicyABC.from_config(sess.training_options, session=sess)
if agent_load_path:

16
src/primaite/simulator/core.py
View File

@@ -153,22 +153,18 @@ class SimComponent(BaseModel):
uuid: str = Field(default_factory=lambda: str(uuid4()))
"""The component UUID."""
_original_state: Dict = {}
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._request_manager: RequestManager = self._init_request_manager()
self._parent: Optional["SimComponent"] = None
# @abstractmethod
def set_original_state(self):
"""Sets the original state."""
pass
def setup_for_episode(self, episode: int):
"""
Perform any additional setup on this component that can't happen during __init__.
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
for key, value in self._original_state.items():
self.__setattr__(key, value)
For instance, some components may require for the entire network to exist before some configuration can be set.
"""
pass
def _init_request_manager(self) -> RequestManager:
"""

13
src/primaite/simulator/domain/account.py
View File

@@ -42,19 +42,6 @@ class Account(SimComponent):
"Account Type, currently this can be service account (used by apps) or user account."
enabled: bool = True
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {
"num_logons",
"num_logoffs",
"num_group_changes",
"username",
"password",
"account_type",
"enabled",
}
self._original_state = self.model_dump(include=vals_to_include)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

14
src/primaite/simulator/file_system/file.py
View File

@@ -73,20 +73,6 @@ class File(FileSystemItemABC):
self.sys_log.info(f"Created file /{self.path} (id: {self.uuid})")
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting File ({self.path}) original state on node {self.sys_log.hostname}")
super().set_original_state()
vals_to_include = {"folder_id", "folder_name", "file_type", "sim_size", "real", "sim_path", "sim_root"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting File ({self.path}) state on node {self.sys_log.hostname}")
super().reset_component_for_episode(episode)
@property
def path(self) -> str:
"""

37
src/primaite/simulator/file_system/file_system.py
View File

@@ -34,43 +34,6 @@ class FileSystem(SimComponent):
if not self.folders:
self.create_folder("root")
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting FileSystem original state on node {self.sys_log.hostname}")
for folder in self.folders.values():
folder.set_original_state()
# Capture a list of all 'original' file uuids
original_keys = list(self.folders.keys())
vals_to_include = {"sim_root"}
self._original_state.update(self.model_dump(include=vals_to_include))
self._original_state["original_folder_uuids"] = original_keys
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting FileSystem state on node {self.sys_log.hostname}")
# Move any 'original' folder that have been deleted back to folders
original_folder_uuids = self._original_state["original_folder_uuids"]
for uuid in original_folder_uuids:
if uuid in self.deleted_folders:
folder = self.deleted_folders[uuid]
self.deleted_folders.pop(uuid)
self.folders[uuid] = folder
# Clear any other deleted folders that aren't original (have been created by agent)
self.deleted_folders.clear()
# Now clear all non-original folders created by agent
current_folder_uuids = list(self.folders.keys())
for uuid in current_folder_uuids:
if uuid not in original_folder_uuids:
folder = self.folders[uuid]
self.folders.pop(uuid)
# Now reset all remaining folders
for folder in self.folders.values():
folder.reset_component_for_episode(episode)
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()

5
src/primaite/simulator/file_system/file_system_item_abc.py
View File

@@ -85,11 +85,6 @@ class FileSystemItemABC(SimComponent):
deleted: bool = False
"If true, the FileSystemItem was deleted."
def set_original_state(self):
"""Sets the original state."""
vals_to_keep = {"name", "health_status", "visible_health_status", "previous_hash", "revealed_to_red", "deleted"}
self._original_state = self.model_dump(include=vals_to_keep)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

43
src/primaite/simulator/file_system/folder.py
View File

@@ -49,49 +49,6 @@ class Folder(FileSystemItemABC):
self.sys_log.info(f"Created file /{self.name} (id: {self.uuid})")
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting Folder ({self.name}) original state on node {self.sys_log.hostname}")
for file in self.files.values():
file.set_original_state()
super().set_original_state()
vals_to_include = {
"scan_duration",
"scan_countdown",
"red_scan_duration",
"red_scan_countdown",
"restore_duration",
"restore_countdown",
}
self._original_state.update(self.model_dump(include=vals_to_include))
self._original_state["original_file_uuids"] = list(self.files.keys())
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting Folder ({self.name}) state on node {self.sys_log.hostname}")
# Move any 'original' file that have been deleted back to files
original_file_uuids = self._original_state["original_file_uuids"]
for uuid in original_file_uuids:
if uuid in self.deleted_files:
file = self.deleted_files[uuid]
self.deleted_files.pop(uuid)
self.files[uuid] = file
# Clear any other deleted files that aren't original (have been created by agent)
self.deleted_files.clear()
# Now clear all non-original files created by agent
current_file_uuids = list(self.files.keys())
for uuid in current_file_uuids:
if uuid not in original_file_uuids:
file = self.files[uuid]
self.files.pop(uuid)
# Now reset all remaining files
for file in self.files.values():
file.reset_component_for_episode(episode)
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(

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

@@ -273,11 +273,6 @@ class IPWirelessNetworkInterface(WirelessNetworkInterface, Layer3Interface, ABC)
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def enable(self):
"""
Enables this wired network interface and attempts to send a "hello" message to the default gateway.

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

@@ -45,19 +45,12 @@ class Network(SimComponent):
self._nx_graph = MultiGraph()
def set_original_state(self):
"""Sets the original state."""
for node in self.nodes.values():
node.set_original_state()
for link in self.links.values():
link.set_original_state()
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
for node in self.nodes.values():
node.reset_component_for_episode(episode)
node.setup_for_episode(episode=episode)
for link in self.links.values():
link.reset_component_for_episode(episode)
link.setup_for_episode(episode=episode)
for node in self.nodes.values():
node.power_on()
@@ -179,7 +172,7 @@ class Network(SimComponent):
def clear_links(self):
"""Clear all the links in the network by resetting their component state for the episode."""
for link in self.links.values():
link.reset_component_for_episode()
link.setup_for_episode(episode=0) # TODO: shouldn't be using this method here.
def draw(self, seed: int = 123):
"""

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

@@ -17,6 +17,15 @@ from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.domain.account import Account
from primaite.simulator.file_system.file_system import FileSystem
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
from primaite.simulator.network.nmne import (
CAPTURE_BY_DIRECTION,
CAPTURE_BY_IP_ADDRESS,
CAPTURE_BY_KEYWORD,
CAPTURE_BY_PORT,
CAPTURE_BY_PROTOCOL,
CAPTURE_NMNE,
NMNE_CAPTURE_KEYWORDS,
)
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.core.packet_capture import PacketCapture
@@ -88,6 +97,18 @@ class NetworkInterface(SimComponent, ABC):
pcap: Optional[PacketCapture] = None
"A PacketCapture instance for capturing and analysing packets passing through this interface."
nmne: Dict = Field(default_factory=lambda: {})
"A dict containing details of the number of malicious network events captured."
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().setup_for_episode(episode=episode)
self.nmne = {}
if episode and self.pcap:
self.pcap.current_episode = episode
self.pcap.setup_logger()
self.enable()
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
@@ -111,16 +132,10 @@ class NetworkInterface(SimComponent, ABC):
"enabled": self.enabled,
}
)
if CAPTURE_NMNE:
state.update({"nmne": self.nmne})
return state
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().reset_component_for_episode(episode)
if episode and self.pcap:
self.pcap.current_episode = episode
self.pcap.setup_logger()
self.enable()
@abstractmethod
def enable(self):
"""Enable the interface."""
@@ -131,6 +146,82 @@ class NetworkInterface(SimComponent, ABC):
"""Disable the interface."""
pass
def _capture_nmne(self, frame: Frame, inbound: bool = True) -> None:
"""
Processes and captures network frame data based on predefined global NMNE settings.
This method updates the NMNE structure with counts of malicious network events based on the frame content and
direction. The structure is dynamically adjusted according to the enabled capture settings.
.. note::
While there is a lot of logic in this code that defines a multi-level hierarchical NMNE structure,
most of it is unused for now as a result of all `CAPTURE_BY_<>` variables in
``primaite.simulator.network.nmne`` being hardcoded and set as final. Once they're 'released' and made
configurable, this function will be updated to properly explain the dynamic data structure.
:param frame: The network frame to process, containing IP, TCP/UDP, and payload information.
:param inbound: Boolean indicating if the frame direction is inbound. Defaults to True.
"""
# Exit function if NMNE capturing is disabled
if not CAPTURE_NMNE:
return
# Initialise basic frame data variables
direction = "inbound" if inbound else "outbound" # Direction of the traffic
ip_address = str(frame.ip.src_ip_address if inbound else frame.ip.dst_ip_address) # Source or destination IP
protocol = frame.ip.protocol.name # Network protocol used in the frame
# Initialise port variable; will be determined based on protocol type
port = None
# Determine the source or destination port based on the protocol (TCP/UDP)
if frame.tcp:
port = frame.tcp.src_port.value if inbound else frame.tcp.dst_port.value
elif frame.udp:
port = frame.udp.src_port.value if inbound else frame.udp.dst_port.value
# Convert frame payload to string for keyword checking
frame_str = str(frame.payload)
# Proceed only if any NMNE keyword is present in the frame payload
if any(keyword in frame_str for keyword in NMNE_CAPTURE_KEYWORDS):
# Start with the root of the NMNE capture structure
current_level = self.nmne
# Update NMNE structure based on enabled settings
if CAPTURE_BY_DIRECTION:
# Set or get the dictionary for the current direction
current_level = current_level.setdefault("direction", {})
current_level = current_level.setdefault(direction, {})
if CAPTURE_BY_IP_ADDRESS:
# Set or get the dictionary for the current IP address
current_level = current_level.setdefault("ip_address", {})
current_level = current_level.setdefault(ip_address, {})
if CAPTURE_BY_PROTOCOL:
# Set or get the dictionary for the current protocol
current_level = current_level.setdefault("protocol", {})
current_level = current_level.setdefault(protocol, {})
if CAPTURE_BY_PORT:
# Set or get the dictionary for the current port
current_level = current_level.setdefault("port", {})
current_level = current_level.setdefault(port, {})
# Ensure 'KEYWORD' level is present in the structure
keyword_level = current_level.setdefault("keywords", {})
# Increment the count for detected keywords in the payload
if CAPTURE_BY_KEYWORD:
for keyword in NMNE_CAPTURE_KEYWORDS:
if keyword in frame_str:
# Update the count for each keyword found
keyword_level[keyword] = keyword_level.get(keyword, 0) + 1
else:
# Increment a generic counter if keyword capturing is not enabled
keyword_level["*"] = keyword_level.get("*", 0) + 1
@abstractmethod
def send_frame(self, frame: Frame) -> bool:
"""
@@ -139,7 +230,7 @@ class NetworkInterface(SimComponent, ABC):
:param frame: The network frame to be sent.
:return: A boolean indicating whether the frame was successfully sent.
"""
pass
self._capture_nmne(frame, inbound=False)
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
@@ -149,7 +240,7 @@ class NetworkInterface(SimComponent, ABC):
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
self._capture_nmne(frame, inbound=True)
def __str__(self) -> str:
"""
@@ -263,6 +354,7 @@ class WiredNetworkInterface(NetworkInterface, ABC):
:param frame: The network frame to be sent.
:return: True if the frame is sent, False if the Network Interface is disabled or not connected to a link.
"""
super().send_frame(frame)
if self.enabled:
frame.set_sent_timestamp()
self.pcap.capture_outbound(frame)
@@ -279,7 +371,7 @@ class WiredNetworkInterface(NetworkInterface, ABC):
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
return super().receive_frame(frame)
class Layer3Interface(BaseModel, ABC):
@@ -409,7 +501,7 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC):
except AttributeError:
pass
# @abstractmethod
@abstractmethod
def receive_frame(self, frame: Frame) -> bool:
"""
Receives a network frame on the network interface.
@@ -417,7 +509,7 @@ class IPWiredNetworkInterface(WiredNetworkInterface, Layer3Interface, ABC):
:param frame: The network frame being received.
:return: A boolean indicating whether the frame was successfully received.
"""
pass
return super().receive_frame(frame)
class Link(SimComponent):
@@ -455,14 +547,6 @@ class Link(SimComponent):
self.endpoint_b.connect_link(self)
self.endpoint_up()
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"bandwidth", "current_load"}
self._original_state = self.model_dump(include=vals_to_include)
super().set_original_state()
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -648,50 +732,20 @@ class Node(SimComponent):
self.session_manager.node = self
self.session_manager.software_manager = self.software_manager
self._install_system_software()
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
for software in self.software_manager.software.values():
software.set_original_state()
self.file_system.set_original_state()
for network_interface in self.network_interfaces.values():
network_interface.set_original_state()
vals_to_include = {
"hostname",
"default_gateway",
"operating_state",
"revealed_to_red",
"start_up_duration",
"start_up_countdown",
"shut_down_duration",
"shut_down_countdown",
"is_resetting",
"node_scan_duration",
"node_scan_countdown",
"red_scan_countdown",
}
self._original_state = self.model_dump(include=vals_to_include)
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
super().reset_component_for_episode(episode)
# Reset Session Manager
self.session_manager.clear()
super().setup_for_episode(episode=episode)
# Reset File System
self.file_system.reset_component_for_episode(episode)
self.file_system.setup_for_episode(episode=episode)
# Reset all Nics
for network_interface in self.network_interfaces.values():
network_interface.reset_component_for_episode(episode)
network_interface.setup_for_episode(episode=episode)
for software in self.software_manager.software.values():
software.reset_component_for_episode(episode)
software.setup_for_episode(episode=episode)
if episode and self.sys_log:
self.sys_log.current_episode = episode

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

@@ -205,19 +205,10 @@ class NIC(IPWiredNetworkInterface):
state = super().describe_state()
# Update the state with NIC-specific information
state.update(
{
"wake_on_lan": self.wake_on_lan,
}
)
state.update({"wake_on_lan": self.wake_on_lan})
return state
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"ip_address", "subnet_mask", "mac_address", "speed", "mtu", "wake_on_lan", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
def receive_frame(self, frame: Frame) -> bool:
"""
Attempt to receive and process a network frame from the connected Link.
@@ -248,6 +239,7 @@ class NIC(IPWiredNetworkInterface):
accept_frame = True
if accept_frame:
super().receive_frame(frame)
self._connected_node.receive_frame(frame=frame, from_network_interface=self)
return True
return False

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

@@ -109,24 +109,6 @@ class Firewall(Router):
sys_log=kwargs["sys_log"], implicit_action=ACLAction.PERMIT, name=f"{hostname} - External Outbound"
)
self.set_original_state()
def set_original_state(self):
"""Set the original state for the Firewall."""
super().set_original_state()
vals_to_include = {
"internal_port",
"external_port",
"dmz_port",
"internal_inbound_acl",
"internal_outbound_acl",
"dmz_inbound_acl",
"dmz_outbound_acl",
"external_inbound_acl",
"external_outbound_acl",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def describe_state(self) -> Dict:
"""
Describes the current state of the Firewall.

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

@@ -136,11 +136,6 @@ class ACLRule(SimComponent):
rule_strings.append(f"{key}={value}")
return ", ".join(rule_strings)
def set_original_state(self):
"""Sets the original state."""
vals_to_keep = {"action", "protocol", "src_ip_address", "src_port", "dst_ip_address", "dst_port"}
self._original_state = self.model_dump(include=vals_to_keep, exclude_none=True)
def describe_state(self) -> Dict:
"""
Describes the current state of the ACLRule.
@@ -296,48 +291,6 @@ class AccessControlList(SimComponent):
super().__init__(**kwargs)
self._acl = [None] * (self.max_acl_rules - 1)
self.set_original_state()
def set_original_state(self):
"""Sets the original state."""
self.implicit_rule.set_original_state()
vals_to_keep = {"implicit_action", "max_acl_rules", "acl"}
self._original_state = self.model_dump(include=vals_to_keep, exclude_none=True)
for i, rule in enumerate(self._acl):
if not rule:
continue
self._default_config[i] = {"action": rule.action.name}
if rule.src_ip_address:
self._default_config[i]["src_ip"] = str(rule.src_ip_address)
if rule.dst_ip_address:
self._default_config[i]["dst_ip"] = str(rule.dst_ip_address)
if rule.src_port:
self._default_config[i]["src_port"] = rule.src_port.name
if rule.dst_port:
self._default_config[i]["dst_port"] = rule.dst_port.name
if rule.protocol:
self._default_config[i]["protocol"] = rule.protocol.name
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.implicit_rule.reset_component_for_episode(episode)
super().reset_component_for_episode(episode)
self._reset_rules_to_default()
def _reset_rules_to_default(self) -> None:
"""Clear all ACL rules and set them to the default rules config."""
self._acl = [None] * (self.max_acl_rules - 1)
for r_num, r_cfg in self._default_config.items():
self.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
position=r_num,
)
def _init_request_manager(self) -> RequestManager:
# TODO: Add src and dst wildcard masks as positional args in this request.
@@ -616,11 +569,6 @@ class RouteEntry(SimComponent):
metric: float = 0.0
"The cost metric for this route. Default is 0.0."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"address", "subnet_mask", "next_hop_ip_address", "metric"}
self._original_values = self.model_dump(include=vals_to_include)
def describe_state(self) -> Dict:
"""
Describes the current state of the RouteEntry.
@@ -653,17 +601,6 @@ class RouteTable(SimComponent):
default_route: Optional[RouteEntry] = None
sys_log: SysLog
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
self._original_state["routes_orig"] = self.routes
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.routes.clear()
self.routes = self._original_state["routes_orig"]
super().reset_component_for_episode(episode)
def describe_state(self) -> Dict:
"""
Describes the current state of the RouteTable.
@@ -1104,8 +1041,6 @@ class Router(NetworkNode):
self._set_default_acl()
self.set_original_state()
def _install_system_software(self):
"""
Installs essential system software and network services on the router.
@@ -1131,20 +1066,7 @@ class Router(NetworkNode):
self.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
self.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
def set_original_state(self):
"""
Sets or resets the router to its original configuration state.
This method is called to initialize the router's state or to revert it to a known good configuration during
network simulations or after configuration changes.
"""
self.acl.set_original_state()
self.route_table.set_original_state()
super().set_original_state()
vals_to_include = {"num_ports"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""
Resets the router's components for a new network simulation episode.
@@ -1154,13 +1076,10 @@ class Router(NetworkNode):
:param episode: The episode number for which the router is being reset.
"""
self.software_manager.arp.clear()
self.acl.reset_component_for_episode(episode)
self.route_table.reset_component_for_episode(episode)
for i, network_interface in self.network_interface.items():
network_interface.reset_component_for_episode(episode)
for i, _ in self.network_interface.items():
self.enable_port(i)
super().reset_component_for_episode(episode)
super().setup_for_episode(episode=episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
@@ -1391,7 +1310,6 @@ class Router(NetworkNode):
network_interface.ip_address = ip_address
network_interface.subnet_mask = subnet_mask
self.sys_log.info(f"Configured Network Interface {network_interface}")
self.set_original_state()
def enable_port(self, port: int):
"""
@@ -1493,6 +1411,14 @@ class Router(NetworkNode):
subnet_mask=port_cfg["subnet_mask"],
)
if "acl" in cfg:
new.acl._default_config = cfg["acl"] # save the config to allow resetting
new.acl._reset_rules_to_default() # read the config and apply rules
for r_num, r_cfg in cfg["acl"].items():
new.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else Port[p],
dst_port=None if not (p := r_cfg.get("dst_port")) else Port[p],
protocol=None if not (p := r_cfg.get("protocol")) else IPProtocol[p],
src_ip_address=r_cfg.get("src_ip"),
dst_ip_address=r_cfg.get("dst_ip"),
position=r_num,
)
return new

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

@@ -32,12 +32,6 @@ class SwitchPort(WiredNetworkInterface):
_connected_node: Optional[Switch] = None
"The Switch to which the SwitchPort is connected."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {"port_num", "mac_address", "speed", "mtu", "enabled"}
self._original_state = self.model_dump(include=vals_to_include)
super().set_original_state()
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.

3
src/primaite/simulator/network/hardware/nodes/network/wireless_router.py
View File

@@ -122,8 +122,6 @@ class WirelessRouter(Router):
self.connect_nic(RouterInterface(ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0"))
self.set_original_state()
@property
def wireless_access_point(self) -> WirelessAccessPoint:
"""
@@ -166,7 +164,6 @@ class WirelessRouter(Router):
network_interface.ip_address = ip_address
network_interface.subnet_mask = subnet_mask
self.sys_log.info(f"Configured WAP {network_interface}")
self.set_original_state()
self.wireless_access_point.frequency = frequency # Set operating frequency
self.wireless_access_point.enable() # Re-enable the WAP with new settings

47
src/primaite/simulator/network/nmne.py Normal file
View File

@@ -0,0 +1,47 @@
from typing import Dict, Final, List
CAPTURE_NMNE: bool = True
"""Indicates whether Malicious Network Events (MNEs) should be captured. Default is True."""
NMNE_CAPTURE_KEYWORDS: List[str] = []
"""List of keywords to identify malicious network events."""
# TODO: Remove final and make configurable after example layout when the NicObservation creates nmne structure dynamically
CAPTURE_BY_DIRECTION: Final[bool] = True
"""Flag to determine if captures should be organized by traffic direction (inbound/outbound)."""
CAPTURE_BY_IP_ADDRESS: Final[bool] = False
"""Flag to determine if captures should be organized by source or destination IP address."""
CAPTURE_BY_PROTOCOL: Final[bool] = False
"""Flag to determine if captures should be organized by network protocol (e.g., TCP, UDP)."""
CAPTURE_BY_PORT: Final[bool] = False
"""Flag to determine if captures should be organized by source or destination port."""
CAPTURE_BY_KEYWORD: Final[bool] = False
"""Flag to determine if captures should be filtered and categorised based on specific keywords."""
def set_nmne_config(nmne_config: Dict):
"""
Sets the configuration for capturing Malicious Network Events (MNEs) based on a provided dictionary.
This function updates global settings related to NMNE capture, including whether to capture NMNEs and what
keywords to use for identifying NMNEs.
The function ensures that the settings are updated only if they are provided in the `nmne_config` dictionary,
and maintains type integrity by checking the types of the provided values.
:param nmne_config: A dictionary containing the NMNE configuration settings. Possible keys include:
"capture_nmne" (bool) to indicate whether NMNEs should be captured, "nmne_capture_keywords" (list of strings)
to specify keywords for NMNE identification.
"""
global NMNE_CAPTURE_KEYWORDS
global CAPTURE_NMNE
# Update the NMNE capture flag, defaulting to False if not specified or if the type is incorrect
CAPTURE_NMNE = nmne_config.get("capture_nmne", False)
if not isinstance(CAPTURE_NMNE, bool):
CAPTURE_NMNE = True # Revert to default True if the provided value is not a boolean
# Update the NMNE capture keywords, appending new keywords if provided
NMNE_CAPTURE_KEYWORDS += nmne_config.get("nmne_capture_keywords", [])
if not isinstance(NMNE_CAPTURE_KEYWORDS, list):
NMNE_CAPTURE_KEYWORDS = [] # Reset to empty list if the provided value is not a list

8
src/primaite/simulator/sim_container.py
View File

@@ -21,13 +21,9 @@ class Simulation(SimComponent):
super().__init__(**kwargs)
def set_original_state(self):
"""Sets the original state."""
self.network.set_original_state()
def reset_component_for_episode(self, episode: int):
def setup_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.network.reset_component_for_episode(episode)
self.network.setup_for_episode(episode=episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()

6
src/primaite/simulator/system/applications/application.py
View File

@@ -38,12 +38,6 @@ class Application(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
vals_to_include = {"operating_state", "execution_control_status", "num_executions", "groups"}
self._original_state.update(self.model_dump(include=vals_to_include))
@abstractmethod
def describe_state(self) -> Dict:
"""

14
src/primaite/simulator/system/applications/database_client.py
View File

@@ -34,7 +34,6 @@ class DatabaseClient(Application):
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
self.set_original_state()
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
@@ -50,19 +49,6 @@ class DatabaseClient(Application):
self._connections_status.append(can_connect)
return can_connect
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting DatabaseClient WebServer original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"server_ip_address", "server_password", "connected", "_query_success_tracker"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting DataBaseClient state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
self._query_success_tracker.clear()
def describe_state(self) -> Dict:
"""
Describes the current state of the ACLRule.

20
src/primaite/simulator/system/applications/red_applications/data_manipulation_bot.py
View File

@@ -49,26 +49,6 @@ class DataManipulationBot(DatabaseClient):
super().__init__(**kwargs)
self.name = "DataManipulationBot"
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting DataManipulationBot original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {
"server_ip_address",
"payload",
"server_password",
"port_scan_p_of_success",
"data_manipulation_p_of_success",
"attack_stage",
"repeat",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting DataManipulationBot state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()

21
src/primaite/simulator/system/applications/red_applications/dos_bot.py
View File

@@ -57,27 +57,6 @@ class DoSBot(DatabaseClient, Application):
self.name = "DoSBot"
self.max_sessions = 1000 # override normal max sessions
def set_original_state(self):
"""Set the original state of the Denial of Service Bot."""
_LOGGER.debug(f"Setting {self.name} original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {
"target_ip_address",
"target_port",
"payload",
"repeat",
"attack_stage",
"max_sessions",
"port_scan_p_of_success",
"dos_intensity",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting {self.name} state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()

16
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -47,21 +47,8 @@ class WebBrowser(Application):
kwargs["port"] = Port.HTTP
super().__init__(**kwargs)
self.set_original_state()
self.run()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting WebBrowser original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"target_url", "domain_name_ip_address", "latest_response"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting WebBrowser state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(
@@ -80,9 +67,6 @@ class WebBrowser(Application):
state["history"] = [hist_item.state() for hist_item in self.history]
return state
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
def get_webpage(self, url: Optional[str] = None) -> bool:
"""
Retrieve the webpage.

6
src/primaite/simulator/system/processes/process.py
View File

@@ -24,12 +24,6 @@ class Process(Software):
operating_state: ProcessOperatingState
"The current operating state of the Process."
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
vals_to_include = {"operating_state"}
self._original_state.update(self.model_dump(include=vals_to_include))
@abstractmethod
def describe_state(self) -> Dict:
"""

19
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -40,25 +40,6 @@ class DatabaseService(Service):
super().__init__(**kwargs)
self._create_db_file()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting DatabaseService original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {
"password",
"connections",
"backup_server_ip",
"latest_backup_directory",
"latest_backup_file_name",
}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug("Resetting DatabaseService original state on node {self.software_manager.node.hostname}")
self.clear_connections()
super().reset_component_for_episode(episode)
def configure_backup(self, backup_server: IPv4Address):
"""
Set up the database backup.

12
src/primaite/simulator/system/services/dns/dns_client.py
View File

@@ -29,18 +29,6 @@ class DNSClient(Service):
super().__init__(**kwargs)
self.start()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting DNSClient original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"dns_server"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.dns_cache.clear()
super().reset_component_for_episode(episode)
def describe_state(self) -> Dict:
"""
Describes the current state of the software.

14
src/primaite/simulator/system/services/dns/dns_server.py
View File

@@ -28,20 +28,6 @@ class DNSServer(Service):
super().__init__(**kwargs)
self.start()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting DNSServer original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"dns_table"}
self._original_state["dns_table_orig"] = self.model_dump(include=vals_to_include)["dns_table"]
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
self.dns_table.clear()
for key, value in self._original_state["dns_table_orig"].items():
self.dns_table[key] = value
super().reset_component_for_episode(episode)
def describe_state(self) -> Dict:
"""
Describes the current state of the software.

12
src/primaite/simulator/system/services/ftp/ftp_client.py
View File

@@ -27,18 +27,6 @@ class FTPClient(FTPServiceABC):
super().__init__(**kwargs)
self.start()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting FTPClient original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"connected"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting FTPClient state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
"""
Process the command in the FTP Packet.

13
src/primaite/simulator/system/services/ftp/ftp_server.py
View File

@@ -27,19 +27,6 @@ class FTPServer(FTPServiceABC):
super().__init__(**kwargs)
self.start()
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting FTPServer original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"server_password"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting FTPServer state on node {self.software_manager.node.hostname}")
self.clear_connections()
super().reset_component_for_episode(episode)
def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
"""
Process the command in the FTP Packet.

11
src/primaite/simulator/system/services/ntp/ntp_client.py
View File

@@ -49,21 +49,12 @@ class NTPClient(Service):
state = super().describe_state()
return state
def reset_component_for_episode(self, episode: int):
"""
Resets the Service component for a new episode.
This method ensures the Service is ready for a new episode, including resetting any
stateful properties or statistics, and clearing any message queues.
"""
pass
def send(
self,
payload: NTPPacket,
session_id: Optional[str] = None,
dest_ip_address: IPv4Address = None,
dest_port: [Port] = Port.NTP,
dest_port: Port = Port.NTP,
**kwargs,
) -> bool:
"""Requests NTP data from NTP server.

10
src/primaite/simulator/system/services/ntp/ntp_server.py
View File

@@ -34,16 +34,6 @@ class NTPServer(Service):
state = super().describe_state()
return state
def reset_component_for_episode(self, episode: int):
"""
Resets the Service component for a new episode.
This method ensures the Service is ready for a new episode, including
resetting any stateful properties or statistics, and clearing any message
queues.
"""
pass
def receive(
self,
payload: NTPPacket,

6
src/primaite/simulator/system/services/service.py
View File

@@ -78,12 +78,6 @@ class Service(IOSoftware):
"""
return super().receive(payload=payload, session_id=session_id, **kwargs)
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
vals_to_include = {"operating_state", "restart_duration", "restart_countdown"}
self._original_state.update(self.model_dump(include=vals_to_include))
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request("scan", RequestType(func=lambda request, context: self.scan()))

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

@@ -23,18 +23,6 @@ class WebServer(Service):
last_response_status_code: Optional[HttpStatusCode] = None
def set_original_state(self):
"""Sets the original state."""
_LOGGER.debug(f"Setting WebServer original state on node {self.software_manager.node.hostname}")
super().set_original_state()
vals_to_include = {"last_response_status_code"}
self._original_state.update(self.model_dump(include=vals_to_include))
def reset_component_for_episode(self, episode: int):
"""Reset the original state of the SimComponent."""
_LOGGER.debug(f"Resetting WebServer state on node {self.software_manager.node.hostname}")
super().reset_component_for_episode(episode)
def describe_state(self) -> Dict:
"""
Produce a dictionary describing the current state of this object.
@@ -130,7 +118,7 @@ class WebServer(Service):
self.set_health_state(SoftwareHealthState.COMPROMISED)
return response
except Exception:
except Exception: # TODO: refactor this. Likely to cause silent bugs. (ADO ticket #2345 )
# something went wrong on the server
response.status_code = HttpStatusCode.INTERNAL_SERVER_ERROR
return response

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

@@ -3,7 +3,7 @@ from abc import abstractmethod
from datetime import datetime
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Optional, TYPE_CHECKING, Union
from primaite.simulator.core import _LOGGER, RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file_system import FileSystem, Folder
@@ -13,6 +13,9 @@ from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.session_manager import Session
from primaite.simulator.system.core.sys_log import SysLog
if TYPE_CHECKING:
from primaite.simulator.system.core.software_manager import SoftwareManager
class SoftwareType(Enum):
"""
@@ -84,7 +87,7 @@ class Software(SimComponent):
"The count of times the software has been scanned, defaults to 0."
revealed_to_red: bool = False
"Indicates if the software has been revealed to red agent, defaults is False."
software_manager: Any = None
software_manager: "SoftwareManager" = None
"An instance of Software Manager that is used by the parent node."
sys_log: SysLog = None
"An instance of SysLog that is used by the parent node."
@@ -97,19 +100,6 @@ class Software(SimComponent):
_patching_countdown: Optional[int] = None
"Current number of ticks left to patch the software."
def set_original_state(self):
"""Sets the original state."""
vals_to_include = {
"name",
"health_state_actual",
"health_state_visible",
"criticality",
"patching_count",
"scanning_count",
"revealed_to_red",
}
self._original_state = self.model_dump(include=vals_to_include)
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(
@@ -248,12 +238,6 @@ class IOSoftware(Software):
_connections: Dict[str, Dict] = {}
"Active connections."
def set_original_state(self):
"""Sets the original state."""
super().set_original_state()
vals_to_include = {"installing_count", "max_sessions", "tcp", "udp", "port"}
self._original_state.update(self.model_dump(include=vals_to_include))
@abstractmethod
def describe_state(self) -> Dict:
"""