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

Merge branch '4.0.0a1-dev' into feature/3075_Migrate_notebooks_to_MilPac_(Core_changes)

Browse Source
This commit is contained in:
Nick Todd
2025-02-06 09:17:34 +00:00
parent 7b2a9c1d2a 0d1edf0362
commit 0d8672cdba
210 changed files with 3022 additions and 2983 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -24,7 +24,7 @@ game:
agents:
- ref: client_2_green_user
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.3
@@ -34,33 +34,33 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client_2
application_name: WebBrowser
application_name: web-browser
2:
action: node_application_execute
action: node-application-execute
options:
node_name: client_2
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
weight: 0.25
options:
node_hostname: client_2
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 0.05
options:
node_hostname: client_2
- ref: client_1_green_user
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.3
@@ -70,26 +70,26 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client_1
application_name: WebBrowser
application_name: web-browser
2:
action: node_application_execute
action: node-application-execute
options:
node_name: client_1
application_name: DatabaseClient
application_name: web-browser
reward_function:
reward_components:
- type: WEBPAGE_UNAVAILABLE_PENALTY
- type: webpage-unavailable-penalty
weight: 0.25
options:
node_hostname: client_1
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 0.05
options:
node_hostname: client_1
@@ -100,31 +100,31 @@ agents:
- ref: data_manipulation_attacker
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client_1, client_2]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 25
frequency: 20
variance: 5
- ref: defender
team: BLUE
type: ProxyAgent
type: proxy-agent
observation_space:
type: CUSTOM
type: custom
options:
components:
- type: NODES
- type: nodes
label: NODES
options:
hosts:
- hostname: domain_controller
- hostname: web_server
services:
- service_name: WebServer
- service_name: web-server
- hostname: database_server
folders:
- folder_name: database
@@ -169,7 +169,7 @@ agents:
- UDP
num_rules: 10
- type: LINKS
- type: links
label: LINKS
options:
link_references:
@@ -183,222 +183,222 @@ agents:
- switch_2:eth-1<->client_1:eth-1
- switch_2:eth-2<->client_2:eth-1
- switch_2:eth-7<->security_suite:eth-2
- type: "NONE"
- type: "none"
label: ICS
options: {}
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
# scan webapp service
1:
action: node_service_scan
action: node-service-scan
options:
node_name: web_server
service_name: WebServer
service_name: web-server
# stop webapp service
2:
action: node_service_stop
action: node-service-stop
options:
node_name: web_server
service_name: WebServer
service_name: web-server
# start webapp service
3:
action: "node_service_start"
action: "node-service-start"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
4:
action: "node_service_pause"
action: "node-service-pause"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
5:
action: "node_service_resume"
action: "node-service-resume"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
6:
action: "node_service_restart"
action: "node-service-restart"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
7:
action: "node_service_disable"
action: "node-service-disable"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
8:
action: "node_service_enable"
action: "node-service-enable"
options:
node_name: web_server
service_name: WebServer
service_name: web-server
9: # check database.db file
action: "node_file_scan"
action: "node-file-scan"
options:
node_name: database_server
folder_name: database
file_name: database.db
10:
action: "node_file_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
action: "node-file-scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
options:
node_name: database_server
folder_name: database
file_name: database.db
11:
action: "node_file_delete"
action: "node-file-delete"
options:
node_name: database_server
folder_name: database
file_name: database.db
12:
action: "node_file_repair"
action: "node-file-repair"
options:
node_name: database_server
folder_name: database
file_name: database.db
13:
action: "node_service_fix"
action: "node-service-fix"
options:
node_name: database_server
service_name: DatabaseService
service_name: database-service
14:
action: "node_folder_scan"
action: "node-folder-scan"
options:
node_name: database_server
folder_name: database
15:
action: "node_folder_scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
action: "node-folder-scan" # CHECKHASH replaced by SCAN - but the behaviour is the same in this context.
options:
node_name: database_server
folder_name: database
16:
action: "node_folder_repair"
action: "node-folder-repair"
options:
node_name: database_server
folder_name: database
17:
action: "node_folder_restore"
action: "node-folder-restore"
options:
node_name: database_server
folder_name: database
18:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: domain_controller
19:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: domain_controller
20:
action: node_startup
action: node-startup
options:
node_name: domain_controller
21:
action: node_reset
action: node-reset
options:
node_name: domain_controller
22:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: web_server
23:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: web_server
24:
action: node_startup
action: node-startup
options:
node_name: web_server
25:
action: node_reset
action: node-reset
options:
node_name: web_server
26: # old action num: 18
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: database_server
27:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: database_server
28:
action: node_startup
action: node-startup
options:
node_name: database_server
29:
action: node_reset
action: node-reset
options:
node_name: database_server
30:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: backup_server
31:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: backup_server
32:
action: node_startup
action: node-startup
options:
node_name: backup_server
33:
action: node_reset
action: node-reset
options:
node_name: backup_server
34:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: security_suite
35:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: security_suite
36:
action: node_startup
action: node-startup
options:
node_name: security_suite
37:
action: node_reset
action: node-reset
options:
node_name: security_suite
38:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: client_1
39: # old action num: 19 # shutdown client 1
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: client_1
40: # old action num: 20
action: node_startup
action: node-startup
options:
node_name: client_1
41: # old action num: 21
action: node_reset
action: node-reset
options:
node_name: client_1
42:
action: "node_os_scan"
action: "node-os-scan"
options:
node_name: client_2
43:
action: "node_shutdown"
action: "node-shutdown"
options:
node_name: client_2
44:
action: node_startup
action: node-startup
options:
node_name: client_2
45:
action: node_reset
action: node-reset
options:
node_name: client_2
46: # old action num: 22 # "ACL: ADDRULE - Block outgoing traffic from client 1"
action: "router_acl_add_rule"
46: # old action num: 22 # "acl: ADDRULE - Block outgoing traffic from client 1"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 1
@@ -410,8 +410,8 @@ agents:
protocol_name: ALL
src_wildcard: NONE
dst_wildcard: NONE
47: # old action num: 23 # "ACL: ADDRULE - Block outgoing traffic from client 2"
action: "router_acl_add_rule"
47: # old action num: 23 # "acl: ADDRULE - Block outgoing traffic from client 2"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 2
@@ -424,7 +424,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
48: # old action num: 24 # block tcp traffic from client 1 to web app
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 3
@@ -437,7 +437,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
49: # old action num: 25 # block tcp traffic from client 2 to web app
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 4
@@ -450,7 +450,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
50: # old action num: 26
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 5
@@ -463,7 +463,7 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
51: # old action num: 27
action: "router_acl_add_rule"
action: "router-acl-add-rule"
options:
target_router: router_1
position: 6
@@ -476,132 +476,132 @@ agents:
src_wildcard: NONE
dst_wildcard: NONE
52: # old action num: 28
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 0
53: # old action num: 29
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 1
54: # old action num: 30
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 2
55: # old action num: 31
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 3
56: # old action num: 32
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 4
57: # old action num: 33
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 5
58: # old action num: 34
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 6
59: # old action num: 35
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 7
60: # old action num: 36
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 8
61: # old action num: 37
action: "router_acl_remove_rule"
action: "router-acl-remove-rule"
options:
target_router: router_1
position: 9
62: # old action num: 38
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: domain_controller
nic_num: 1
63: # old action num: 39
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: domain_controller
nic_num: 1
64: # old action num: 40
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: web_server
nic_num: 1
65: # old action num: 41
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: web_server
nic_num: 1
66: # old action num: 42
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: database_server
nic_num: 1
67: # old action num: 43
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: database_server
nic_num: 1
68: # old action num: 44
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: backup_server
nic_num: 1
69: # old action num: 45
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: backup_server
nic_num: 1
70: # old action num: 46
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: security_suite
nic_num: 1
71: # old action num: 47
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: security_suite
nic_num: 1
72: # old action num: 48
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: security_suite
nic_num: 2
73: # old action num: 49
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: security_suite
nic_num: 2
74: # old action num: 50
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: client_1
nic_num: 1
75: # old action num: 51
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: client_1
nic_num: 1
76: # old action num: 52
action: "host_nic_disable"
action: "host-nic-disable"
options:
node_name: client_2
nic_num: 1
77: # old action num: 53
action: "host_nic_enable"
action: "host-nic-enable"
options:
node_name: client_2
nic_num: 1
@@ -611,19 +611,19 @@ agents:
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.40
options:
node_hostname: database_server
folder_name: database
file_name: database.db
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: client_1_green_user
- type: SHARED_REWARD
- type: shared-reward
weight: 1.0
options:
agent_name: client_2_green_user
@@ -693,7 +693,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
services:
- type: DNSServer
- type: dns-server
options:
domain_mapping:
arcd.com: 192.168.1.12 # web server
@@ -705,9 +705,9 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: WebServer
- type: web-server
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
@@ -719,10 +719,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: DatabaseService
- type: database-service
options:
backup_server_ip: 192.168.1.16
- type: FTPClient
- type: ftp-client
- hostname: backup_server
type: server
@@ -731,7 +731,7 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 192.168.1.10
services:
- type: FTPServer
- type: ftp-server
- hostname: security_suite
type: server
@@ -751,20 +751,20 @@ simulation:
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- type: DataManipulationBot
- type: data-manipulation-bot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- type: WebBrowser
- type: web-browser
options:
target_url: http://arcd.com/users/
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
services:
- type: DNSClient
- type: dns-client
- hostname: client_2
type: computer
@@ -773,20 +773,20 @@ simulation:
default_gateway: 192.168.10.1
dns_server: 192.168.1.10
applications:
- type: WebBrowser
- type: web-browser
options:
target_url: http://arcd.com/users/
- type: DataManipulationBot
- type: data-manipulation-bot
options:
port_scan_p_of_success: 0.8
data_manipulation_p_of_success: 0.8
payload: "DELETE"
server_ip: 192.168.1.14
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.14
services:
- type: DNSClient
- type: dns-client
links:
- endpoint_a_hostname: router_1

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

File diff suppressed because it is too large Load Diff

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

@@ -5,46 +5,46 @@ game:
agents:
- ref: RL_Agent
type: ProxyAgent
type: proxy-agent
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_shutdown
action: node-shutdown
options:
node_name: client_1
2:
action: node_shutdown
action: node-shutdown
options:
node_name: server
3:
action: node_startup
action: node-startup
options:
node_name: client_1
4:
action: node_startup
action: node-startup
options:
node_name: server
5:
action: host_nic_disable
action: host-nic-disable
options:
node_name: client_1
nic_num: 1
6:
action: host_nic_disable
action: host-nic-disable
options:
node_name: server
nic_num: 1
7:
action: host_nic_enable
action: host-nic-enable
options:
node_name: client_1
nic_num: 1
8:
action: host_nic_enable
action: host-nic-enable
options:
node_name: server
nic_num: 1

4
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_1.yaml
View File

@@ -1,5 +1,5 @@
server_services: &server_services
- type: DatabaseService
- type: database-service
client_applications: &client_applications
- type: DatabaseClient
- type: database-client

4
src/primaite/config/_package_data/mini_scenario_with_simulation_variation/simulation_variant_2.yaml
View File

@@ -1,5 +1,5 @@
server_services: &server_services
- type: FTPServer
- type: ftp-server
client_applications: &client_applications
- type: RansomwareScript
- type: ransomware-script

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

@@ -20,10 +20,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -34,10 +34,10 @@ simulation:
default_gateway: 192.168.1.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -102,8 +102,7 @@ simulation:
subnet_mask: 255.255.255.252
default_gateway: 8.8.8.1
services:
- ref: dns_server
type: DNSServer
- type: dns-server
options:
domain_mapping:
sometech.ai: 94.10.180.6
@@ -150,7 +149,7 @@ simulation:
dst_ip: 94.10.180.6
dst_port: POSTGRES_SERVER
dst_wildcard_mask: 0.0.0.0
8: # Permit SomeTech DMZ to use ARP
8: # Permit SomeTech DMZ to use arp
action: PERMIT
src_port: ARP
dst_port: ARP
@@ -170,7 +169,7 @@ simulation:
dst_ip: 10.10.1.11
dst_port: POSTGRES_SERVER
dst_wildcard_mask: 0.0.0.0
8: # Permit SomeTech DMZ to use ARP
8: # Permit SomeTech DMZ to use arp
action: PERMIT
src_port: ARP
dst_port: ARP
@@ -196,10 +195,9 @@ simulation:
default_gateway: 94.10.180.5
dns_server: 8.8.8.2
services:
- ref: web_server
type: WebServer
- type: web-server
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
@@ -269,12 +267,12 @@ simulation:
action: PERMIT
src_port: HTTP
dst_port: HTTP
18: # Allow the SomeTech internal network to use ARP
18: # Allow the SomeTech internal network to use arp
action: PERMIT
src_ip: 10.10.0.0
src_wildcard_mask: 0.0.255.255
src_port: ARP
19: # Allow the SomeTech internal network to use ICMP
19: # Allow the SomeTech internal network to use icmp
action: PERMIT
src_ip: 10.10.0.0
src_wildcard_mask: 0.0.255.255
@@ -318,10 +316,10 @@ simulation:
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: DatabaseService
- type: database-service
options:
backup_server_ip: 10.10.1.12 # The some_tech_storage_srv server
- type: FTPClient
- type: ftp-client
- hostname: some_tech_storage_srv
type: server
@@ -330,7 +328,7 @@ simulation:
default_gateway: 10.10.1.1
dns_server: 8.8.8.2
services:
- type: FTPServer
- type: ftp-server
- hostname: some_tech_hr_1
type: computer
@@ -339,10 +337,10 @@ simulation:
default_gateway: 10.10.3.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -353,10 +351,10 @@ simulation:
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/
@@ -367,10 +365,10 @@ simulation:
default_gateway: 10.10.2.1
dns_server: 8.8.8.2
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 10.10.1.11
- type: WebBrowser
- type: web-browser
options:
target_url: http://sometech.ai/users/

10
src/primaite/config/_package_data/scenario_with_placeholders/greens_1.yaml
View File

@@ -1,7 +1,7 @@
agents: &greens
- ref: green_A
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.2
@@ -10,17 +10,17 @@ agents: &greens
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 1.0
options:
node_hostname: client

10
src/primaite/config/_package_data/scenario_with_placeholders/greens_2.yaml
View File

@@ -1,7 +1,7 @@
agents: &greens
- ref: green_B
team: GREEN
type: ProbabilisticAgent
type: probabilistic-agent
agent_settings:
action_probabilities:
0: 0.95
@@ -10,17 +10,17 @@ agents: &greens
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_application_execute
action: node-application-execute
options:
node_name: client
application_name: DatabaseClient
application_name: database-client
reward_function:
reward_components:
- type: GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY
- type: green-admin-database-unreachable-penalty
weight: 1.0
options:
node_hostname: client

4
src/primaite/config/_package_data/scenario_with_placeholders/reds_1.yaml
View File

@@ -1,11 +1,11 @@
reds: &reds
- ref: red_A
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client,]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 10
frequency: 10
variance: 0

4
src/primaite/config/_package_data/scenario_with_placeholders/reds_2.yaml
View File

@@ -1,11 +1,11 @@
reds: &reds
- ref: red_B
team: RED
type: RedDatabaseCorruptingAgent
type: red-database-corrupting-agent
agent_settings:
possible_start_nodes: [client_1]
target_application: DataManipulationBot
target_application: data-manipulation-bot
start_step: 3
frequency: 2
variance: 1

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

@@ -26,12 +26,12 @@ agents:
- ref: defender
team: BLUE
type: ProxyAgent
type: proxy-agent
observation_space:
type: CUSTOM
type: custom
options:
components:
- type: NODES
- type: nodes
label: NODES
options:
routers: []
@@ -46,7 +46,7 @@ agents:
include_num_access: false
include_nmne: true
- type: LINKS
- type: links
label: LINKS
options:
link_references:
@@ -56,48 +56,48 @@ agents:
action_space:
action_map:
0:
action: do_nothing
action: do-nothing
options: {}
1:
action: node_shutdown
action: node-shutdown
options:
node_name: client_1
2:
action: node_shutdown
action: node-shutdown
options:
node_name: server
3:
action: node_startup
action: node-startup
options:
node_name: client_1
4:
action: node_startup
action: node-startup
options:
node_name: server
5:
action: host_nic_disable
action: host-nic-disable
options:
node_name: client_1
nic_num: 1
6:
action: host_nic_disable
action: host-nic-disable
options:
node_name: server
nic_num: 1
7:
action: host_nic_enable
action: host-nic-enable
options:
node_name: client_1
nic_num: 1
8:
action: host_nic_enable
action: host-nic-enable
options:
node_name: server
nic_num: 1
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.40
options:
node_hostname: database_server
@@ -121,10 +121,10 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
applications:
- type: DatabaseClient
- type: database-client
options:
db_server_ip: 192.168.1.3
- type: DataManipulationBot
- type: data-manipulation-bot
options:
server_ip: 192.168.1.3
payload: "DELETE"
@@ -139,7 +139,7 @@ simulation:
subnet_mask: 255.255.255.0
default_gateway: 192.168.1.1
services:
- type: DatabaseService
- type: database-service
links:
- endpoint_a_hostname: client

19
src/primaite/game/agent/actions/abstract.py
View File

@@ -12,8 +12,6 @@ from primaite.interface.request import RequestFormat
class AbstractAction(BaseModel, ABC):
"""Base class for actions."""
config: "AbstractAction.ConfigSchema"
class ConfigSchema(BaseModel, ABC):
"""Base configuration schema for Actions."""
@@ -22,13 +20,20 @@ class AbstractAction(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAction]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an action type.
:param discriminator: discriminator used to uniquely specify action types.
:type discriminator: str
:raises ValueError: When attempting to create an action with a name that is already in use.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Cannot create new action under reserved name {discriminator}")
cls._registry[discriminator] = cls
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:

28
src/primaite/game/agent/actions/acl.py
View File

@@ -21,9 +21,7 @@ __all__ = (
class ACLAddRuleAbstractAction(AbstractAction, ABC):
"""Base abstract class for ACL add rule actions."""
config: ConfigSchema = "ACLAddRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Configuration Schema base for ACL add rule abstract actions."""
src_ip: Union[IPV4Address, Literal["ALL"]]
@@ -37,19 +35,17 @@ class ACLAddRuleAbstractAction(AbstractAction, ABC):
dst_wildcard: Union[IPV4Address, Literal["NONE"]]
class ACLRemoveRuleAbstractAction(AbstractAction, identifier="acl_remove_rule_abstract_action"):
"""Base abstract class for ACL remove rule actions."""
class ACLRemoveRuleAbstractAction(AbstractAction, ABC):
"""Base abstract class for acl remove rule actions."""
config: ConfigSchema = "ACLRemoveRuleAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Configuration Schema base for ACL remove rule abstract actions."""
position: int
class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_add_rule"):
"""Action which adds a rule to a router's ACL."""
class RouterACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="router-acl-add-rule"):
"""Action which adds a rule to a router's acl."""
config: "RouterACLAddRuleAction.ConfigSchema"
@@ -79,8 +75,8 @@ class RouterACLAddRuleAction(ACLAddRuleAbstractAction, identifier="router_acl_ad
]
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_acl_remove_rule"):
"""Action which removes a rule from a router's ACL."""
class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="router-acl-remove-rule"):
"""Action which removes a rule from a router's acl."""
config: "RouterACLRemoveRuleAction.ConfigSchema"
@@ -95,8 +91,8 @@ class RouterACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="router_
return ["network", "node", config.target_router, "acl", "remove_rule", config.position]
class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_acl_add_rule"):
"""Action which adds a rule to a firewall port's ACL."""
class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, discriminator="firewall-acl-add-rule"):
"""Action which adds a rule to a firewall port's acl."""
config: "FirewallACLAddRuleAction.ConfigSchema"
@@ -130,8 +126,8 @@ class FirewallACLAddRuleAction(ACLAddRuleAbstractAction, identifier="firewall_ac
]
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, identifier="firewall_acl_remove_rule"):
"""Action which removes a rule from a firewall port's ACL."""
class FirewallACLRemoveRuleAction(ACLRemoveRuleAbstractAction, discriminator="firewall-acl-remove-rule"):
"""Action which removes a rule from a firewall port's acl."""
config: "FirewallACLRemoveRuleAction.ConfigSchema"

16
src/primaite/game/agent/actions/application.py
View File

@@ -23,9 +23,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
inherit from this base class.
"""
config: "NodeApplicationAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for Node Application actions."""
node_name: str
@@ -45,7 +43,7 @@ class NodeApplicationAbstractAction(AbstractAction, ABC):
]
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="node_application_execute"):
class NodeApplicationExecuteAction(NodeApplicationAbstractAction, discriminator="node-application-execute"):
"""Action which executes an application."""
config: "NodeApplicationExecuteAction.ConfigSchema"
@@ -56,7 +54,7 @@ class NodeApplicationExecuteAction(NodeApplicationAbstractAction, identifier="no
verb: str = "execute"
class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_application_scan"):
class NodeApplicationScanAction(NodeApplicationAbstractAction, discriminator="node-application-scan"):
"""Action which scans an application."""
config: "NodeApplicationScanAction.ConfigSchema"
@@ -67,7 +65,7 @@ class NodeApplicationScanAction(NodeApplicationAbstractAction, identifier="node_
verb: str = "scan"
class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node_application_close"):
class NodeApplicationCloseAction(NodeApplicationAbstractAction, discriminator="node-application-close"):
"""Action which closes an application."""
config: "NodeApplicationCloseAction.ConfigSchema"
@@ -78,7 +76,7 @@ class NodeApplicationCloseAction(NodeApplicationAbstractAction, identifier="node
verb: str = "close"
class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_application_fix"):
class NodeApplicationFixAction(NodeApplicationAbstractAction, discriminator="node-application-fix"):
"""Action which fixes an application."""
config: "NodeApplicationFixAction.ConfigSchema"
@@ -89,7 +87,7 @@ class NodeApplicationFixAction(NodeApplicationAbstractAction, identifier="node_a
verb: str = "fix"
class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="node_application_install"):
class NodeApplicationInstallAction(NodeApplicationAbstractAction, discriminator="node-application-install"):
"""Action which installs an application."""
config: "NodeApplicationInstallAction.ConfigSchema"
@@ -113,7 +111,7 @@ class NodeApplicationInstallAction(NodeApplicationAbstractAction, identifier="no
]
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, identifier="node_application_remove"):
class NodeApplicationRemoveAction(NodeApplicationAbstractAction, discriminator="node-application-remove"):
"""Action which removes/uninstalls an application."""
config: "NodeApplicationRemoveAction.ConfigSchema"

28
src/primaite/game/agent/actions/file.py
View File

@@ -24,9 +24,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
only three parameters can inherit from this base class.
"""
config: "NodeFileAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Configuration Schema for NodeFileAbstractAction."""
node_name: str
@@ -38,7 +36,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None or config.file_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -52,7 +50,7 @@ class NodeFileAbstractAction(AbstractAction, ABC):
]
class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"):
class NodeFileCreateAction(NodeFileAbstractAction, discriminator="node-file-create"):
"""Action which creates a new file in a given folder."""
config: "NodeFileCreateAction.ConfigSchema"
@@ -67,7 +65,7 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None or config.file_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -81,7 +79,7 @@ class NodeFileCreateAction(NodeFileAbstractAction, identifier="node_file_create"
]
class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
class NodeFileScanAction(NodeFileAbstractAction, discriminator="node-file-scan"):
"""Action which scans a file."""
config: "NodeFileScanAction.ConfigSchema"
@@ -92,7 +90,7 @@ class NodeFileScanAction(NodeFileAbstractAction, identifier="node_file_scan"):
verb: ClassVar[str] = "scan"
class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"):
class NodeFileDeleteAction(NodeFileAbstractAction, discriminator="node-file-delete"):
"""Action which deletes a file."""
config: "NodeFileDeleteAction.ConfigSchema"
@@ -106,7 +104,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None or config.file_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -119,7 +117,7 @@ class NodeFileDeleteAction(NodeFileAbstractAction, identifier="node_file_delete"
]
class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restore"):
class NodeFileRestoreAction(NodeFileAbstractAction, discriminator="node-file-restore"):
"""Action which restores a file."""
config: "NodeFileRestoreAction.ConfigSchema"
@@ -130,7 +128,7 @@ class NodeFileRestoreAction(NodeFileAbstractAction, identifier="node_file_restor
verb: ClassVar[str] = "restore"
class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrupt"):
class NodeFileCorruptAction(NodeFileAbstractAction, discriminator="node-file-corrupt"):
"""Action which corrupts a file."""
config: "NodeFileCorruptAction.ConfigSchema"
@@ -141,7 +139,7 @@ class NodeFileCorruptAction(NodeFileAbstractAction, identifier="node_file_corrup
verb: ClassVar[str] = "corrupt"
class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"):
class NodeFileAccessAction(NodeFileAbstractAction, discriminator="node-file-access"):
"""Action which increases a file's access count."""
config: "NodeFileAccessAction.ConfigSchema"
@@ -155,7 +153,7 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None or config.file_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -167,7 +165,7 @@ class NodeFileAccessAction(NodeFileAbstractAction, identifier="node_file_access"
]
class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_checkhash"):
class NodeFileCheckhashAction(NodeFileAbstractAction, discriminator="node-file-checkhash"):
"""Action which checks the hash of a file."""
config: "NodeFileCheckhashAction.ConfigSchema"
@@ -178,7 +176,7 @@ class NodeFileCheckhashAction(NodeFileAbstractAction, identifier="node_file_chec
verb: ClassVar[str] = "checkhash"
class NodeFileRepairAction(NodeFileAbstractAction, identifier="node_file_repair"):
class NodeFileRepairAction(NodeFileAbstractAction, discriminator="node-file-repair"):
"""Action which repairs a file."""
config: "NodeFileRepairAction.ConfigSchema"

18
src/primaite/game/agent/actions/folder.py
View File

@@ -22,9 +22,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
this base class.
"""
config: "NodeFolderAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NodeFolder actions."""
node_name: str
@@ -35,7 +33,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -47,7 +45,7 @@ class NodeFolderAbstractAction(AbstractAction, ABC):
]
class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_scan"):
class NodeFolderScanAction(NodeFolderAbstractAction, discriminator="node-folder-scan"):
"""Action which scans a folder."""
config: "NodeFolderScanAction.ConfigSchema"
@@ -58,7 +56,7 @@ class NodeFolderScanAction(NodeFolderAbstractAction, identifier="node_folder_sca
verb: ClassVar[str] = "scan"
class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folder_checkhash"):
class NodeFolderCheckhashAction(NodeFolderAbstractAction, discriminator="node-folder-checkhash"):
"""Action which checks the hash of a folder."""
config: "NodeFolderCheckhashAction.ConfigSchema"
@@ -69,7 +67,7 @@ class NodeFolderCheckhashAction(NodeFolderAbstractAction, identifier="node_folde
verb: ClassVar[str] = "checkhash"
class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_repair"):
class NodeFolderRepairAction(NodeFolderAbstractAction, discriminator="node-folder-repair"):
"""Action which repairs a folder."""
config: "NodeFolderRepairAction.ConfigSchema"
@@ -80,7 +78,7 @@ class NodeFolderRepairAction(NodeFolderAbstractAction, identifier="node_folder_r
verb: ClassVar[str] = "repair"
class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_restore"):
class NodeFolderRestoreAction(NodeFolderAbstractAction, discriminator="node-folder-restore"):
"""Action which restores a folder."""
config: "NodeFolderRestoreAction.ConfigSchema"
@@ -91,7 +89,7 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction, identifier="node_folder_
verb: ClassVar[str] = "restore"
class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_create"):
class NodeFolderCreateAction(NodeFolderAbstractAction, discriminator="node-folder-create"):
"""Action which creates a new folder."""
config: "NodeFolderCreateAction.ConfigSchema"
@@ -105,7 +103,7 @@ class NodeFolderCreateAction(NodeFolderAbstractAction, identifier="node_folder_c
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.folder_name is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",

10
src/primaite/game/agent/actions/host_nic.py
View File

@@ -16,9 +16,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
base class.
"""
config: "HostNICAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for HostNIC actions."""
node_name: str
@@ -29,7 +27,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.nic_num is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -40,7 +38,7 @@ class HostNICAbstractAction(AbstractAction, ABC):
]
class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
class HostNICEnableAction(HostNICAbstractAction, discriminator="host-nic-enable"):
"""Action which enables a NIC."""
config: "HostNICEnableAction.ConfigSchema"
@@ -51,7 +49,7 @@ class HostNICEnableAction(HostNICAbstractAction, identifier="host_nic_enable"):
verb: ClassVar[str] = "enable"
class HostNICDisableAction(HostNICAbstractAction, identifier="host_nic_disable"):
class HostNICDisableAction(HostNICAbstractAction, discriminator="host-nic-disable"):
"""Action which disables a NIC."""
config: "HostNICDisableAction.ConfigSchema"

12
src/primaite/game/agent/actions/manager.py
View File

@@ -5,9 +5,9 @@ agents:
- name: agent_1
action_space:
actions:
- do_nothing
- node_service_start
- node_service_stop
- do-nothing
- node-service-start
- node-service-stop
action_map:
"""
@@ -24,18 +24,18 @@ from primaite.interface.request import RequestFormat
__all__ = ("DoNothingAction", "ActionManager")
class DoNothingAction(AbstractAction, identifier="do_nothing"):
class DoNothingAction(AbstractAction, discriminator="do-nothing"):
"""Do Nothing Action."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for do_nothingAction."""
type: str = "do_nothing"
type: str = "do-nothing"
@classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["do_nothing"]
return ["do-nothing"]
class _ActionMapItem(BaseModel):

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

@@ -1,5 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
@@ -8,12 +9,10 @@ from primaite.interface.request import RequestFormat
__all__ = ("NetworkPortEnableAction", "NetworkPortDisableAction")
class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstract"):
class NetworkPortAbstractAction(AbstractAction, ABC):
"""Base class for Network port actions."""
config: "NetworkPortAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NetworkPort actions."""
target_nodename: str
@@ -24,7 +23,7 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.target_nodename is None or config.port_num is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
@@ -35,7 +34,7 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac
]
class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_port_enable"):
class NetworkPortEnableAction(NetworkPortAbstractAction, discriminator="network-port-enable"):
"""Action which enables are port on a router or a firewall."""
config: "NetworkPortEnableAction.ConfigSchema"
@@ -46,7 +45,7 @@ class NetworkPortEnableAction(NetworkPortAbstractAction, identifier="network_por
verb: ClassVar[str] = "enable"
class NetworkPortDisableAction(NetworkPortAbstractAction, identifier="network_port_disable"):
class NetworkPortDisableAction(NetworkPortAbstractAction, discriminator="network-port-disable"):
"""Action which disables are port on a router or a firewall."""
config: "NetworkPortDisableAction.ConfigSchema"

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

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from abc import ABC, abstractmethod
from typing import ClassVar, List, Optional, Union
from primaite.game.agent.actions.manager import AbstractAction
@@ -18,16 +18,14 @@ __all__ = (
)
class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
class NodeAbstractAction(AbstractAction, ABC):
"""
Abstract base class for node actions.
Any action which applies to a node and uses node_name as its only parameter can inherit from this base class.
"""
config: "NodeAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration schema for Node actions."""
node_name: str
@@ -39,7 +37,7 @@ class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
class NodeOSScanAction(NodeAbstractAction, discriminator="node-os-scan"):
"""Action which scans a node's OS."""
config: "NodeOSScanAction.ConfigSchema"
@@ -50,7 +48,7 @@ class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):
verb: ClassVar[str] = "scan"
class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
class NodeShutdownAction(NodeAbstractAction, discriminator="node-shutdown"):
"""Action which shuts down a node."""
config: "NodeShutdownAction.ConfigSchema"
@@ -61,7 +59,7 @@ class NodeShutdownAction(NodeAbstractAction, identifier="node_shutdown"):
verb: ClassVar[str] = "shutdown"
class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
class NodeStartupAction(NodeAbstractAction, discriminator="node-startup"):
"""Action which starts up a node."""
config: "NodeStartupAction.ConfigSchema"
@@ -72,7 +70,7 @@ class NodeStartupAction(NodeAbstractAction, identifier="node_startup"):
verb: ClassVar[str] = "startup"
class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
class NodeResetAction(NodeAbstractAction, discriminator="node-reset"):
"""Action which resets a node."""
config: "NodeResetAction.ConfigSchema"
@@ -83,12 +81,10 @@ class NodeResetAction(NodeAbstractAction, identifier="node_reset"):
verb: ClassVar[str] = "reset"
class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_action"):
class NodeNMAPAbstractAction(AbstractAction, ABC):
"""Base class for NodeNMAP actions."""
config: "NodeNMAPAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base Configuration Schema for NodeNMAP actions."""
target_ip_address: Union[str, List[str]]
@@ -103,8 +99,8 @@ class NodeNMAPAbstractAction(AbstractAction, identifier="node_nmap_abstract_acti
pass
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_scan"):
"""Action which performs an NMAP ping scan."""
class NodeNMAPPingScanAction(NodeNMAPAbstractAction, discriminator="node-nmap-ping-scan"):
"""Action which performs an nmap ping scan."""
config: "NodeNMAPPingScanAction.ConfigSchema"
@@ -116,14 +112,14 @@ class NodeNMAPPingScanAction(NodeNMAPAbstractAction, identifier="node_nmap_ping_
"node",
config.source_node,
"application",
"NMAP",
"nmap",
"ping_scan",
{"target_ip_address": config.target_ip_address, "show": config.show},
]
class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_scan"):
"""Action which performs an NMAP port scan."""
class NodeNMAPPortScanAction(NodeNMAPAbstractAction, discriminator="node-nmap-port-scan"):
"""Action which performs an nmap port scan."""
config: "NodeNMAPPortScanAction.ConfigSchema"
@@ -143,7 +139,7 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
"node",
config.source_node,
"application",
"NMAP",
"nmap",
"port_scan",
{
"target_ip_address": config.target_ip_address,
@@ -154,8 +150,8 @@ class NodeNMAPPortScanAction(NodeNMAPAbstractAction, identifier="node_nmap_port_
]
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_network_service_recon"):
"""Action which performs an NMAP network service recon (ping scan followed by port scan)."""
class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, discriminator="node-network-service-recon"):
"""Action which performs an nmap network service recon (ping scan followed by port scan)."""
config: "NodeNetworkServiceReconAction.ConfigSchema"
@@ -174,7 +170,7 @@ class NodeNetworkServiceReconAction(NodeNMAPAbstractAction, identifier="node_net
"node",
config.source_node,
"application",
"NMAP",
"nmap",
"network_service_recon",
{
"target_ip_address": config.target_ip_address,

25
src/primaite/game/agent/actions/service.py
View File

@@ -1,4 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import ABC
from typing import ClassVar
from primaite.game.agent.actions.manager import AbstractAction
@@ -17,15 +18,13 @@ __all__ = (
)
class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstract"):
class NodeServiceAbstractAction(AbstractAction, ABC):
"""Abstract Action for Node Service related actions.
Any actions which use node_name and service_name can inherit from this class.
"""
config: "NodeServiceAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
node_name: str
service_name: str
verb: ClassVar[str]
@@ -36,7 +35,7 @@ class NodeServiceAbstractAction(AbstractAction, identifier="node_service_abstrac
return ["network", "node", config.node_name, "service", config.service_name, config.verb]
class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_scan"):
class NodeServiceScanAction(NodeServiceAbstractAction, discriminator="node-service-scan"):
"""Action which scans a service."""
config: "NodeServiceScanAction.ConfigSchema"
@@ -47,7 +46,7 @@ class NodeServiceScanAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "scan"
class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_stop"):
class NodeServiceStopAction(NodeServiceAbstractAction, discriminator="node-service-stop"):
"""Action which stops a service."""
config: "NodeServiceStopAction.ConfigSchema"
@@ -58,7 +57,7 @@ class NodeServiceStopAction(NodeServiceAbstractAction, identifier="node_service_
verb: ClassVar[str] = "stop"
class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service_start"):
class NodeServiceStartAction(NodeServiceAbstractAction, discriminator="node-service-start"):
"""Action which starts a service."""
config: "NodeServiceStartAction.ConfigSchema"
@@ -69,7 +68,7 @@ class NodeServiceStartAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "start"
class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service_pause"):
class NodeServicePauseAction(NodeServiceAbstractAction, discriminator="node-service-pause"):
"""Action which pauses a service."""
config: "NodeServicePauseAction.ConfigSchema"
@@ -80,7 +79,7 @@ class NodeServicePauseAction(NodeServiceAbstractAction, identifier="node_service
verb: ClassVar[str] = "pause"
class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_service_resume"):
class NodeServiceResumeAction(NodeServiceAbstractAction, discriminator="node-service-resume"):
"""Action which resumes a service."""
config: "NodeServiceResumeAction.ConfigSchema"
@@ -91,7 +90,7 @@ class NodeServiceResumeAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "resume"
class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_service_restart"):
class NodeServiceRestartAction(NodeServiceAbstractAction, discriminator="node-service-restart"):
"""Action which restarts a service."""
config: "NodeServiceRestartAction.ConfigSchema"
@@ -102,7 +101,7 @@ class NodeServiceRestartAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "restart"
class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_service_disable"):
class NodeServiceDisableAction(NodeServiceAbstractAction, discriminator="node-service-disable"):
"""Action which disables a service."""
config: "NodeServiceDisableAction.ConfigSchema"
@@ -113,7 +112,7 @@ class NodeServiceDisableAction(NodeServiceAbstractAction, identifier="node_servi
verb: ClassVar[str] = "disable"
class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_service_enable"):
class NodeServiceEnableAction(NodeServiceAbstractAction, discriminator="node-service-enable"):
"""Action which enables a service."""
config: "NodeServiceEnableAction.ConfigSchema"
@@ -124,7 +123,7 @@ class NodeServiceEnableAction(NodeServiceAbstractAction, identifier="node_servic
verb: ClassVar[str] = "enable"
class NodeServiceFixAction(NodeServiceAbstractAction, identifier="node_service_fix"):
class NodeServiceFixAction(NodeServiceAbstractAction, discriminator="node-service-fix"):
"""Action which fixes a service."""
config: "NodeServiceFixAction.ConfigSchema"

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

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from abc import ABC, abstractmethod
from primaite.game.agent.actions.manager import AbstractAction
from primaite.interface.request import RequestFormat
@@ -11,12 +11,10 @@ __all__ = (
)
class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstract"):
class NodeSessionAbstractAction(AbstractAction, ABC):
"""Base class for NodeSession actions."""
config: "NodeSessionAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
class ConfigSchema(AbstractAction.ConfigSchema, ABC):
"""Base configuration schema for NodeSessionAbstractActions."""
node_name: str
@@ -33,7 +31,7 @@ class NodeSessionAbstractAction(AbstractAction, identifier="node_session_abstrac
pass
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_session_remote_login"):
class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, discriminator="node-session-remote-login"):
"""Action which performs a remote session login."""
config: "NodeSessionsRemoteLoginAction.ConfigSchema"
@@ -48,21 +46,21 @@ class NodeSessionsRemoteLoginAction(NodeSessionAbstractAction, identifier="node_
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.remote_ip is None:
return ["do_nothing"]
return ["do-nothing"]
return [
"network",
"node",
config.node_name,
"service",
"Terminal",
"node_session_remote_login",
"terminal",
"node-session-remote-login",
config.username,
config.password,
config.remote_ip,
]
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node_session_remote_logoff"):
class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, discriminator="node-session-remote-logoff"):
"""Action which performs a remote session logout."""
config: "NodeSessionsRemoteLogoutAction.ConfigSchema"
@@ -76,11 +74,11 @@ class NodeSessionsRemoteLogoutAction(NodeSessionAbstractAction, identifier="node
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
if config.node_name is None or config.remote_ip is None:
return ["do_nothing"]
return ["network", "node", config.node_name, "service", "Terminal", config.verb, config.remote_ip]
return ["do-nothing"]
return ["network", "node", config.node_name, "service", "terminal", config.verb, config.remote_ip]
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="node_account_change_password"):
class NodeAccountChangePasswordAction(NodeSessionAbstractAction, discriminator="node-account-change-password"):
"""Action which changes the password for a user."""
config: "NodeAccountChangePasswordAction.ConfigSchema"
@@ -100,7 +98,7 @@ class NodeAccountChangePasswordAction(NodeSessionAbstractAction, identifier="nod
"node",
config.node_name,
"service",
"UserManager",
"user-manager",
"change_password",
config.username,
config.current_password,

48
src/primaite/game/agent/actions/software.py
View File

@@ -22,7 +22,7 @@ __all__ = (
)
class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_ransomware_script"):
class ConfigureRansomwareScriptAction(AbstractAction, discriminator="configure-ransomware-script"):
"""Action which sets config parameters for a ransomware script on a node."""
config: "ConfigureRansomwareScriptAction.ConfigSchema"
@@ -39,16 +39,18 @@ class ConfigureRansomwareScriptAction(AbstractAction, identifier="configure_rans
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["do-nothing"]
data = dict(
server_ip_address=config.server_ip_address,
server_password=config.server_password,
payload=config.payload,
)
return ["network", "node", config.node_name, "application", "RansomwareScript", "configure", data]
return ["network", "node", config.node_name, "application", "ransomware-script", "configure", data]
class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identifier="c2_server_ransomware_configure"):
class RansomwareConfigureC2ServerAction(
ConfigureRansomwareScriptAction, discriminator="c2-server-ransomware-configure"
):
"""Action which causes a C2 server to send a command to set options on a ransomware script remotely."""
@classmethod
@@ -56,10 +58,10 @@ class RansomwareConfigureC2ServerAction(ConfigureRansomwareScriptAction, identif
data = dict(
server_ip_address=config.server_ip_address, server_password=config.server_password, payload=config.payload
)
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_configure", data]
return ["network", "node", config.node_name, "application", "c2-server", "ransomware_configure", data]
class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
class ConfigureDoSBotAction(AbstractAction, discriminator="configure-dos-bot"):
"""Action which sets config parameters for a DoS bot on a node."""
class ConfigSchema(AbstractAction.ConfigSchema):
@@ -88,10 +90,10 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
max_sessions=config.max_sessions,
)
data = {k: v for k, v in data.items() if v is not None}
return ["network", "node", config.node_name, "application", "DoSBot", "configure", data]
return ["network", "node", config.node_name, "application", "dos-bot", "configure", data]
class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
class ConfigureC2BeaconAction(AbstractAction, discriminator="configure-c2-beacon"):
"""Action which configures a C2 Beacon based on the parameters given."""
class ConfigSchema(AbstractAction.ConfigSchema):
@@ -112,10 +114,10 @@ class ConfigureC2BeaconAction(AbstractAction, identifier="configure_c2_beacon"):
masquerade_protocol=config.masquerade_protocol,
masquerade_port=config.masquerade_port,
)
return ["network", "node", config.node_name, "application", "C2Beacon", "configure", data]
return ["network", "node", config.node_name, "application", "c2-beacon", "configure", data]
class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_command"):
class NodeSendRemoteCommandAction(AbstractAction, discriminator="node-send-remote-command"):
"""Action which sends a terminal command to a remote node via SSH."""
config: "NodeSendRemoteCommandAction.ConfigSchema"
@@ -135,14 +137,14 @@ class NodeSendRemoteCommandAction(AbstractAction, identifier="node_send_remote_c
"node",
config.node_name,
"service",
"Terminal",
"terminal",
"send_remote_command",
config.remote_ip,
{"command": config.command},
]
class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_command"):
class TerminalC2ServerAction(AbstractAction, discriminator="c2-server-terminal-command"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to execute the terminal command passed."""
config: "TerminalC2ServerAction.ConfigSchema"
@@ -160,7 +162,7 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["do-nothing"]
command_model = {
"commands": config.commands,
@@ -168,10 +170,10 @@ class TerminalC2ServerAction(AbstractAction, identifier="c2_server_terminal_comm
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "terminal_command", command_model]
return ["network", "node", config.node_name, "application", "c2-server", "terminal_command", command_model]
class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ransomware_launch"):
class RansomwareLaunchC2ServerAction(AbstractAction, discriminator="c2-server-ransomware-launch"):
"""Action which causes the C2 Server to send a command to the C2 Beacon to launch the RansomwareScript."""
config: "RansomwareLaunchC2ServerAction.ConfigSchema"
@@ -185,12 +187,12 @@ class RansomwareLaunchC2ServerAction(AbstractAction, identifier="c2_server_ranso
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["do-nothing"]
# This action currently doesn't require any further configuration options.
return ["network", "node", config.node_name, "application", "C2Server", "ransomware_launch"]
return ["network", "node", config.node_name, "application", "c2-server", "ransomware_launch"]
class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfiltrate"):
class ExfiltrationC2ServerAction(AbstractAction, discriminator="c2-server-data-exfiltrate"):
"""Action which exfiltrates a target file from a certain node onto the C2 beacon and then the C2 Server."""
config: "ExfiltrationC2ServerAction.ConfigSchema"
@@ -210,7 +212,7 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["do-nothing"]
command_model = {
"target_file_name": config.target_file_name,
@@ -220,10 +222,10 @@ class ExfiltrationC2ServerAction(AbstractAction, identifier="c2_server_data_exfi
"username": config.username,
"password": config.password,
}
return ["network", "node", config.node_name, "application", "C2Server", "exfiltrate", command_model]
return ["network", "node", config.node_name, "application", "c2-server", "exfiltrate", command_model]
class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_database_client"):
class ConfigureDatabaseClientAction(AbstractAction, discriminator="configure-database-client"):
"""Action which sets config parameters for a database client on a node."""
config: "ConfigureDatabaseClientAction.ConfigSchema"
@@ -239,6 +241,6 @@ class ConfigureDatabaseClientAction(AbstractAction, identifier="configure_databa
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request that can be ingested by the simulation."""
if config.node_name is None:
return ["do_nothing"]
return ["do-nothing"]
data = {"server_ip_address": config.server_ip_address, "server_password": config.server_password}
return ["network", "node", config.node_name, "application", "DatabaseClient", "configure", data]
return ["network", "node", config.node_name, "application", "database-client", "configure", data]

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

@@ -68,11 +68,7 @@ class AbstractAgent(BaseModel, ABC):
)
reward_function: RewardFunction.ConfigSchema = Field(default_factory=lambda: RewardFunction.ConfigSchema())
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.logger: AgentLog = AgentLog(agent_name=self.config.ref)
config: "AbstractAgent.ConfigSchema" = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: AbstractAgent.ConfigSchema())
logger: AgentLog = AgentLog(agent_name="Abstract_Agent")
history: List[AgentHistoryItem] = []
@@ -83,13 +79,13 @@ class AbstractAgent(BaseModel, ABC):
_registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {discriminator}")
cls._registry[discriminator] = cls
def model_post_init(self, __context: Any) -> None:
"""Overwrite the default empty action, observation, and rewards with ones defined through the config."""
@@ -134,7 +130,7 @@ class AbstractAgent(BaseModel, ABC):
"""
# in RL agent, this method will send CAOS observation to RL agent, then receive a int 0-39,
# then use a bespoke conversion to take 1-40 int back into CAOS action
return ("do_nothing", {})
return ("do-nothing", {})
def format_request(self, action: Tuple[str, Dict], options: Dict[str, int]) -> RequestFormat:
# this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator.
@@ -165,23 +161,23 @@ class AbstractAgent(BaseModel, ABC):
return agent_class(config=config)
class AbstractScriptedAgent(AbstractAgent, identifier="AbstractScriptedAgent"):
class AbstractScriptedAgent(AbstractAgent, ABC):
"""Base class for actors which generate their own behaviour."""
config: "AbstractScriptedAgent.ConfigSchema" = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema())
class ConfigSchema(AbstractAgent.ConfigSchema):
class ConfigSchema(AbstractAgent.ConfigSchema, ABC):
"""Configuration Schema for AbstractScriptedAgents."""
type: str = "AbstractScriptedAgent"
config: ConfigSchema = Field(default_factory=lambda: AbstractScriptedAgent.ConfigSchema())
@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)
class ProxyAgent(AbstractAgent, identifier="ProxyAgent"):
class ProxyAgent(AbstractAgent, discriminator="proxy-agent"):
"""Agent that sends observations to an RL model and receives actions from that model."""
config: "ProxyAgent.ConfigSchema" = Field(default_factory=lambda: ProxyAgent.ConfigSchema())

2
src/primaite/game/agent/observations/acl_observation.py
View File

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class ACLObservation(AbstractObservation, identifier="ACL"):
class ACLObservation(AbstractObservation, discriminator="acl"):
"""ACL observation, provides information about access control lists within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

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

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__)
class FileObservation(AbstractObservation, identifier="FILE"):
class FileObservation(AbstractObservation, discriminator="file"):
"""File observation, provides status information about a file within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -125,7 +125,7 @@ class FileObservation(AbstractObservation, identifier="FILE"):
)
class FolderObservation(AbstractObservation, identifier="FOLDER"):
class FolderObservation(AbstractObservation, discriminator="folder"):
"""Folder observation, provides status information about a folder within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

4
src/primaite/game/agent/observations/firewall_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
class FirewallObservation(AbstractObservation, discriminator="firewall"):
"""Firewall observation, provides status information about a firewall within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -181,7 +181,7 @@ class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
},
}
if self.include_users:
sess = firewall_state["services"]["UserSessionManager"]
sess = firewall_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

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

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class HostObservation(AbstractObservation, identifier="HOST"):
class HostObservation(AbstractObservation, discriminator="host"):
"""Host observation, provides status information about a host within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -209,7 +209,7 @@ class HostObservation(AbstractObservation, identifier="HOST"):
obs["num_file_creations"] = node_state["file_system"]["num_file_creations"]
obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"]
if self.include_users:
sess = node_state["services"]["UserSessionManager"]
sess = node_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

4
src/primaite/game/agent/observations/link_observation.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_ST
_LOGGER = getLogger(__name__)
class LinkObservation(AbstractObservation, identifier="LINK"):
class LinkObservation(AbstractObservation, discriminator="link"):
"""Link observation, providing information about a specific link within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -90,7 +90,7 @@ class LinkObservation(AbstractObservation, identifier="LINK"):
return cls(where=where)
class LinksObservation(AbstractObservation, identifier="LINKS"):
class LinksObservation(AbstractObservation, discriminator="links"):
"""Collection of link observations representing multiple links within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

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

@@ -12,7 +12,7 @@ from primaite.utils.validation.ip_protocol import IPProtocol
from primaite.utils.validation.port import Port
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
class NICObservation(AbstractObservation, discriminator="network-interface"):
"""Status information about a network interface within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -227,7 +227,7 @@ class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
)
class PortObservation(AbstractObservation, identifier="PORT"):
class PortObservation(AbstractObservation, discriminator="port"):
"""Port observation, provides status information about a network port within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

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

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class NodesObservation(AbstractObservation, identifier="NODES"):
class NodesObservation(AbstractObservation, discriminator="nodes"):
"""Nodes observation, provides status information about nodes within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

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

@@ -11,7 +11,7 @@ from pydantic import BaseModel, computed_field, ConfigDict, Field, model_validat
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
class NestedObservation(AbstractObservation, identifier="CUSTOM"):
class NestedObservation(AbstractObservation, discriminator="custom"):
"""Observation type that allows combining other observations into a gymnasium.spaces.Dict space."""
class NestedObservationItem(BaseModel):
@@ -19,7 +19,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
model_config = ConfigDict(extra="forbid")
type: str
"""Select observation class. It maps to the identifier of the obs class by checking the registry."""
"""Select observation class. It maps to the discriminator of the obs class by checking the registry."""
label: str
"""Dict key in the final observation space."""
options: Dict
@@ -48,7 +48,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
def __init__(self, components: Dict[str, AbstractObservation]) -> None:
"""Initialise nested observation."""
self.components: Dict[str, AbstractObservation] = components
"""Maps label: observation object"""
"""Maps label observation object"""
self.default_observation = {label: obs.default_observation for label, obs in self.components.items()}
"""Default observation is just the default observations of constituents."""
@@ -84,7 +84,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
```yaml
observation_space:
- type: CUSTOM
- type: custom
options:
components:
@@ -119,7 +119,7 @@ class NestedObservation(AbstractObservation, identifier="CUSTOM"):
return cls(components=instances)
class NullObservation(AbstractObservation, identifier="NONE"):
class NullObservation(AbstractObservation, discriminator="none"):
"""Empty observation that acts as a placeholder."""
def __init__(self) -> None:
@@ -157,8 +157,8 @@ class ObservationManager(BaseModel):
"""Config Schema for Observation Manager."""
model_config = ConfigDict(extra="forbid")
type: str = "NONE"
"""Identifier name for the top-level observation."""
type: str = "none"
"""discriminator name for the top-level observation."""
options: AbstractObservation.ConfigSchema = Field(
default_factory=lambda: NullObservation.ConfigSchema(), validate_default=True
)
@@ -187,7 +187,7 @@ class ObservationManager(BaseModel):
return data
# (TODO: duplicate default definition between here and the actual model)
obs_type = data["type"] if "type" in data else "NONE"
obs_type = data["type"] if "type" in data else "none"
obs_class = AbstractObservation._registry[obs_type]
# if no options are passed in, try to create a default schema. Only works if there are no mandatory fields
@@ -235,7 +235,7 @@ class ObservationManager(BaseModel):
:param config: Dictionary containing the configuration for this observation space.
If None, a blank observation space is created.
Otherwise, this must be a Dict with a type field and options field.
type: string that corresponds to one of the observation identifiers that are provided when subclassing
type: string that corresponds to one of the observation discriminators that are provided when subclassing
AbstractObservation
options: this must adhere to the chosen observation type's ConfigSchema nested class.
:type config: Dict

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

@@ -31,20 +31,20 @@ class AbstractObservation(ABC):
"""Initialise an observation. This method must be overwritten."""
self.default_observation: ObsType
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an observation type.
:param identifier: Identifier used to uniquely specify observation component types.
:type identifier: str
:param discriminator: discriminator used to uniquely specify observation component types.
:type discriminator: str
:raises ValueError: When attempting to create a component with a name that is already in use.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate observation component type {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate observation component type {discriminator}")
cls._registry[discriminator] = cls
@abstractmethod
def observe(self, state: Dict) -> Any:

4
src/primaite/game/agent/observations/router_observation.py
View File

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port
_LOGGER = getLogger(__name__)
class RouterObservation(AbstractObservation, identifier="ROUTER"):
class RouterObservation(AbstractObservation, discriminator="router"):
"""Router observation, provides status information about a router within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -113,7 +113,7 @@ class RouterObservation(AbstractObservation, identifier="ROUTER"):
if self.ports:
obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
if self.include_users:
sess = router_state["services"]["UserSessionManager"]
sess = router_state["services"]["user-session-manager"]
obs["users"] = {
"local_login": 1 if sess["current_local_user"] else 0,
"remote_sessions": min(self.max_users, len(sess["active_remote_sessions"])),

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

@@ -10,7 +10,7 @@ from primaite.game.agent.observations.observations import AbstractObservation, W
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
class ServiceObservation(AbstractObservation, discriminator="service"):
"""Service observation, shows status of a service in the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
@@ -73,7 +73,7 @@ class ServiceObservation(AbstractObservation, identifier="SERVICE"):
return cls(where=parent_where + ["services", config.service_name])
class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
class ApplicationObservation(AbstractObservation, discriminator="application"):
"""Application observation, shows the status of an application within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):

58
src/primaite/game/agent/rewards.py
View File

@@ -12,7 +12,7 @@ the structure:
```yaml
reward_function:
reward_components:
- type: DATABASE_FILE_INTEGRITY
- type: database-file-integrity
weight: 0.5
options:
node_name: database_server
@@ -20,7 +20,7 @@ the structure:
file_name: database.db
- type: WEB_SERVER_404_PENALTY
- type: web-server-404-penalty
weight: 0.5
options:
node_name: web_server
@@ -43,25 +43,25 @@ _LOGGER = getLogger(__name__)
WhereType = Optional[Iterable[Union[str, int]]]
class AbstractReward(BaseModel):
class AbstractReward(BaseModel, ABC):
"""Base class for reward function components."""
config: "AbstractReward.ConfigSchema"
class ConfigSchema(BaseModel, ABC):
"""Config schema for AbstractReward."""
type: str = ""
config: ConfigSchema
_registry: ClassVar[Dict[str, Type["AbstractReward"]]] = {}
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate reward {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate reward {discriminator}")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "AbstractReward":
@@ -92,7 +92,7 @@ class AbstractReward(BaseModel):
return 0.0
class DummyReward(AbstractReward, identifier="DUMMY"):
class DummyReward(AbstractReward, discriminator="dummy"):
"""Dummy reward function component which always returns 0.0."""
def calculate(self, state: Dict, last_action_response: "AgentHistoryItem") -> float:
@@ -108,7 +108,7 @@ class DummyReward(AbstractReward, identifier="DUMMY"):
return 0.0
class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"):
class DatabaseFileIntegrity(AbstractReward, discriminator="database-file-integrity"):
"""Reward function component which rewards the agent for maintaining the integrity of a database file."""
config: "DatabaseFileIntegrity.ConfigSchema"
@@ -118,7 +118,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for DatabaseFileIntegrity."""
type: str = "DATABASE_FILE_INTEGRITY"
type: str = "database-file-integrity"
node_hostname: str
folder_name: str
file_name: str
@@ -161,7 +161,7 @@ class DatabaseFileIntegrity(AbstractReward, identifier="DATABASE_FILE_INTEGRITY"
return 0
class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
class WebServer404Penalty(AbstractReward, discriminator="web-server-404-penalty"):
"""Reward function component which penalises the agent when the web server returns a 404 error."""
config: "WebServer404Penalty.ConfigSchema"
@@ -171,7 +171,7 @@ class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for WebServer404Penalty."""
type: str = "WEB_SERVER_404_PENALTY"
type: str = "web-server-404-penalty"
node_hostname: str
service_name: str
sticky: bool = True
@@ -215,7 +215,7 @@ class WebServer404Penalty(AbstractReward, identifier="WEB_SERVER_404_PENALTY"):
return self.reward
class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_PENALTY"):
class WebpageUnavailablePenalty(AbstractReward, discriminator="webpage-unavailable-penalty"):
"""Penalises the agent when the web browser fails to fetch a webpage."""
config: "WebpageUnavailablePenalty.ConfigSchema"
@@ -225,7 +225,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for WebpageUnavailablePenalty."""
type: str = "WEBPAGE_UNAVAILABLE_PENALTY"
type: str = "webpage-unavailable-penalty"
node_hostname: str = ""
sticky: bool = True
@@ -248,7 +248,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
"nodes",
self.config.node_hostname,
"applications",
"WebBrowser",
"web-browser",
]
web_browser_state = access_from_nested_dict(state, self.location_in_state)
@@ -261,7 +261,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
"node",
self.config.node_hostname,
"application",
"WebBrowser",
"web-browser",
"execute",
]
@@ -289,7 +289,7 @@ class WebpageUnavailablePenalty(AbstractReward, identifier="WEBPAGE_UNAVAILABLE_
return self.reward
class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"):
class GreenAdminDatabaseUnreachablePenalty(AbstractReward, discriminator="green-admin-database-unreachable-penalty"):
"""Penalises the agent when the green db clients fail to connect to the database."""
config: "GreenAdminDatabaseUnreachablePenalty.ConfigSchema"
@@ -298,7 +298,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
class ConfigSchema(AbstractReward.ConfigSchema):
"""ConfigSchema for GreenAdminDatabaseUnreachablePenalty."""
type: str = "GREEN_ADMIN_DATABASE_UNREACHABLE_PENALTY"
type: str = "green-admin-database-unreachable-penalty"
node_hostname: str
sticky: bool = True
@@ -322,7 +322,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
"node",
self.config.node_hostname,
"application",
"DatabaseClient",
"database-client",
"execute",
]
@@ -339,7 +339,7 @@ class GreenAdminDatabaseUnreachablePenalty(AbstractReward, identifier="GREEN_ADM
return self.reward
class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
class SharedReward(AbstractReward, discriminator="shared-reward"):
"""Adds another agent's reward to the overall reward."""
config: "SharedReward.ConfigSchema"
@@ -347,7 +347,7 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
class ConfigSchema(AbstractReward.ConfigSchema):
"""Config schema for SharedReward."""
type: str = "SHARED_REWARD"
type: str = "shared-reward"
agent_name: str
def default_callback(agent_name: str) -> Never:
@@ -376,17 +376,17 @@ class SharedReward(AbstractReward, identifier="SHARED_REWARD"):
return self.callback(self.config.agent_name)
class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"):
"""Apply a negative reward when taking any action except do_nothing."""
class ActionPenalty(AbstractReward, discriminator="action-penalty"):
"""Apply a negative reward when taking any action except do-nothing."""
config: "ActionPenalty.ConfigSchema"
class ConfigSchema(AbstractReward.ConfigSchema):
"""Config schema for ActionPenalty.
:param action_penalty: Reward to give agents for taking any action except do_nothing
:param action_penalty: Reward to give agents for taking any action except do-nothing
:type action_penalty: float
:param do_nothing_penalty: Reward to give agent for taking the do_nothing action
:param do_nothing_penalty: Reward to give agent for taking the do-nothing action
:type do_nothing_penalty: float
"""
@@ -403,7 +403,7 @@ class ActionPenalty(AbstractReward, identifier="ACTION_PENALTY"):
:return: Reward value
:rtype: float
"""
if last_action_response.action == "do_nothing":
if last_action_response.action == "do-nothing":
return self.config.do_nothing_penalty
else:

10
src/primaite/game/agent/scripted_agents/abstract_tap.py
View File

@@ -2,7 +2,7 @@
from __future__ import annotations
import random
from abc import abstractmethod
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple
from gymnasium.core import ObsType
@@ -13,21 +13,21 @@ from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "AbstractTAPAgent"
class AbstractTAPAgent(PeriodicAgent, identifier="AbstractTAP"):
class AbstractTAPAgent(PeriodicAgent, ABC):
"""Base class for TAP agents to inherit from."""
config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema())
next_execution_timestep: int = 0
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema, ABC):
"""Schema for the `agent_settings` part of the agent config."""
possible_starting_nodes: List[str] = Field(default_factory=list)
class ConfigSchema(PeriodicAgent.ConfigSchema):
class ConfigSchema(PeriodicAgent.ConfigSchema, ABC):
"""Configuration schema for Abstract TAP agents."""
type: str = "AbstractTAP"
type: str = "abstract-tap"
agent_settings: AbstractTAPAgent.AgentSettingsSchema = Field(
default_factory=lambda: AbstractTAPAgent.AgentSettingsSchema()
)

12
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -6,21 +6,19 @@ from pydantic import Field
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "DataManipulationAgent"
class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"):
class DataManipulationAgent(PeriodicAgent, discriminator="red-database-corrupting-agent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
target_application: str = "DataManipulationBot"
target_application: str = "data-manipulation-bot"
class ConfigSchema(PeriodicAgent.ConfigSchema):
"""Configuration Schema for DataManipulationAgent."""
type: str = "RedDatabaseCorruptingAgent"
type: str = "red-database-corrupting-agent"
agent_settings: "DataManipulationAgent.AgentSettingsSchema" = Field(
default_factory=lambda: DataManipulationAgent.AgentSettingsSchema()
)
@@ -43,13 +41,13 @@ class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgen
"""
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
return "do-nothing", {}
self._set_next_execution_timestep(
timestep=timestep + self.config.agent_settings.frequency, variance=self.config.agent_settings.variance
)
self.logger.info(msg="Performing a data manipulation attack!")
return "node_application_execute", {
return "node-application-execute", {
"node_name": self.start_node,
"application_name": self.config.agent_settings.target_application,
}

4
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -13,7 +13,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = "ProbabilisticAgent"
class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"):
class ProbabilisticAgent(AbstractScriptedAgent, discriminator="probabilistic-agent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535)))
@@ -46,7 +46,7 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration schema for Probabilistic Agent."""
type: str = "ProbabilisticAgent"
type: str = "probabilistic-agent"
agent_settings: "ProbabilisticAgent.AgentSettingsSchema" = Field(
default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema()
)

12
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -11,7 +11,7 @@ from primaite.game.agent.interface import AbstractScriptedAgent
__all__ = ("RandomAgent", "PeriodicAgent")
class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
class RandomAgent(AbstractScriptedAgent, discriminator="random-agent"):
"""Agent that ignores its observation and acts completely at random."""
config: "RandomAgent.ConfigSchema" = Field(default_factory=lambda: RandomAgent.ConfigSchema())
@@ -19,7 +19,7 @@ class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Random Agents."""
type: str = "RandomAgent"
type: str = "random-agent"
def get_action(self) -> Tuple[str, Dict]:
"""Sample the action space randomly.
@@ -34,7 +34,7 @@ class RandomAgent(AbstractScriptedAgent, identifier="RandomAgent"):
return self.action_manager.get_action(self.action_manager.space.sample())
class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
class PeriodicAgent(AbstractScriptedAgent, discriminator="periodic-agent"):
"""Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""
config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema())
@@ -72,7 +72,7 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Periodic Agent."""
type: str = "PeriodicAgent"
type: str = "periodic-agent"
"""Name of the agent."""
agent_settings: "PeriodicAgent.AgentSettingsSchema" = Field(
default_factory=lambda: PeriodicAgent.AgentSettingsSchema()
@@ -113,9 +113,9 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
self._set_next_execution_timestep(
timestep + self.config.agent_settings.frequency, self.config.agent_settings.variance
)
return "node_application_execute", {
return "node-application-execute", {
"node_name": self.start_node,
"application_name": self.config.agent_settings.target_application,
}
return "do_nothing", {}
return "do-nothing", {}

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

@@ -14,6 +14,7 @@ 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.switch import Switch
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
from primaite.simulator.network.nmne import NMNEConfig
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.application import Application
@@ -268,9 +269,11 @@ class PrimaiteGame:
new_node = None
if n_type in Node._registry:
if n_type == "wireless_router":
node_cfg["airspace"] = net.airspace
new_node = Node._registry[n_type].from_config(config=node_cfg)
n_class = Node._registry[n_type]
if issubclass(n_class, WirelessRouter):
new_node = n_class.from_config(config=node_cfg, airspace=net.airspace)
else:
new_node = Node._registry[n_type].from_config(config=node_cfg)
else:
msg = f"invalid node type {n_type} in config"
_LOGGER.error(msg)
@@ -288,8 +291,8 @@ class PrimaiteGame:
if "folder_restore_duration" in defaults_config:
new_node.file_system._default_folder_restore_duration = defaults_config["folder_restore_duration"]
if "users" in node_cfg and new_node.software_manager.software.get("UserManager"):
user_manager: UserManager = new_node.software_manager.software["UserManager"] # noqa
if "users" in node_cfg and new_node.software_manager.software.get("user-manager"):
user_manager: UserManager = new_node.software_manager.software["user-manager"] # noqa
for user_cfg in node_cfg["users"]:
user_manager.add_user(**user_cfg, bypass_can_perform_action=True)
@@ -321,7 +324,7 @@ class PrimaiteGame:
if service_class is not None:
_LOGGER.debug(f"installing {service_type} on node {new_node.config.hostname}")
new_node.software_manager.install(service_class, software_config=service_cfg.get("options", {}))
new_service = new_node.software_manager.software[service_class.__name__]
new_service = new_node.software_manager.software[service_type]
# fixing duration for the service
if "fixing_duration" in service_cfg.get("options", {}):

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

@@ -95,7 +95,7 @@
" - install\n",
" - RansomwareScript\n",
" 5:\n",
" action: c2_server_ransomware_configure\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_name: client_1\n",
" server_ip_address: 192.168.1.14\n",
@@ -392,7 +392,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### **Command and Control** | C2 Server Actions | c2_server_ransomware_configure\n",
"### **Command and Control** | C2 Server Actions | c2-server-ransomware-configure\n",
"\n",
"Another action the C2 Server grants is the ability for a Red Agent to configure the RansomwareScript via the C2 Server rather than the note directly.\n",
"\n",
@@ -411,7 +411,7 @@
" ...\n",
" action_map:\n",
" 5:\n",
" action: c2_server_ransomware_configure\n",
" action: c2-server-ransomware-configure\n",
" options:\n",
" node_id: 1\n",
" config:\n",

2
src/primaite/simulator/_package_data/create-simulation_demo.ipynb
View File

@@ -172,7 +172,7 @@
"\n",
"\n",
"# no applications exist yet so we will create our own.\n",
"class MSPaint(Application, identifier=\"MSPaint\"):\n",
"class MSPaint(Application, discriminator=\"MSPaint\"):\n",
" def describe_state(self):\n",
" return super().describe_state()"
]

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

@@ -244,7 +244,7 @@ class SimComponent(BaseModel):
..code::python
class WebBrowser(Application, identifier="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager() # all requests generic to any Application get initialised
rm.add_request(...) # initialise any requests specific to the web browser

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

@@ -130,8 +130,8 @@ class File(FileSystemItemABC):
Return False if corruption is detected, otherwise True
"""
warnings.warn("node_file_checkhash is currently not implemented.")
self.sys_log.warning("node_file_checkhash is currently not implemented.")
warnings.warn("node-file-checkhash is currently not implemented.")
self.sys_log.warning("node-file-checkhash is currently not implemented.")
return False
if self.deleted:

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

@@ -396,8 +396,8 @@ class Folder(FileSystemItemABC):
Return False if corruption is detected, otherwise True
"""
warnings.warn("node_folder_checkhash is currently not implemented.")
self.sys_log.error("node_folder_checkhash is currently not implemented.")
warnings.warn("node-folder-checkhash is currently not implemented.")
self.sys_log.error("node-folder-checkhash is currently not implemented.")
return False
if self.deleted:

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

@@ -180,7 +180,7 @@ class Network(SimComponent):
table.align = "l"
table.title = "Nodes"
for node in self.nodes.values():
table.add_row((node.config.hostname, type(node)._identifier, node.operating_state.name))
table.add_row((node.config.hostname, type(node)._discriminator, node.operating_state.name))
print(table)
if ip_addresses:

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

@@ -22,7 +22,7 @@ class NetworkNodeAdder(BaseModel):
Here is a template that users can use to define custom node adders:
```
class YourNodeAdder(NetworkNodeAdder, identifier="your_name"):
class YourNodeAdder(NetworkNodeAdder, discriminator="your-name"):
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
property_1 : str
property_2 : int
@@ -40,8 +40,8 @@ class NetworkNodeAdder(BaseModel):
"""
Base schema for node adders.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The identifier is used
by the from_config method to select the correct node adder at runtime.
Child classes of NetworkNodeAdder must define a schema which inherits from this schema. The discriminator is
used by the from_config method to select the correct node adder at runtime.
"""
model_config = ConfigDict(extra="forbid")
@@ -50,20 +50,20 @@ class NetworkNodeAdder(BaseModel):
_registry: ClassVar[Dict[str, Type["NetworkNodeAdder"]]] = {}
def __init_subclass__(cls, identifier: Optional[str], **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str], **kwargs: Any) -> None:
"""
Register a network node adder class.
:param identifier: Unique name for the node adder to use for matching against primaite config entries.
:type identifier: str
:param discriminator: Unique name for the node adder to use for matching against primaite config entries.
:type discriminator: str
:raises ValueError: When attempting to register a name that is already reserved.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Duplicate node adder {identifier}")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Duplicate node adder {discriminator}")
cls._registry[discriminator] = cls
@classmethod
@abstractmethod
@@ -99,13 +99,13 @@ class NetworkNodeAdder(BaseModel):
adder_class.add_nodes_to_net(config=adder_class.ConfigSchema(**config), network=network)
class OfficeLANAdder(NetworkNodeAdder, identifier="office_lan"):
class OfficeLANAdder(NetworkNodeAdder, discriminator="office-lan"):
"""Creates an office LAN."""
class ConfigSchema(NetworkNodeAdder.ConfigSchema):
"""Configuration schema for OfficeLANAdder."""
type: Literal["office_lan"] = "office_lan"
type: Literal["office-lan"] = "office-lan"
lan_name: str
"""Name of lan used for generating hostnames for new nodes."""
subnet_base: int

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

@@ -824,7 +824,7 @@ class User(SimComponent):
return self.model_dump()
class UserManager(Service, identifier="UserManager"):
class UserManager(Service, discriminator="user-manager"):
"""
Manages users within the PrimAITE system, handling creation, authentication, and administration.
@@ -836,7 +836,7 @@ class UserManager(Service, identifier="UserManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserManager."""
type: str = "UserManager"
type: str = "user-manager"
config: "UserManager.ConfigSchema" = Field(default_factory=lambda: UserManager.ConfigSchema())
@@ -849,7 +849,7 @@ class UserManager(Service, identifier="UserManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserManager"
kwargs["name"] = "user-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1037,7 +1037,7 @@ class UserManager(Service, identifier="UserManager"):
@property
def _user_session_manager(self) -> "UserSessionManager":
return self.software_manager.software["UserSessionManager"] # noqa
return self.software_manager.software["user-session-manager"] # noqa
class UserSession(SimComponent):
@@ -1137,7 +1137,7 @@ class RemoteUserSession(UserSession):
return state
class UserSessionManager(Service, identifier="UserSessionManager"):
class UserSessionManager(Service, discriminator="user-session-manager"):
"""
Manages user sessions on a Node, including local and remote sessions.
@@ -1147,7 +1147,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for UserSessionManager."""
type: str = "UserSessionManager"
type: str = "user-session-manager"
config: "UserSessionManager.ConfigSchema" = Field(default_factory=lambda: UserSessionManager.ConfigSchema())
@@ -1179,7 +1179,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
:param username: The username for the default admin user
:param password: The password for the default admin user
"""
kwargs["name"] = "UserSessionManager"
kwargs["name"] = "user-session-manager"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)
@@ -1289,7 +1289,7 @@ class UserSessionManager(Service, identifier="UserSessionManager"):
:return: The UserManager instance.
"""
return self.software_manager.software["UserManager"] # noqa
return self.software_manager.software["user-manager"] # noqa
def pre_timestep(self, timestep: int) -> None:
"""Apply any pre-timestep logic that helps make sure we have the correct observations."""
@@ -1522,15 +1522,18 @@ class Node(SimComponent, ABC):
_registry: ClassVar[Dict[str, Type["Node"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
_identifier: ClassVar[str] = "unknown"
"""Identifier for this particular class, used for printing and logging. Each subclass redefines this."""
# 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."""
class ConfigSchema(BaseModel, ABC):
"""Configuration Schema for Node based classes."""
model_config = ConfigDict(arbitrary_types_allowed=True)
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
"""Configure pydantic to allow arbitrary types, let the instance have attributes not present in the model."""
type: str
hostname: str
"The node hostname on the network."
@@ -1569,6 +1572,8 @@ class Node(SimComponent, ABC):
operating_state: Any = None
users: Any = None # Temporary to appease "extra=forbid"
config: ConfigSchema = Field(default_factory=lambda: Node.ConfigSchema())
"""Configuration items within Node"""
@@ -1586,22 +1591,22 @@ class Node(SimComponent, ABC):
obj = cls(config=cls.ConfigSchema(**config))
return obj
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a node type.
:param identifier: Uniquely specifies an node class by name. Used for finding items by config.
:type identifier: str
:param discriminator: Uniquely specifies an node class by name. Used for finding items by config.
:type discriminator: str
:raises ValueError: When attempting to register an node with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
identifier = identifier.lower()
if identifier in cls._registry:
raise ValueError(f"Tried to define new node {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
cls._identifier = identifier
discriminator = discriminator.lower()
if discriminator in cls._registry:
raise ValueError(f"Tried to define new node {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
cls._discriminator = discriminator
def __init__(self, **kwargs):
"""
@@ -1635,17 +1640,17 @@ class Node(SimComponent, ABC):
@property
def user_manager(self) -> Optional[UserManager]:
"""The Nodes User Manager."""
return self.software_manager.software.get("UserManager") # noqa
return self.software_manager.software.get("user-manager") # noqa
@property
def user_session_manager(self) -> Optional[UserSessionManager]:
"""The Nodes User Session Manager."""
return self.software_manager.software.get("UserSessionManager") # noqa
return self.software_manager.software.get("user-session-manager") # noqa
@property
def terminal(self) -> Optional[Terminal]:
"""The Nodes Terminal."""
return self.software_manager.software.get("Terminal")
"""The Node's Terminal."""
return self.software_manager.software.get("terminal")
def local_login(self, username: str, password: str) -> Optional[str]:
"""

7
src/primaite/simulator/network/hardware/nodes/host/computer.py
View File

@@ -1,5 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import ClassVar, Dict
from typing import ClassVar, Dict, Literal
from pydantic import Field
@@ -7,7 +7,7 @@ from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
from primaite.simulator.system.services.ftp.ftp_client import FTPClient
class Computer(HostNode, identifier="computer"):
class Computer(HostNode, discriminator="computer"):
"""
A basic Computer class.
@@ -35,11 +35,12 @@ class Computer(HostNode, identifier="computer"):
* Web Browser
"""
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient}
SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "ftp-client": FTPClient}
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Computer class."""
type: Literal["computer"] = "computer"
hostname: str = "Computer"
config: ConfigSchema = Field(default_factory=lambda: Computer.ConfigSchema())

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

@@ -2,7 +2,7 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Any, ClassVar, Dict, Optional
from typing import Any, ClassVar, Dict, Literal, Optional
from pydantic import Field
@@ -267,7 +267,7 @@ class NIC(IPWiredNetworkInterface):
return f"Port {self.port_name if self.port_name else self.port_num}: {self.mac_address}/{self.ip_address}"
class HostNode(Node, identifier="HostNode"):
class HostNode(Node, discriminator="host-node"):
"""
Represents a host node in the network.
@@ -314,14 +314,14 @@ class HostNode(Node, identifier="HostNode"):
SYSTEM_SOFTWARE: ClassVar[Dict] = {
"HostARP": HostARP,
"ICMP": ICMP,
"DNSClient": DNSClient,
"NTPClient": NTPClient,
"WebBrowser": WebBrowser,
"NMAP": NMAP,
"UserSessionManager": UserSessionManager,
"UserManager": UserManager,
"Terminal": Terminal,
"icmp": ICMP,
"dns-client": DNSClient,
"ntp-client": NTPClient,
"web-browser": WebBrowser,
"nmap": NMAP,
"user-session-manager": UserSessionManager,
"user-manager": UserManager,
"terminal": Terminal,
}
"""List of system software that is automatically installed on nodes."""
@@ -333,9 +333,14 @@ class HostNode(Node, identifier="HostNode"):
class ConfigSchema(Node.ConfigSchema):
"""Configuration Schema for HostNode class."""
type: Literal["host-node"]
hostname: str = "HostNode"
subnet_mask: IPV4Address = "255.255.255.0"
ip_address: IPV4Address
services: Any = None # temporarily unset to appease extra="forbid"
applications: Any = None # temporarily unset to appease extra="forbid"
folders: Any = None # temporarily unset to appease extra="forbid"
network_interfaces: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: HostNode.ConfigSchema())
@@ -351,7 +356,7 @@ class HostNode(Node, identifier="HostNode"):
:return: NMAP application installed on the Node.
:rtype: Optional[NMAP]
"""
return self.software_manager.software.get("NMAP")
return self.software_manager.software.get("nmap")
@property
def arp(self) -> Optional[ARP]:
@@ -361,7 +366,7 @@ class HostNode(Node, identifier="HostNode"):
:return: ARP Cache for given HostNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")
def default_gateway_hello(self):
"""
@@ -393,8 +398,8 @@ class HostNode(Node, identifier="HostNode"):
dst_port = frame.udp.dst_port
can_accept_nmap = False
if self.software_manager.software.get("NMAP"):
if self.software_manager.software["NMAP"].operating_state == ApplicationOperatingState.RUNNING:
if self.software_manager.software.get("nmap"):
if self.software_manager.software["nmap"].operating_state == ApplicationOperatingState.RUNNING:
can_accept_nmap = True
accept_nmap = can_accept_nmap and frame.payload.__class__.__name__ == "PortScanPayload"

8
src/primaite/simulator/network/hardware/nodes/host/server.py
View File

@@ -1,11 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Literal
from pydantic import Field
from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
class Server(HostNode, identifier="server"):
class Server(HostNode, discriminator="server"):
"""
A basic Server class.
@@ -36,12 +38,13 @@ class Server(HostNode, identifier="server"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Server class."""
type: Literal["server"] = "server"
hostname: str = "server"
config: ConfigSchema = Field(default_factory=lambda: Server.ConfigSchema())
class Printer(HostNode, identifier="printer"):
class Printer(HostNode, discriminator="printer"):
"""Printer? I don't even know her!."""
# TODO: Implement printer-specific behaviour
@@ -49,6 +52,7 @@ class Printer(HostNode, identifier="printer"):
class ConfigSchema(HostNode.ConfigSchema):
"""Configuration Schema for Printer class."""
type: Literal["printer"] = "printer"
hostname: str = "printer"
config: ConfigSchema = Field(default_factory=lambda: Printer.ConfigSchema())

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

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Dict, Final, Union
from typing import Dict, Final, Literal, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -26,7 +26,7 @@ DMZ_PORT_ID: Final[int] = 3
"""The Firewall port ID of the DMZ port."""
class Firewall(Router, identifier="firewall"):
class Firewall(Router, discriminator="firewall"):
"""
A Firewall class that extends the functionality of a Router.
@@ -103,6 +103,7 @@ class Firewall(Router, identifier="firewall"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for Firewall 'Nodes' within PrimAITE."""
type: Literal["firewall"] = "firewall"
hostname: str = "firewall"
num_ports: int = 0

11
src/primaite/simulator/network/hardware/nodes/network/network_node.py
View File

@@ -1,13 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from abc import abstractmethod
from typing import Optional
from typing import Any, Optional
from primaite.simulator.network.hardware.base import NetworkInterface, Node
from primaite.simulator.network.transmission.data_link_layer import Frame
from primaite.simulator.system.services.arp.arp import ARP
class NetworkNode(Node, identifier="NetworkNode"):
class NetworkNode(Node, discriminator="network-node"):
"""
Represents an abstract base class for a network node that can receive and process network frames.
@@ -16,6 +16,11 @@ class NetworkNode(Node, identifier="NetworkNode"):
provide functionality for receiving and processing frames received on their network interfaces.
"""
class ConfigSchema(Node.ConfigSchema):
"""Config schema for Node baseclass."""
num_ports: Any = None # temporarily unset to appease extra="forbid"
@abstractmethod
def receive_frame(self, frame: Frame, from_network_interface: NetworkInterface):
"""
@@ -40,4 +45,4 @@ class NetworkNode(Node, identifier="NetworkNode"):
:return: ARP Cache for given NetworkNode
:rtype: Optional[ARP]
"""
return self.software_manager.software.get("ARP")
return self.software_manager.software.get("arp")

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

@@ -4,7 +4,7 @@ from __future__ import annotations
import secrets
from enum import Enum
from ipaddress import IPv4Address, IPv4Network
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field, validate_call
@@ -1184,7 +1184,7 @@ class RouterSessionManager(SessionManager):
return outbound_network_interface, dst_mac_address, dst_ip_address, src_port, dst_port, protocol, is_broadcast
class Router(NetworkNode, identifier="router"):
class Router(NetworkNode, discriminator="router"):
"""
Represents a network router, managing routing and forwarding of IP packets across network interfaces.
@@ -1204,8 +1204,13 @@ class Router(NetworkNode, identifier="router"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Routers."""
type: Literal["router"] = "router"
hostname: str = "router"
num_ports: int = 5
acl: Any = None # temporarily unset to appease extra="forbid"
routes: Any = None # temporarily unset to appease extra="forbid"
ports: Any = None # temporarily unset to appease extra="forbid"
default_route: Any = None # temporarily unset to appease extra="forbid"
config: ConfigSchema = Field(default_factory=lambda: Router.ConfigSchema())
@@ -1625,16 +1630,20 @@ class Router(NetworkNode, identifier="router"):
:return: Configured router.
:rtype: Router
"""
ports = config.pop("ports", None)
acl = config.pop("acl", None)
routes = config.pop("routes", None)
default_route = config.pop("default_route", None)
router = Router(config=Router.ConfigSchema(**config))
if "ports" in config:
for port_num, port_cfg in config["ports"].items():
if ports:
for port_num, port_cfg in ports.items():
router.configure_port(
port=port_num,
ip_address=port_cfg["ip_address"],
subnet_mask=IPv4Address(port_cfg.get("subnet_mask", "255.255.255.0")),
)
if "acl" in config:
for r_num, r_cfg in config["acl"].items():
if acl:
for r_num, r_cfg in acl.items():
router.acl.add_rule(
action=ACLAction[r_cfg["action"]],
src_port=None if not (p := r_cfg.get("src_port")) else PORT_LOOKUP[p],
@@ -1646,16 +1655,16 @@ class Router(NetworkNode, identifier="router"):
dst_wildcard_mask=r_cfg.get("dst_wildcard_mask"),
position=r_num,
)
if "routes" in config:
for route in config.get("routes"):
if routes:
for route in routes:
router.route_table.add_route(
address=IPv4Address(route.get("address")),
subnet_mask=IPv4Address(route.get("subnet_mask", "255.255.255.0")),
next_hop_ip_address=IPv4Address(route.get("next_hop_ip_address")),
metric=float(route.get("metric", 0)),
)
if "default_route" in config:
next_hop_ip_address = config["default_route"].get("next_hop_ip_address", None)
if default_route:
next_hop_ip_address = default_route.get("next_hop_ip_address", None)
if next_hop_ip_address:
router.route_table.set_default_route_next_hop_ip_address(next_hop_ip_address)
router.operating_state = (

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

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Optional
from typing import Dict, Literal, Optional
from prettytable import MARKDOWN, PrettyTable
from pydantic import Field
@@ -88,7 +88,7 @@ class SwitchPort(WiredNetworkInterface):
return False
class Switch(NetworkNode, identifier="switch"):
class Switch(NetworkNode, discriminator="switch"):
"""A class representing a Layer 2 network switch."""
network_interfaces: Dict[str, SwitchPort] = {}
@@ -101,6 +101,7 @@ class Switch(NetworkNode, identifier="switch"):
class ConfigSchema(NetworkNode.ConfigSchema):
"""Configuration Schema for Switch nodes within PrimAITE."""
type: Literal["switch"] = "switch"
hostname: str = "Switch"
num_ports: int = 24
"The number of ports on the switch. Default is 24."

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

@@ -1,6 +1,6 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from ipaddress import IPv4Address
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, Literal, Optional, Union
from pydantic import Field, validate_call
@@ -91,7 +91,7 @@ class WirelessAccessPoint(IPWirelessNetworkInterface):
)
class WirelessRouter(Router, identifier="wireless_router"):
class WirelessRouter(Router, discriminator="wireless-router"):
"""
A WirelessRouter class that extends the functionality of a standard Router to include wireless capabilities.
@@ -126,10 +126,13 @@ class WirelessRouter(Router, identifier="wireless_router"):
class ConfigSchema(Router.ConfigSchema):
"""Configuration Schema for WirelessRouter nodes within PrimAITE."""
type: Literal["wireless-router"] = "wireless-router"
hostname: str = "WirelessRouter"
airspace: AirSpace
num_ports: int = 0
router_interface: Any = None # temporarily unset to appease extra="forbid"
wireless_access_point: Any = None # temporarily unset to appease extra="forbid"
airspace: AirSpace
config: ConfigSchema = Field(default_factory=lambda: WirelessRouter.ConfigSchema())
def __init__(self, **kwargs):
@@ -137,7 +140,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
self.connect_nic(
WirelessAccessPoint(
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=kwargs["config"].airspace
ip_address="127.0.0.1", subnet_mask="255.0.0.0", gateway="0.0.0.0", airspace=self.airspace
)
)
@@ -236,7 +239,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
)
@classmethod
def from_config(cls, config: Dict, **kwargs) -> "WirelessRouter":
def from_config(cls, config: Dict, airspace: AirSpace) -> "WirelessRouter":
"""Generate the wireless router from config.
Schema:
@@ -263,7 +266,7 @@ class WirelessRouter(Router, identifier="wireless_router"):
:return: WirelessRouter instance.
:rtype: WirelessRouter
"""
router = cls(config=cls.ConfigSchema(**config))
router = cls(config=cls.ConfigSchema(**config), airspace=airspace)
router.operating_state = (
NodeOperatingState.ON if not (p := config.get("operating_state")) else NodeOperatingState[p.upper()]
)

16
src/primaite/simulator/network/networks.py
View File

@@ -170,13 +170,13 @@ def arcd_uc2_network() -> Network:
client_1.power_on()
network.connect(endpoint_b=client_1.network_interface[1], endpoint_a=switch_2.network_interface[1])
client_1.software_manager.install(DatabaseClient)
db_client_1: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
db_client_1: DatabaseClient = client_1.software_manager.software.get("database-client")
db_client_1.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_1.run()
web_browser_1 = client_1.software_manager.software.get("WebBrowser")
web_browser_1 = client_1.software_manager.software.get("web-browser")
web_browser_1.target_url = "http://arcd.com/users/"
client_1.software_manager.install(DataManipulationBot)
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("data-manipulation-bot")
db_manipulation_bot.configure(
server_ip_address=IPv4Address("192.168.1.14"),
payload="DELETE",
@@ -199,10 +199,10 @@ def arcd_uc2_network() -> Network:
client_2.power_on()
client_2.software_manager.install(DatabaseClient)
db_client_2 = client_2.software_manager.software.get("DatabaseClient")
db_client_2 = client_2.software_manager.software.get("database-client")
db_client_2.configure(server_ip_address=IPv4Address("192.168.1.14"))
db_client_2.run()
web_browser_2 = client_2.software_manager.software.get("WebBrowser")
web_browser_2 = client_2.software_manager.software.get("web-browser")
web_browser_2.target_url = "http://arcd.com/users/"
network.connect(
endpoint_b=client_2.network_interface[1],
@@ -244,7 +244,7 @@ def arcd_uc2_network() -> Network:
network.connect(endpoint_b=database_server.network_interface[1], endpoint_a=switch_1.network_interface[3])
database_server.software_manager.install(DatabaseService)
database_service: DatabaseService = database_server.software_manager.software.get("DatabaseService") # noqa
database_service: DatabaseService = database_server.software_manager.software.get("database-service") # noqa
database_service.start()
database_service.configure_backup(backup_server=IPv4Address("192.168.1.16"))
@@ -264,7 +264,7 @@ def arcd_uc2_network() -> Network:
web_server.power_on()
web_server.software_manager.install(DatabaseClient)
database_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
database_client: DatabaseClient = web_server.software_manager.software.get("database-client")
database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
network.connect(endpoint_b=web_server.network_interface[1], endpoint_a=switch_1.network_interface[2])
database_client.run()
@@ -273,7 +273,7 @@ def arcd_uc2_network() -> Network:
web_server.software_manager.install(WebServer)
# register the web_server to a domain
dns_server_service: DNSServer = domain_controller.software_manager.software.get("DNSServer") # noqa
dns_server_service: DNSServer = domain_controller.software_manager.software.get("dns-server") # noqa
dns_server_service.dns_register("arcd.com", web_server.network_interface[1].ip_address)
# Backup Server

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

@@ -38,8 +38,8 @@ class Simulation(SimComponent):
rm.add_request("network", RequestType(func=self.network._request_manager))
# pass through domain requests to the domain object
rm.add_request("domain", RequestType(func=self.domain._request_manager))
# if 'do_nothing' is requested, just return a success
rm.add_request("do_nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
# if 'do-nothing' is requested, just return a success
rm.add_request("do-nothing", RequestType(func=lambda request, context: RequestResponse(status="success")))
return rm
def describe_state(self) -> Dict:

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

@@ -53,20 +53,20 @@ class Application(IOSoftware, ABC):
_registry: ClassVar[Dict[str, Type["Application"]]] = {}
"""Registry of application types. Automatically populated when subclasses are defined."""
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register an application type.
:param identifier: Uniquely specifies an application class by name. Used for finding items by config.
:type identifier: Optional[str]
:param discriminator: Uniquely specifies an application class by name. Used for finding items by config.
:type discriminator: Optional[str]
:raises ValueError: When attempting to register an application with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
if identifier in cls._registry:
raise ValueError(f"Tried to define new application {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
if discriminator in cls._registry:
raise ValueError(f"Tried to define new application {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "Application":

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

@@ -37,7 +37,7 @@ class DatabaseClientConnection(BaseModel):
@property
def client(self) -> Optional[DatabaseClient]:
"""The DatabaseClient that holds this connection."""
return self.parent_node.software_manager.software.get("DatabaseClient")
return self.parent_node.software_manager.software.get("database-client")
def query(self, sql: str) -> bool:
"""
@@ -61,7 +61,7 @@ class DatabaseClientConnection(BaseModel):
return str(self)
class DatabaseClient(Application, identifier="DatabaseClient"):
class DatabaseClient(Application, discriminator="database-client"):
"""
A DatabaseClient application.
@@ -72,7 +72,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for DatabaseClient."""
type: str = "DatabaseClient"
type: str = "database-client"
db_server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
@@ -97,7 +97,7 @@ class DatabaseClient(Application, identifier="DatabaseClient"):
"""Native Client Connection for using the client directly (similar to psql in a terminal)."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
kwargs["name"] = "database-client"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -44,7 +44,7 @@ class PortScanPayload(SimComponent):
return state
class NMAP(Application, identifier="NMAP"):
class NMAP(Application, discriminator="nmap"):
"""
A class representing the NMAP application for network scanning.
@@ -55,7 +55,7 @@ class NMAP(Application, identifier="NMAP"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for NMAP."""
type: str = "NMAP"
type: str = "nmap"
config: "NMAP.ConfigSchema" = Field(default_factory=lambda: NMAP.ConfigSchema())
@@ -70,7 +70,7 @@ class NMAP(Application, identifier="NMAP"):
}
def __init__(self, **kwargs):
kwargs["name"] = "NMAP"
kwargs["name"] = "nmap"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
super().__init__(**kwargs)

8
src/primaite/simulator/system/applications/red_applications/c2/abstract_c2.py
View File

@@ -162,11 +162,11 @@ class AbstractC2(Application):
:return: An FTPClient object is successful, else None
:rtype: union[FTPClient, None]
"""
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("FTPClient")
ftp_client: Union[FTPClient, None] = self.software_manager.software.get("ftp-client")
if ftp_client is None:
self.sys_log.warning(f"{self.__class__.__name__}: No FTPClient. Attempting to install.")
self.software_manager.install(FTPClient)
ftp_client = self.software_manager.software.get("FTPClient")
ftp_client = self.software_manager.software.get("ftp-client")
# Force start if the service is stopped.
if ftp_client.operating_state == ServiceOperatingState.STOPPED:
@@ -189,11 +189,11 @@ class AbstractC2(Application):
:return: An FTPServer object is successful, else None
:rtype: Optional[FTPServer]
"""
ftp_server: Optional[FTPServer] = self.software_manager.software.get("FTPServer")
ftp_server: Optional[FTPServer] = self.software_manager.software.get("ftp-server")
if ftp_server is None:
self.sys_log.warning(f"{self.__class__.__name__}:No FTPServer installed. Attempting to install FTPServer.")
self.software_manager.install(FTPServer)
ftp_server = self.software_manager.software.get("FTPServer")
ftp_server = self.software_manager.software.get("ftp-server")
# Force start if the service is stopped.
if ftp_server.operating_state == ServiceOperatingState.STOPPED:

26
src/primaite/simulator/system/applications/red_applications/c2/c2_beacon.py
View File

@@ -17,7 +17,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
class C2Beacon(AbstractC2, identifier="C2Beacon"):
class C2Beacon(AbstractC2, discriminator="c2-beacon"):
"""
C2 Beacon Application.
@@ -39,7 +39,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Beacon."""
type: str = "C2Beacon"
type: str = "c2-beacon"
c2_server_ip_address: Optional[IPV4Address] = None
keep_alive_frequency: int = 5
masquerade_protocol: IPProtocol = PROTOCOL_LOOKUP["TCP"]
@@ -54,13 +54,13 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
"The currently in use terminal session."
def __init__(self, **kwargs):
kwargs["name"] = "C2Beacon"
kwargs["name"] = "c2-beacon"
super().__init__(**kwargs)
@property
def _host_terminal(self) -> Optional[Terminal]:
"""Return the Terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("Terminal")
"""Return the terminal that is installed on the same machine as the C2 Beacon."""
host_terminal: Terminal = self.software_manager.software.get("terminal")
if host_terminal is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a terminal on its host.")
return host_terminal
@@ -68,7 +68,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
@property
def _host_ransomware_script(self) -> RansomwareScript:
"""Return the RansomwareScript that is installed on the same machine as the C2 Beacon."""
ransomware_script: RansomwareScript = self.software_manager.software.get("RansomwareScript")
ransomware_script: RansomwareScript = self.software_manager.software.get("ransomware-script")
if ransomware_script is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find installed ransomware on its host.")
return ransomware_script
@@ -300,7 +300,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = RansomwareOpts.model_validate(payload.payload)
@@ -324,7 +324,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ransomware_script is None:
@@ -351,7 +351,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns a tuple containing Request Response returned by the Terminal execute method.
:return: Returns a tuple containing Request Response returned by the terminal execute method.
:rtype: Request Response
"""
if self._host_ftp_server is None:
@@ -372,7 +372,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
)
# Using the terminal to start the FTP Client on the remote machine.
self.terminal_session.execute(command=["service", "start", "FTPClient"])
self.terminal_session.execute(command=["service", "start", "ftp-client"])
# Need to supply to the FTP Client the C2 Beacon's host IP.
host_network_interfaces = self.software_manager.node.network_interfaces
@@ -430,7 +430,7 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
# Using the terminal to send the target data back to the C2 Beacon.
exfil_response: RequestResponse = RequestResponse.from_bool(
self.terminal_session.execute(command=["service", "FTPClient", "send", ftp_opts])
self.terminal_session.execute(command=["service", "ftp-client", "send", ftp_opts])
)
# Validating that we successfully received the target data.
@@ -472,14 +472,14 @@ class C2Beacon(AbstractC2, identifier="C2Beacon"):
def _command_terminal(self, payload: C2Packet) -> RequestResponse:
"""
C2 Command: Terminal.
C2 Command: terminal.
Creates a request that executes a terminal command.
This request is then sent to the terminal service in order to be executed.
:payload C2Packet: The incoming INPUT command.
:type Masquerade Packet: C2Packet.
:return: Returns the Request Response returned by the Terminal execute method.
:return: Returns the Request Response returned by the terminal execute method.
:rtype: Request Response
"""
command_opts = TerminalOpts.model_validate(payload.payload)

6
src/primaite/simulator/system/applications/red_applications/c2/c2_server.py
View File

@@ -16,7 +16,7 @@ from primaite.simulator.system.applications.red_applications.c2 import (
from primaite.simulator.system.applications.red_applications.c2.abstract_c2 import AbstractC2, C2Command, C2Payload
class C2Server(AbstractC2, identifier="C2Server"):
class C2Server(AbstractC2, discriminator="c2-server"):
"""
C2 Server Application.
@@ -37,7 +37,7 @@ class C2Server(AbstractC2, identifier="C2Server"):
class ConfigSchema(AbstractC2.ConfigSchema):
"""ConfigSchema for C2Server."""
type: str = "C2Server"
type: str = "c2-server"
config: ConfigSchema = Field(default_factory=lambda: C2Server.ConfigSchema())
@@ -125,7 +125,7 @@ class C2Server(AbstractC2, identifier="C2Server"):
return rm
def __init__(self, **kwargs):
kwargs["name"] = "C2Server"
kwargs["name"] = "c2-server"
super().__init__(**kwargs)
self.run()

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

@@ -40,13 +40,13 @@ class DataManipulationAttackStage(IntEnum):
"Signifies that the attack has failed."
class DataManipulationBot(Application, identifier="DataManipulationBot"):
class DataManipulationBot(Application, discriminator="data-manipulation-bot"):
"""A bot that simulates a script which performs a SQL injection attack."""
class ConfigSchema(Application.ConfigSchema):
"""Configuration schema for DataManipulationBot."""
type: str = "DataManipulationBot"
type: str = "data-manipulation-bot"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "DELETE"
@@ -65,7 +65,7 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
"Whether to repeat attacking once finished."
def __init__(self, **kwargs):
kwargs["name"] = "DataManipulationBot"
kwargs["name"] = "data-manipulation-bot"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -94,7 +94,7 @@ class DataManipulationBot(Application, identifier="DataManipulationBot"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the DataManipulationBot."""
db_client = self.software_manager.software.get("DatabaseClient")
db_client = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

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

@@ -32,13 +32,13 @@ class DoSAttackStage(IntEnum):
"Attack is completed."
class DoSBot(DatabaseClient, identifier="DoSBot"):
class DoSBot(DatabaseClient, discriminator="dos-bot"):
"""A bot that simulates a Denial of Service attack."""
class ConfigSchema(DatabaseClient.ConfigSchema):
"""ConfigSchema for DoSBot."""
type: str = "DoSBot"
type: str = "dos-bot"
target_ip_address: Optional[IPV4Address] = None
target_port: Port = PORT_LOOKUP["POSTGRES_SERVER"]
payload: Optional[str] = None
@@ -72,7 +72,7 @@ class DoSBot(DatabaseClient, identifier="DoSBot"):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.name = "DoSBot"
self.name = "dos-bot"
self.target_ip_address = self.config.target_ip_address
self.target_port = self.config.target_port
self.payload = self.config.payload

8
src/primaite/simulator/system/applications/red_applications/ransomware_script.py
View File

@@ -14,7 +14,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class RansomwareScript(Application, identifier="RansomwareScript"):
class RansomwareScript(Application, discriminator="ransomware-script"):
"""Ransomware Kill Chain - Designed to be used by the TAP001 Agent on the example layout Network.
:ivar payload: The attack stage query payload. (Default ENCRYPT)
@@ -23,7 +23,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for RansomwareScript."""
type: str = "RansomwareScript"
type: str = "ransomware-script"
server_ip: Optional[IPV4Address] = None
server_password: Optional[str] = None
payload: str = "ENCRYPT"
@@ -38,7 +38,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
"Payload String for the payload stage"
def __init__(self, **kwargs):
kwargs["name"] = "RansomwareScript"
kwargs["name"] = "ransomware-script"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["NONE"]
@@ -63,7 +63,7 @@ class RansomwareScript(Application, identifier="RansomwareScript"):
@property
def _host_db_client(self) -> DatabaseClient:
"""Return the database client that is installed on the same machine as the Ransomware Script."""
db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
db_client: DatabaseClient = self.software_manager.software.get("database-client")
if db_client is None:
self.sys_log.warning(f"{self.__class__.__name__} cannot find a database client on its host.")
return db_client

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

@@ -23,7 +23,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebBrowser(Application, identifier="WebBrowser"):
class WebBrowser(Application, discriminator="web-browser"):
"""
Represents a web browser in the simulation environment.
@@ -33,7 +33,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
class ConfigSchema(Application.ConfigSchema):
"""ConfigSchema for WebBrowser."""
type: str = "WebBrowser"
type: str = "web-browser"
target_url: Optional[str] = None
config: "WebBrowser.ConfigSchema" = Field(default_factory=lambda: WebBrowser.ConfigSchema())
@@ -48,7 +48,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
"""Keep a log of visited websites and information about the visit, such as response code."""
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["name"] = "web-browser"
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
@@ -108,7 +108,7 @@ class WebBrowser(Application, identifier="WebBrowser"):
return False
# get the IP address of the domain name via DNS
dns_client: DNSClient = self.software_manager.software.get("DNSClient")
dns_client: DNSClient = self.software_manager.software.get("dns-client")
domain_exists = dns_client.check_domain_exists(target_domain=parsed_url.hostname)
# if domain does not exist, the request fails

6
src/primaite/simulator/system/core/software_manager.py
View File

@@ -60,12 +60,12 @@ class SoftwareManager:
@property
def arp(self) -> "ARP":
"""Provides access to the ARP service instance, if installed."""
return self.software.get("ARP") # noqa
return self.software.get("arp") # noqa
@property
def icmp(self) -> "ICMP":
"""Provides access to the ICMP service instance, if installed."""
return self.software.get("ICMP") # noqa
return self.software.get("icmp") # noqa
def get_open_ports(self) -> List[Port]:
"""
@@ -244,7 +244,7 @@ class SoftwareManager:
:param session: The transport session the payload originates from.
"""
if payload.__class__.__name__ == "PortScanPayload":
self.software.get("NMAP").receive(payload=payload, session_id=session_id)
self.software.get("nmap").receive(payload=payload, session_id=session_id)
return
main_receiver = self.port_protocol_mapping.get((port, protocol), None)
if main_receiver:

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

@@ -15,7 +15,7 @@ from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import PORT_LOOKUP
class ARP(Service, identifier="ARP"):
class ARP(Service, discriminator="arp"):
"""
The ARP (Address Resolution Protocol) Service.
@@ -26,14 +26,14 @@ class ARP(Service, identifier="ARP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ARP."""
type: str = "ARP"
type: str = "arp"
config: "ARP.ConfigSchema" = Field(default_factory=lambda: ARP.ConfigSchema())
arp: Dict[IPV4Address, ARPEntry] = {}
def __init__(self, **kwargs):
kwargs["name"] = "ARP"
kwargs["name"] = "arp"
kwargs["port"] = PORT_LOOKUP["ARP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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

@@ -19,7 +19,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DatabaseService(Service, identifier="DatabaseService"):
class DatabaseService(Service, discriminator="database-service"):
"""
A class for simulating a generic SQL Server service.
@@ -29,12 +29,12 @@ class DatabaseService(Service, identifier="DatabaseService"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DatabaseService."""
type: str = "DatabaseService"
type: str = "database-service"
backup_server_ip: Optional[IPv4Address] = None
db_password: Optional[str] = None
"""Password that needs to be provided by clients if they want to connect to the DatabaseService."""
config: "DatabaseService.ConfigSchema" = Field(default_factory=lambda: DatabaseService.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DatabaseService.ConfigSchema())
backup_server_ip: IPv4Address = None
"""IP address of the backup server."""
@@ -46,7 +46,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
"""File name of latest backup."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseService"
kwargs["name"] = "database-service"
kwargs["port"] = PORT_LOOKUP["POSTGRES_SERVER"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@@ -70,7 +70,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
"""
super().install()
if not self.parent.software_manager.software.get("FTPClient"):
if not self.parent.software_manager.software.get("ftp-client"):
self.parent.sys_log.info(f"{self.name}: Installing FTPClient to enable database backups")
self.parent.software_manager.install(FTPClient)
@@ -94,7 +94,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(
@@ -128,7 +128,7 @@ class DatabaseService(Service, identifier="DatabaseService"):
return False
software_manager: SoftwareManager = self.software_manager
ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
ftp_client_service: FTPClient = software_manager.software.get("ftp-client")
if not ftp_client_service:
self.sys_log.error(

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

@@ -18,24 +18,22 @@ if TYPE_CHECKING:
_LOGGER = getLogger(__name__)
class DNSClient(Service, identifier="DNSClient"):
class DNSClient(Service, discriminator="dns-client"):
"""Represents a DNS Client as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSClient."""
type: str = "DNSClient"
dns_server: Optional[IPv4Address] = None
type: str = "dns-client"
dns_server: Optional[IPV4Address] = None
"The DNS Server the client sends requests to."
config: ConfigSchema = Field(default_factory=lambda: DNSClient.ConfigSchema())
dns_cache: Dict[str, IPv4Address] = {}
"A dict of known mappings between domain/URLs names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSClient"
kwargs["name"] = "dns-client"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

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

@@ -14,22 +14,22 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class DNSServer(Service, identifier="DNSServer"):
class DNSServer(Service, discriminator="dns-server"):
"""Represents a DNS Server as a Service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for DNSServer."""
type: str = "DNSServer"
type: str = "dns-server"
domain_mapping: dict = {}
config: "DNSServer.ConfigSchema" = Field(default_factory=lambda: DNSServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: DNSServer.ConfigSchema())
dns_table: Dict[str, IPv4Address] = {}
"A dict of mappings between domain names and IPv4 addresses."
def __init__(self, **kwargs):
kwargs["name"] = "DNSServer"
kwargs["name"] = "dns-server"
kwargs["port"] = PORT_LOOKUP["DNS"]
# DNS uses UDP by default
# it switches to TCP when the bytes exceed 512 (or 4096) bytes

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

@@ -18,7 +18,7 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPClient(FTPServiceABC, identifier="FTPClient"):
class FTPClient(FTPServiceABC, discriminator="ftp-client"):
"""
A class for simulating an FTP client service.
@@ -26,15 +26,15 @@ class FTPClient(FTPServiceABC, identifier="FTPClient"):
RFC 959: https://datatracker.ietf.org/doc/html/rfc959
"""
config: "FTPClient.ConfigSchema" = Field(default_factory=lambda: FTPClient.ConfigSchema())
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for FTPClient."""
type: str = "FTPClient"
type: str = "ftp-client"
config: ConfigSchema = Field(default_factory=lambda: FTPClient.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "FTPClient"
kwargs["name"] = "ftp-client"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -12,7 +12,7 @@ from primaite.utils.validation.port import is_valid_port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class FTPServer(FTPServiceABC, identifier="FTPServer"):
class FTPServer(FTPServiceABC, discriminator="ftp-server"):
"""
A class for simulating an FTP server service.
@@ -23,15 +23,14 @@ class FTPServer(FTPServiceABC, identifier="FTPServer"):
class ConfigSchema(FTPServiceABC.ConfigSchema):
"""ConfigSchema for FTPServer."""
type: str = "FTPServer"
type: str = "ftp-server"
server_password: Optional[str] = None
"""Password needed to connect to FTP server. Default is None."""
config: ConfigSchema = Field(default_factory=lambda: FTPServer.ConfigSchema())
server_password: Optional[str] = None
def __init__(self, **kwargs):
kwargs["name"] = "FTPServer"
kwargs["name"] = "ftp-server"
kwargs["port"] = PORT_LOOKUP["FTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)

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

@@ -16,7 +16,7 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class ICMP(Service, identifier="ICMP"):
class ICMP(Service, discriminator="icmp"):
"""
The Internet Control Message Protocol (ICMP) service.
@@ -27,14 +27,14 @@ class ICMP(Service, identifier="ICMP"):
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for ICMP."""
type: str = "ICMP"
type: str = "icmp"
config: "ICMP.ConfigSchema" = Field(default_factory=lambda: ICMP.ConfigSchema())
request_replies: Dict = {}
def __init__(self, **kwargs):
kwargs["name"] = "ICMP"
kwargs["name"] = "icmp"
kwargs["port"] = PORT_LOOKUP["NONE"]
kwargs["protocol"] = PROTOCOL_LOOKUP["ICMP"]
super().__init__(**kwargs)

16
src/primaite/simulator/system/services/icmp/router_icmp.py
View File

@@ -1,13 +1,13 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
# class RouterICMP(ICMP):
# class RouterICMP(icmp):
# """
# A class to represent a router's Internet Control Message Protocol (ICMP) handler.
# A class to represent a router's Internet Control Message Protocol (icmp) handler.
#
# :param sys_log: System log for logging network events and errors.
# :type sys_log: SysLog
# :param arp_cache: The ARP cache for resolving MAC addresses.
# :param arp_cache: The arp cache for resolving MAC addresses.
# :type arp_cache: ARPCache
# :param router: The router to which this ICMP handler belongs.
# :param router: The router to which this icmp handler belongs.
# :type router: Router
# """
#
@@ -19,7 +19,7 @@
#
# def process_icmp(self, frame: Frame, from_network_interface: NIC, is_reattempt: bool = False):
# """
# Process incoming ICMP frames based on ICMP type.
# Process incoming icmp frames based on icmp type.
#
# :param frame: The incoming frame to process.
# :param from_network_interface: The network interface where the frame is coming from.
@@ -36,13 +36,13 @@
# self.sys_log.info(f"Received echo request from {frame.ip.src_ip_address}")
# target_mac_address = self.arp.get_arp_cache_mac_address(frame.ip.src_ip_address)
# src_nic = self.arp.get_arp_cache_network_interface(frame.ip.src_ip_address)
# tcp_header = TCPHeader(src_port=Port["ARP"], dst_port=Port["ARP"])
# tcp_header = TCPHeader(src_port=Port["arp"], dst_port=Port["arp"])
#
# # Network Layer
# ip_packet = IPPacket(
# src_ip_address=network_interface.ip_address,
# dst_ip_address=frame.ip.src_ip_address,
# protocol=IPProtocol["ICMP"],
# protocol=IPProtocol["icmp"],
# )
# # Data Link Layer
# ethernet_header = EthernetHeader(
@@ -54,7 +54,7 @@
# identifier=frame.icmp.identifier,
# sequence=frame.icmp.sequence + 1,
# )
# payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard ICMP 32 bytes size
# payload = secrets.token_urlsafe(int(32 / 1.3)) # Standard icmp 32 bytes size
# frame = Frame(
# ethernet=ethernet_header,
# ip=ip_packet,

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

@@ -9,20 +9,20 @@ from primaite import getLogger
from primaite.simulator.network.protocols.ntp import NTPPacket
from primaite.simulator.system.services.service import Service, ServiceOperatingState
from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPClient(Service, identifier="NTPClient"):
class NTPClient(Service, discriminator="ntp-client"):
"""Represents a NTP client as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPClient."""
type: str = "NTPClient"
ntp_server_ip: Optional[IPv4Address] = None
type: str = "ntp-client"
ntp_server_ip: Optional[IPV4Address] = None
"The NTP server the client sends requests to."
config: ConfigSchema = Field(default_factory=lambda: NTPClient.ConfigSchema())
@@ -30,7 +30,7 @@ class NTPClient(Service, identifier="NTPClient"):
time: Optional[datetime] = None
def __init__(self, **kwargs):
kwargs["name"] = "NTPClient"
kwargs["name"] = "ntp-client"
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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

@@ -13,18 +13,18 @@ from primaite.utils.validation.port import PORT_LOOKUP
_LOGGER = getLogger(__name__)
class NTPServer(Service, identifier="NTPServer"):
class NTPServer(Service, discriminator="ntp-server"):
"""Represents a NTP server as a service."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for NTPServer."""
type: str = "NTPServer"
type: str = "ntp-server"
config: "NTPServer.ConfigSchema" = Field(default_factory=lambda: NTPServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: NTPServer.ConfigSchema())
def __init__(self, **kwargs):
kwargs["name"] = "NTPServer"
kwargs["name"] = "ntp-server"
kwargs["port"] = PORT_LOOKUP["NTP"]
kwargs["protocol"] = PROTOCOL_LOOKUP["UDP"]
super().__init__(**kwargs)

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

@@ -61,22 +61,22 @@ class Service(IOSoftware):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def __init_subclass__(cls, identifier: Optional[str] = None, **kwargs: Any) -> None:
def __init_subclass__(cls, discriminator: Optional[str] = None, **kwargs: Any) -> None:
"""
Register a hostnode type.
:param identifier: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type identifier: str
:param discriminator: Uniquely specifies an hostnode class by name. Used for finding items by config.
:type discriminator: str
:raises ValueError: When attempting to register an hostnode with a name that is already allocated.
"""
super().__init_subclass__(**kwargs)
if identifier is None:
if discriminator is None:
return
# Enforce lowercase registry entries because it makes comparisons everywhere else much easier.
identifier = identifier.lower()
if identifier in cls._registry:
raise ValueError(f"Tried to define new hostnode {identifier}, but this name is already reserved.")
cls._registry[identifier] = cls
discriminator = discriminator.lower()
if discriminator in cls._registry:
raise ValueError(f"Tried to define new hostnode {discriminator}, but this name is already reserved.")
cls._registry[discriminator] = cls
@classmethod
def from_config(cls, config: Dict) -> "Service":

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

@@ -129,13 +129,13 @@ class RemoteTerminalConnection(TerminalClientConnection):
return self.parent_terminal.send(payload=payload, session_id=self.ssh_session_id)
class Terminal(Service, identifier="Terminal"):
class Terminal(Service, discriminator="terminal"):
"""Class used to simulate a generic terminal service. Can be interacted with by other terminals via SSH."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for Terminal."""
type: str = "Terminal"
type: str = "terminal"
config: "Terminal.ConfigSchema" = Field(default_factory=lambda: Terminal.ConfigSchema())
@@ -143,7 +143,7 @@ class Terminal(Service, identifier="Terminal"):
"""Dictionary of connect requests made to remote nodes."""
def __init__(self, **kwargs):
kwargs["name"] = "Terminal"
kwargs["name"] = "terminal"
kwargs["port"] = PORT_LOOKUP["SSH"]
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
super().__init__(**kwargs)
@@ -186,7 +186,7 @@ class Terminal(Service, identifier="Terminal"):
return RequestResponse(status="failure", data={})
rm.add_request(
"node_session_remote_login",
"node-session-remote-login",
request_type=RequestType(func=_remote_login),
)

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

@@ -21,15 +21,15 @@ from primaite.utils.validation.port import Port, PORT_LOOKUP
_LOGGER = getLogger(__name__)
class WebServer(Service, identifier="WebServer"):
class WebServer(Service, discriminator="web-server"):
"""Class used to represent a Web Server Service in simulation."""
class ConfigSchema(Service.ConfigSchema):
"""ConfigSchema for WebServer."""
type: str = "WebServer"
type: str = "web-server"
config: "WebServer.ConfigSchema" = Field(default_factory=lambda: WebServer.ConfigSchema())
config: ConfigSchema = Field(default_factory=lambda: WebServer.ConfigSchema())
response_codes_this_timestep: List[HttpStatusCode] = []
@@ -57,7 +57,7 @@ class WebServer(Service, identifier="WebServer"):
return super().pre_timestep(timestep)
def __init__(self, **kwargs):
kwargs["name"] = "WebServer"
kwargs["name"] = "web-server"
kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
# default for web is port 80
if kwargs.get("port") is None:
@@ -146,7 +146,7 @@ class WebServer(Service, identifier="WebServer"):
def _establish_db_connection(self) -> None:
"""Establish a connection to db."""
db_client = self.software_manager.software.get("DatabaseClient")
db_client = self.software_manager.software.get("database-client")
self.db_connection: DatabaseClientConnection = db_client.get_new_connection()
def send(