Merge agents and actions branches + fix import / subclass errors
This commit is contained in:
@@ -205,8 +205,6 @@ simulation:
|
||||
port_scan_p_of_success: 0.8
|
||||
services:
|
||||
- type: DNSClient
|
||||
options:
|
||||
dns_server: 192.168.1.10
|
||||
- type: DNSServer
|
||||
options:
|
||||
domain_mapping:
|
||||
|
||||
@@ -33,7 +33,7 @@ agents:
|
||||
observation_space: null
|
||||
action_space:
|
||||
action_list:
|
||||
- type: DONOTHING
|
||||
- type: do_nothing
|
||||
- type: NODE_APPLICATION_EXECUTE
|
||||
options:
|
||||
nodes:
|
||||
@@ -47,7 +47,7 @@ agents:
|
||||
max_applications_per_node: 2
|
||||
action_map:
|
||||
0:
|
||||
action: DONOTHING
|
||||
action: do_nothing
|
||||
options: {}
|
||||
1:
|
||||
action: NODE_APPLICATION_EXECUTE
|
||||
@@ -82,7 +82,7 @@ agents:
|
||||
observation_space: null
|
||||
action_space:
|
||||
action_list:
|
||||
- type: DONOTHING
|
||||
- type: do_nothing
|
||||
- type: NODE_APPLICATION_EXECUTE
|
||||
options:
|
||||
nodes:
|
||||
@@ -96,7 +96,7 @@ agents:
|
||||
max_applications_per_node: 2
|
||||
action_map:
|
||||
0:
|
||||
action: DONOTHING
|
||||
action: do_nothing
|
||||
options: {}
|
||||
1:
|
||||
action: NODE_APPLICATION_EXECUTE
|
||||
@@ -132,7 +132,7 @@ agents:
|
||||
|
||||
action_space:
|
||||
action_list:
|
||||
- type: DONOTHING
|
||||
- type: do_nothing
|
||||
- type: NODE_APPLICATION_EXECUTE
|
||||
options:
|
||||
nodes:
|
||||
@@ -235,7 +235,7 @@ agents:
|
||||
|
||||
action_space:
|
||||
action_list:
|
||||
- type: DONOTHING
|
||||
- type: do_nothing
|
||||
- type: NODE_SERVICE_SCAN
|
||||
- type: NODE_SERVICE_STOP
|
||||
- type: NODE_SERVICE_START
|
||||
@@ -265,7 +265,7 @@ agents:
|
||||
|
||||
action_map:
|
||||
0:
|
||||
action: DONOTHING
|
||||
action: do_nothing
|
||||
options: {}
|
||||
# scan webapp service
|
||||
1:
|
||||
|
||||
@@ -96,155 +96,158 @@ agents:
|
||||
|
||||
action_space:
|
||||
action_list:
|
||||
- type: DONOTHING
|
||||
- type: FIREWALL_ACL_ADDRULE
|
||||
- type: FIREWALL_ACL_REMOVERULE
|
||||
- type: NETWORK_PORT_DISABLE
|
||||
- type: NETWORK_PORT_ENABLE
|
||||
- type: do_nothing
|
||||
- type: firewall_acl_add_rule
|
||||
- type: firewall_acl_remove_rule
|
||||
- type: network_port_disable
|
||||
- type: network_port_enable
|
||||
action_map:
|
||||
0:
|
||||
action: DONOTHING
|
||||
action: do_nothing
|
||||
options: {}
|
||||
1:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
type: firewall_acl_add_rule
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: internal
|
||||
firewall_port_direction: inbound
|
||||
position: 1
|
||||
permission: 1
|
||||
source_ip_id: 2 # client 1
|
||||
dest_ip_id: 1 # ALL
|
||||
source_port_id: 1
|
||||
dest_port_id: 1
|
||||
protocol_id: 1
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
permission: PERMIT
|
||||
src_ip: 192.168.0.10
|
||||
dst_ip: 0.0.0.0
|
||||
src_port: 80
|
||||
dst_port: HTTP
|
||||
protocol_name: TCP
|
||||
src_wildcard: 0
|
||||
dst_wildcard: 0
|
||||
2:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: internal
|
||||
firewall_port_direction: inbound
|
||||
position: 1
|
||||
3:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: internal
|
||||
firewall_port_direction: outbound
|
||||
position: 1
|
||||
permission: 2
|
||||
source_ip_id: 2 # client 1
|
||||
dest_ip_id: 1 # ALL
|
||||
source_port_id: 2
|
||||
dest_port_id: 3
|
||||
protocol_id: 2
|
||||
permission: DENY
|
||||
src_ip: 192.168.0.10 # client 1
|
||||
dest_ip: ALL
|
||||
src_port: ARP
|
||||
dst_port: DNS
|
||||
protocol_name: ICMP
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
4:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: internal
|
||||
firewall_port_direction: outbound
|
||||
position: 1
|
||||
5:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: dmz
|
||||
firewall_port_direction: inbound
|
||||
position: 1
|
||||
permission: 2
|
||||
source_ip_id: 3 # dmz_server
|
||||
dest_ip_id: 2 # client_1
|
||||
source_port_id: 4
|
||||
dest_port_id: 4
|
||||
protocol_id: 4
|
||||
permission: DENY
|
||||
src_ip: 192.168.10.10 # dmz_server
|
||||
dest_ip: 192.168.0.10 # client_1
|
||||
src_port: HTTP
|
||||
dst_port: HTTP
|
||||
protocol_name: UDP
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
6:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: dmz
|
||||
firewall_port_direction: inbound
|
||||
position: 1
|
||||
7:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: dmz
|
||||
firewall_port_direction: outbound
|
||||
position: 2
|
||||
permission: 2
|
||||
source_ip_id: 3 # dmz_server
|
||||
dest_ip_id: 2 # client_1
|
||||
source_port_id: 4
|
||||
dest_port_id: 4
|
||||
protocol_id: 3
|
||||
permission: DENY
|
||||
src_ip: 192.168.10.10 # dmz_server
|
||||
dest_ip: 192.168.0.10 # client_1
|
||||
src_port: HTTP
|
||||
dst_port: HTTP
|
||||
protocol_name: TCP
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
8:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: dmz
|
||||
firewall_port_direction: outbound
|
||||
position: 2
|
||||
9:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: external
|
||||
firewall_port_direction: inbound
|
||||
position: 10
|
||||
permission: 2
|
||||
source_ip_id: 4 # external_computer
|
||||
dest_ip_id: 3 # dmz
|
||||
source_port_id: 5
|
||||
dest_port_id: 5
|
||||
protocol_id: 2
|
||||
permission: DENY
|
||||
src_ip: 192.168.20.10 # external_computer
|
||||
dest_ip: 192.168.10.10 # dmz
|
||||
src_port: POSTGRES_SERVER
|
||||
dst_port: POSTGRES_SERVER
|
||||
protocol_name: ICMP
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
10:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: external
|
||||
firewall_port_direction: inbound
|
||||
position: 10
|
||||
11:
|
||||
action: FIREWALL_ACL_ADDRULE
|
||||
action: firewall_acl_add_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: external
|
||||
firewall_port_direction: outbound
|
||||
position: 1
|
||||
permission: 2
|
||||
source_ip_id: 4 # external_computer
|
||||
dest_ip_id: 2 # client_1
|
||||
source_port_id: 1
|
||||
dest_port_id: 1
|
||||
protocol_id: 1
|
||||
permission: DENY
|
||||
src_ip: 192.168.20.10 # external_computer
|
||||
dest_ip: 192.168.0.10 # client_1
|
||||
src_port: NONE
|
||||
dst_port: NONE
|
||||
protocol_name: none
|
||||
source_wildcard_id: 0
|
||||
dest_wildcard_id: 0
|
||||
12:
|
||||
action: FIREWALL_ACL_REMOVERULE
|
||||
action: firewall_acl_remove_rule
|
||||
options:
|
||||
target_firewall_nodename: firewall
|
||||
firewall_port_name: external
|
||||
firewall_port_direction: outbound
|
||||
position: 1
|
||||
13:
|
||||
action: NETWORK_PORT_DISABLE
|
||||
action: network_port_disable
|
||||
options:
|
||||
type: network_port_disable
|
||||
target_nodename: firewall
|
||||
port_id: 3
|
||||
14:
|
||||
action: NETWORK_PORT_ENABLE
|
||||
action: network_port_enable
|
||||
options:
|
||||
type: network_port_enable
|
||||
target_nodename: firewall
|
||||
port_id: 3
|
||||
options:
|
||||
|
||||
@@ -201,8 +201,6 @@ simulation:
|
||||
port_scan_p_of_success: 0.8
|
||||
services:
|
||||
- type: DNSClient
|
||||
options:
|
||||
dns_server: 192.168.1.10
|
||||
- type: DNSServer
|
||||
options:
|
||||
domain_mapping:
|
||||
@@ -233,8 +231,6 @@ simulation:
|
||||
server_password: arcd
|
||||
services:
|
||||
- type: DNSClient
|
||||
options:
|
||||
dns_server: 192.168.1.10
|
||||
|
||||
links:
|
||||
- endpoint_a_hostname: switch_1
|
||||
|
||||
@@ -34,15 +34,16 @@ agents:
|
||||
max_services_per_node: 1
|
||||
max_applications_per_node: 1
|
||||
action_list:
|
||||
- type: NODE_NMAP_NETWORK_SERVICE_RECON
|
||||
- type: node_network_service_recon
|
||||
action_map:
|
||||
0:
|
||||
action: NODE_NMAP_NETWORK_SERVICE_RECON
|
||||
action: node_network_service_recon
|
||||
options:
|
||||
source_node: client_1
|
||||
target_ip_address: 192.168.10.0/24
|
||||
target_port: 80
|
||||
target_protocol: tcp
|
||||
show: false
|
||||
|
||||
reward_function:
|
||||
reward_components:
|
||||
|
||||
@@ -34,13 +34,14 @@ agents:
|
||||
max_services_per_node: 1
|
||||
max_applications_per_node: 1
|
||||
action_list:
|
||||
- type: NODE_NMAP_PING_SCAN
|
||||
- type: node_nmap_ping_scan
|
||||
action_map:
|
||||
0:
|
||||
action: NODE_NMAP_PING_SCAN
|
||||
action: node_nmap_ping_scan
|
||||
options:
|
||||
source_node: client_1
|
||||
node_name: client_1
|
||||
target_ip_address: 192.168.1.0/24
|
||||
show: False
|
||||
|
||||
reward_function:
|
||||
reward_components:
|
||||
|
||||
@@ -34,19 +34,21 @@ agents:
|
||||
max_services_per_node: 1
|
||||
max_applications_per_node: 1
|
||||
action_list:
|
||||
- type: NODE_NMAP_PORT_SCAN
|
||||
- type: node_nmap_port_scan
|
||||
action_map:
|
||||
0:
|
||||
action: NODE_NMAP_PORT_SCAN
|
||||
action: node_nmap_port_scan
|
||||
options:
|
||||
source_node: client_1
|
||||
target_ip_address: 192.168.10.0/24
|
||||
target_protocol: tcp
|
||||
target_port:
|
||||
- 21
|
||||
- 53
|
||||
- 80
|
||||
- 123
|
||||
- 219
|
||||
show: false
|
||||
|
||||
reward_function:
|
||||
reward_components:
|
||||
|
||||
@@ -210,7 +210,6 @@ simulation:
|
||||
services:
|
||||
- type: DNSClient
|
||||
options:
|
||||
dns_server: 192.168.1.10
|
||||
fix_duration: 3
|
||||
- type: DNSServer
|
||||
options:
|
||||
@@ -251,8 +250,6 @@ simulation:
|
||||
server_password: arcd
|
||||
services:
|
||||
- type: DNSClient
|
||||
options:
|
||||
dns_server: 192.168.1.10
|
||||
|
||||
links:
|
||||
- endpoint_a_hostname: switch_1
|
||||
|
||||
@@ -415,85 +415,58 @@ def game_and_agent():
|
||||
install_stuff_to_sim(sim)
|
||||
|
||||
actions = [
|
||||
{"type": "DONOTHING"},
|
||||
{"type": "NODE_SERVICE_SCAN"},
|
||||
{"type": "NODE_SERVICE_STOP"},
|
||||
{"type": "NODE_SERVICE_START"},
|
||||
{"type": "NODE_SERVICE_PAUSE"},
|
||||
{"type": "NODE_SERVICE_RESUME"},
|
||||
{"type": "NODE_SERVICE_RESTART"},
|
||||
{"type": "NODE_SERVICE_DISABLE"},
|
||||
{"type": "NODE_SERVICE_ENABLE"},
|
||||
{"type": "NODE_SERVICE_FIX"},
|
||||
{"type": "NODE_APPLICATION_EXECUTE"},
|
||||
{"type": "NODE_APPLICATION_SCAN"},
|
||||
{"type": "NODE_APPLICATION_CLOSE"},
|
||||
{"type": "NODE_APPLICATION_FIX"},
|
||||
{"type": "NODE_APPLICATION_INSTALL"},
|
||||
{"type": "NODE_APPLICATION_REMOVE"},
|
||||
{"type": "NODE_FILE_CREATE"},
|
||||
{"type": "NODE_FILE_SCAN"},
|
||||
{"type": "NODE_FILE_CHECKHASH"},
|
||||
{"type": "NODE_FILE_DELETE"},
|
||||
{"type": "NODE_FILE_REPAIR"},
|
||||
{"type": "NODE_FILE_RESTORE"},
|
||||
{"type": "NODE_FILE_CORRUPT"},
|
||||
{"type": "NODE_FILE_ACCESS"},
|
||||
{"type": "NODE_FOLDER_CREATE"},
|
||||
{"type": "NODE_FOLDER_SCAN"},
|
||||
{"type": "NODE_FOLDER_CHECKHASH"},
|
||||
{"type": "NODE_FOLDER_REPAIR"},
|
||||
{"type": "NODE_FOLDER_RESTORE"},
|
||||
{"type": "NODE_OS_SCAN"},
|
||||
{"type": "NODE_SHUTDOWN"},
|
||||
{"type": "NODE_STARTUP"},
|
||||
{"type": "NODE_RESET"},
|
||||
{"type": "ROUTER_ACL_ADDRULE"},
|
||||
{"type": "ROUTER_ACL_REMOVERULE"},
|
||||
{"type": "HOST_NIC_ENABLE"},
|
||||
{"type": "HOST_NIC_DISABLE"},
|
||||
{"type": "NETWORK_PORT_ENABLE"},
|
||||
{"type": "NETWORK_PORT_DISABLE"},
|
||||
{"type": "CONFIGURE_C2_BEACON"},
|
||||
{"type": "C2_SERVER_RANSOMWARE_LAUNCH"},
|
||||
{"type": "C2_SERVER_RANSOMWARE_CONFIGURE"},
|
||||
{"type": "C2_SERVER_TERMINAL_COMMAND"},
|
||||
{"type": "C2_SERVER_DATA_EXFILTRATE"},
|
||||
{"type": "NODE_ACCOUNTS_CHANGE_PASSWORD"},
|
||||
{"type": "SSH_TO_REMOTE"},
|
||||
{"type": "SESSIONS_REMOTE_LOGOFF"},
|
||||
{"type": "NODE_SEND_REMOTE_COMMAND"},
|
||||
{"type": "do_nothing"},
|
||||
{"type": "node_service_scan"},
|
||||
{"type": "node_service_stop"},
|
||||
{"type": "node_service_start"},
|
||||
{"type": "node_service_pause"},
|
||||
{"type": "node_service_resume"},
|
||||
{"type": "node_service_restart"},
|
||||
{"type": "node_service_disable"},
|
||||
{"type": "node_service_enable"},
|
||||
{"type": "node_service_fix"},
|
||||
{"type": "node_application_execute"},
|
||||
{"type": "node_application_scan"},
|
||||
{"type": "node_application_close"},
|
||||
{"type": "node_application_fix"},
|
||||
{"type": "node_application_install"},
|
||||
{"type": "node_application_remove"},
|
||||
{"type": "node_file_create"},
|
||||
{"type": "node_file_scan"},
|
||||
{"type": "node_file_checkhash"},
|
||||
{"type": "node_file_delete"},
|
||||
{"type": "node_file_repair"},
|
||||
{"type": "node_file_restore"},
|
||||
{"type": "node_file_corrupt"},
|
||||
{"type": "node_file_access"},
|
||||
{"type": "node_folder_create"},
|
||||
{"type": "node_folder_scan"},
|
||||
{"type": "node_folder_checkhash"},
|
||||
{"type": "node_folder_repair"},
|
||||
{"type": "node_folder_restore"},
|
||||
{"type": "node_os_scan"},
|
||||
{"type": "node_shutdown"},
|
||||
{"type": "node_startup"},
|
||||
{"type": "node_reset"},
|
||||
{"type": "router_acl_add_rule"},
|
||||
{"type": "router_acl_remove_rule"},
|
||||
{"type": "host_nic_enable"},
|
||||
{"type": "host_nic_disable"},
|
||||
{"type": "network_port_enable"},
|
||||
{"type": "network_port_disable"},
|
||||
{"type": "configure_c2_beacon"},
|
||||
{"type": "c2_server_ransomware_launch"},
|
||||
{"type": "c2_server_ransomware_configure"},
|
||||
{"type": "c2_server_terminal_command"},
|
||||
{"type": "c2_server_data_exfiltrate"},
|
||||
{"type": "node_account_change_password"},
|
||||
{"type": "node_session_remote_login"},
|
||||
{"type": "node_session_remote_logoff"},
|
||||
{"type": "node_send_remote_command"},
|
||||
]
|
||||
|
||||
action_space = ActionManager(
|
||||
actions=actions, # ALL POSSIBLE ACTIONS
|
||||
nodes=[
|
||||
{
|
||||
"node_name": "client_1",
|
||||
"applications": [
|
||||
{"application_name": "WebBrowser"},
|
||||
{"application_name": "DoSBot"},
|
||||
{"application_name": "C2Server"},
|
||||
],
|
||||
"folders": [{"folder_name": "downloads", "files": [{"file_name": "cat.png"}]}],
|
||||
},
|
||||
{
|
||||
"node_name": "server_1",
|
||||
"services": [{"service_name": "DNSServer"}],
|
||||
"applications": [{"application_name": "C2Beacon"}],
|
||||
},
|
||||
{"node_name": "server_2", "services": [{"service_name": "WebServer"}]},
|
||||
{"node_name": "router"},
|
||||
],
|
||||
max_folders_per_node=2,
|
||||
max_files_per_folder=2,
|
||||
max_services_per_node=2,
|
||||
max_applications_per_node=3,
|
||||
max_nics_per_node=2,
|
||||
max_acl_rules=10,
|
||||
protocols=["TCP", "UDP", "ICMP"],
|
||||
ports=["HTTP", "DNS", "ARP"],
|
||||
ip_list=["10.0.1.1", "10.0.1.2", "10.0.2.1", "10.0.2.2", "10.0.2.3"],
|
||||
act_map={},
|
||||
)
|
||||
observation_space = ObservationManager(NestedObservation(components={}))
|
||||
|
||||
@@ -5,6 +5,7 @@ from primaite.config.load import get_extended_config_path
|
||||
from primaite.simulator.network.container import Network
|
||||
from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
|
||||
from primaite.simulator.network.hardware.nodes.host.computer import Computer
|
||||
from tests import TEST_ASSETS_ROOT
|
||||
from tests.integration_tests.configuration_file_parsing import BASIC_CONFIG, DMZ_NETWORK, load_config
|
||||
from tests.integration_tests.extensions.applications.extended_application import ExtendedApplication
|
||||
from tests.integration_tests.extensions.nodes.giga_switch import GigaSwitch
|
||||
@@ -13,11 +14,12 @@ from tests.integration_tests.extensions.nodes.giga_switch import GigaSwitch
|
||||
from tests.integration_tests.extensions.nodes.super_computer import SuperComputer
|
||||
from tests.integration_tests.extensions.services.extended_service import ExtendedService
|
||||
|
||||
CONFIG_PATH = TEST_ASSETS_ROOT / "configs/extended_config.yaml"
|
||||
|
||||
|
||||
def test_extended_example_config():
|
||||
"""Test that the example config can be parsed properly."""
|
||||
config_path = os.path.join("tests", "assets", "configs", "extended_config.yaml")
|
||||
game = load_config(config_path)
|
||||
game = load_config(CONFIG_PATH)
|
||||
network: Network = game.simulation.network
|
||||
|
||||
assert len(network.nodes) == 10 # 10 nodes in example network
|
||||
|
||||
@@ -134,7 +134,7 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA
|
||||
|
||||
# Stepping a few timesteps to allow for the RansowmareScript to finish installing.
|
||||
|
||||
action = ("DONOTHING", {})
|
||||
action = ("do_nothing", {})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
game.step()
|
||||
|
||||
@@ -4,7 +4,7 @@ from ipaddress import IPv4Address
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from primaite.game.agent.actions import (
|
||||
from primaite.game.agent.actions.software import (
|
||||
ConfigureDatabaseClientAction,
|
||||
ConfigureDoSBotAction,
|
||||
ConfigureRansomwareScriptAction,
|
||||
@@ -35,10 +35,10 @@ class TestConfigureDatabaseAction:
|
||||
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
|
||||
|
||||
action = (
|
||||
"CONFIGURE_DATABASE_CLIENT",
|
||||
"configure_database_client",
|
||||
{
|
||||
"node_id": 0,
|
||||
"config": {
|
||||
"node_name": "client_1",
|
||||
"model_config": {
|
||||
"server_ip_address": "192.168.1.99",
|
||||
"server_password": "admin123",
|
||||
},
|
||||
@@ -53,7 +53,7 @@ class TestConfigureDatabaseAction:
|
||||
def test_configure_ip(self, game_and_agent):
|
||||
game, agent = game_and_agent
|
||||
agent: ControlledAgent
|
||||
agent.action_manager.actions["CONFIGURE_DATABASE_CLIENT"] = ConfigureDatabaseClientAction(agent.action_manager)
|
||||
agent.action_manager.actions["configure_database_client"] = ConfigureDatabaseClientAction(agent.action_manager)
|
||||
|
||||
# make sure there is a database client on this node
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1")
|
||||
|
||||
@@ -36,7 +36,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
|
||||
assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN
|
||||
|
||||
for i in range(client_1.shut_down_duration + 1):
|
||||
action = ("DONOTHING", {"node_id": 0})
|
||||
action = ("do_nothing", {"node_id": 0})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -50,7 +50,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
|
||||
assert client_1.operating_state == NodeOperatingState.BOOTING
|
||||
|
||||
for i in range(client_1.start_up_duration + 1):
|
||||
action = ("DONOTHING", {"node_id": 0})
|
||||
action = ("do_nothing", {"node_id": 0})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -80,7 +80,7 @@ def test_node_cannot_be_shut_down_if_node_is_already_off(game_and_agent_fixture:
|
||||
client_1.power_off()
|
||||
|
||||
for i in range(client_1.shut_down_duration + 1):
|
||||
action = ("DONOTHING", {"node_id": 0})
|
||||
action = ("do_nothing", {"node_id": 0})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
|
||||
@@ -24,12 +24,12 @@ def test_rng_seed_set(create_env):
|
||||
env.reset(seed=3)
|
||||
for i in range(100):
|
||||
env.step(0)
|
||||
a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"]
|
||||
a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"]
|
||||
|
||||
env.reset(seed=3)
|
||||
for i in range(100):
|
||||
env.step(0)
|
||||
b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"]
|
||||
b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"]
|
||||
|
||||
assert a == b
|
||||
|
||||
@@ -40,11 +40,11 @@ def test_rng_seed_unset(create_env):
|
||||
env.reset()
|
||||
for i in range(100):
|
||||
env.step(0)
|
||||
a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"]
|
||||
a = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"]
|
||||
|
||||
env.reset()
|
||||
for i in range(100):
|
||||
env.step(0)
|
||||
b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "DONOTHING"]
|
||||
b = [item.timestep for item in env.game.agents["client_2_green_user"].history if item.action != "do_nothing"]
|
||||
|
||||
assert a != b
|
||||
|
||||
@@ -91,7 +91,7 @@ def test_mask_contents_correct():
|
||||
assert mask[action_num]
|
||||
node_obj.operating_state = NodeOperatingState.ON
|
||||
|
||||
if act_type == "DONOTHING":
|
||||
if act_type == "do_nothing":
|
||||
assert mask[action_num]
|
||||
|
||||
if act_type == "NODE_SERVICE_DISABLE":
|
||||
|
||||
@@ -32,10 +32,10 @@ FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.
|
||||
|
||||
|
||||
def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
|
||||
"""Test that the DoNothingAction can form a request and that it is accepted by the simulation."""
|
||||
"""Test that the do_nothingAction can form a request and that it is accepted by the simulation."""
|
||||
game, agent = game_and_agent
|
||||
|
||||
action = ("DONOTHING", {})
|
||||
action = ("do_nothing", {})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -56,7 +56,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy
|
||||
assert svc.health_state_visible == SoftwareHealthState.UNUSED
|
||||
|
||||
# 2: Scan and check that the visible state is now correct
|
||||
action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0})
|
||||
action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert svc.health_state_actual == SoftwareHealthState.GOOD
|
||||
@@ -67,7 +67,7 @@ def test_node_service_scan_integration(game_and_agent: Tuple[PrimaiteGame, Proxy
|
||||
assert svc.health_state_visible == SoftwareHealthState.GOOD
|
||||
|
||||
# 4: Scan and check that the visible state is now correct
|
||||
action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0})
|
||||
action = ("node_service_scan", {"type": "node_service_scan", "node_name": "server_1", "service_name": "DNSServer"})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert svc.health_state_actual == SoftwareHealthState.COMPROMISED
|
||||
@@ -88,7 +88,7 @@ def test_node_service_fix_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
|
||||
svc.health_state_actual = SoftwareHealthState.COMPROMISED
|
||||
|
||||
# 2: Apply a patch action
|
||||
action = ("NODE_SERVICE_FIX", {"node_id": 1, "service_id": 0})
|
||||
action = ("node_service_fix", {"type": "node_service_fix", "node_name": "server_1", "service_name": "DNSServer"})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -96,7 +96,7 @@ def test_node_service_fix_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
|
||||
assert svc.health_state_actual == SoftwareHealthState.FIXING
|
||||
|
||||
# 4: perform a few do-nothing steps and check that the service is now in the good state
|
||||
action = ("DONOTHING", {})
|
||||
action = ("do_nothing", {})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert svc.health_state_actual == SoftwareHealthState.GOOD
|
||||
@@ -121,18 +121,19 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
|
||||
|
||||
# 2: Add a rule to block client 1 from reaching server 2 on router
|
||||
action = (
|
||||
"ROUTER_ACL_ADDRULE",
|
||||
"router_acl_add_rule",
|
||||
{
|
||||
"type": "router_acl_add_rule",
|
||||
"target_router": "router",
|
||||
"position": 4, # 4th rule
|
||||
"permission": 2, # DENY
|
||||
"source_ip_id": 3, # 10.0.1.2 (client_1)
|
||||
"dest_ip_id": 6, # 10.0.2.3 (server_2)
|
||||
"dest_port_id": 1, # ALL
|
||||
"source_port_id": 1, # ALL
|
||||
"protocol_id": 1, # ALL
|
||||
"source_wildcard_id": 0,
|
||||
"dest_wildcard_id": 0,
|
||||
"position": 4,
|
||||
"permission": "DENY",
|
||||
"src_ip": "10.0.1.2",
|
||||
"src_wildcard": 0,
|
||||
"src_port": "HTTP",
|
||||
"dst_ip": "10.0.2.3",
|
||||
"dst_wildcard": 0,
|
||||
"dst_port": "HTTP",
|
||||
"protocol_name": "udp",
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -148,24 +149,27 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
|
||||
|
||||
# 4: Add a rule to block server_1 from reaching server_2 on router (this should not affect comms as they are on same subnet)
|
||||
action = (
|
||||
"ROUTER_ACL_ADDRULE",
|
||||
"router_acl_add_rule",
|
||||
{
|
||||
"type": "router_acl_add_rule",
|
||||
"target_router": "router",
|
||||
"position": 5, # 5th rule
|
||||
"permission": 2, # DENY
|
||||
"source_ip_id": 5, # 10.0.2.2 (server_1)
|
||||
"dest_ip_id": 6, # 10.0.2.3 (server_2)
|
||||
"dest_port_id": 1, # ALL
|
||||
"source_port_id": 1, # ALL
|
||||
"protocol_id": 1, # ALL
|
||||
"source_wildcard_id": 0,
|
||||
"dest_wildcard_id": 0,
|
||||
"permission": "DENY", # DENY
|
||||
"src_ip": "10.0.2.2", # 10.0.2.2 (server_1)
|
||||
"src_wildcard": 0,
|
||||
"source_port": "ALL", # ALL
|
||||
"dst_ip": "10.0.2.3", # 10.0.2.3 (server_2)
|
||||
"dst_wildcard": 0,
|
||||
"dst_port": "ALL", # ALL
|
||||
"protocol_name": "ALL", # ALL
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
print(agent.most_recent_action)
|
||||
game.step()
|
||||
|
||||
print(agent.most_recent_action)
|
||||
# 5: Check that the ACL now has 6 rules, but that server_1 can still ping server_2
|
||||
print(router.acl.show())
|
||||
assert router.acl.num_rules == 6
|
||||
assert server_1.ping("10.0.2.3") # Can ping server_2
|
||||
|
||||
@@ -186,8 +190,9 @@ def test_router_acl_removerule_integration(game_and_agent: Tuple[PrimaiteGame, P
|
||||
|
||||
# 2: Remove rule that allows HTTP traffic across the network
|
||||
action = (
|
||||
"ROUTER_ACL_REMOVERULE",
|
||||
"router_acl_remove_rule",
|
||||
{
|
||||
"type": "router_acl_remove_rule",
|
||||
"target_router": "router",
|
||||
"position": 3, # 4th rule
|
||||
},
|
||||
@@ -219,10 +224,11 @@ def test_host_nic_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
|
||||
|
||||
# 2: Disable the NIC on client_1
|
||||
action = (
|
||||
"HOST_NIC_DISABLE",
|
||||
"host_nic_disable",
|
||||
{
|
||||
"node_id": 0, # client_1
|
||||
"nic_id": 0, # the only nic (eth-1)
|
||||
"type": "host_nic_disable",
|
||||
"node_name": "client_1", # client_1
|
||||
"nic_num": 1, # the only nic (eth-1)
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -250,10 +256,11 @@ def test_host_nic_enable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAg
|
||||
|
||||
# 2: Use action to enable nic
|
||||
action = (
|
||||
"HOST_NIC_ENABLE",
|
||||
"host_nic_enable",
|
||||
{
|
||||
"node_id": 0, # client_1
|
||||
"nic_id": 0, # the only nic (eth-1)
|
||||
"type": "host_nic_enable",
|
||||
"node_name": "client_1", # client_1
|
||||
"nic_num": 1, # the only nic (eth-1)
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -277,11 +284,12 @@ def test_node_file_scan_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAge
|
||||
|
||||
# 2: perform a scan and make sure nothing has changed
|
||||
action = (
|
||||
"NODE_FILE_SCAN",
|
||||
"node_file_scan",
|
||||
{
|
||||
"node_id": 0, # client_1,
|
||||
"folder_id": 0, # downloads,
|
||||
"file_id": 0, # cat.png
|
||||
"type": "node_file_scan",
|
||||
"node_name": "client_1", # client_1,
|
||||
"folder_name": "downloads", # downloads,
|
||||
"file_name": "cat.png", # cat.png
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -314,11 +322,12 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
|
||||
|
||||
# 2: delete the file
|
||||
action = (
|
||||
"NODE_FILE_DELETE",
|
||||
"node_file_delete",
|
||||
{
|
||||
"node_id": 0, # client_1
|
||||
"folder_id": 0, # downloads
|
||||
"file_id": 0, # cat.png
|
||||
"type": "node_file_delete",
|
||||
"node_name": "client_1", # client_1
|
||||
"folder_name": "downloads", # downloads
|
||||
"file_name": "cat.png", # cat.png
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -334,14 +343,16 @@ def test_node_file_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
|
||||
"""Test that a file is created."""
|
||||
game, agent = game_and_agent
|
||||
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1")
|
||||
|
||||
action = (
|
||||
"NODE_FILE_CREATE",
|
||||
"node_file_create",
|
||||
{
|
||||
"node_id": 0,
|
||||
"type": "node_file_create",
|
||||
"node_name": "client_1",
|
||||
"folder_name": "test",
|
||||
"file_name": "file.txt",
|
||||
"force": "False",
|
||||
},
|
||||
)
|
||||
agent.store_action(action)
|
||||
@@ -357,9 +368,10 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
|
||||
|
||||
action = (
|
||||
"NODE_FILE_CREATE",
|
||||
"node_file_create",
|
||||
{
|
||||
"node_id": 0,
|
||||
"type": "node_file_create",
|
||||
"node_name": "client_1",
|
||||
"folder_name": "test",
|
||||
"file_name": "file.txt",
|
||||
},
|
||||
@@ -370,9 +382,10 @@ def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
|
||||
assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 0
|
||||
|
||||
action = (
|
||||
"NODE_FILE_ACCESS",
|
||||
"node_file_access",
|
||||
{
|
||||
"node_id": 0,
|
||||
"type": "node_file_access",
|
||||
"node_name": "client_1",
|
||||
"folder_name": "test",
|
||||
"file_name": "file.txt",
|
||||
},
|
||||
@@ -390,9 +403,10 @@ def test_node_folder_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
|
||||
|
||||
action = (
|
||||
"NODE_FOLDER_CREATE",
|
||||
"node_folder_create",
|
||||
{
|
||||
"node_id": 0,
|
||||
"type": "node_folder_create",
|
||||
"node_name": "client_1",
|
||||
"folder_name": "test",
|
||||
},
|
||||
)
|
||||
@@ -418,8 +432,9 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG
|
||||
|
||||
# 2: Disable the NIC on client_1
|
||||
action = (
|
||||
"NETWORK_PORT_DISABLE",
|
||||
"network_port_disable",
|
||||
{
|
||||
"type": "network_port_disable",
|
||||
"target_nodename": "router", # router
|
||||
"port_id": 1, # port 1
|
||||
},
|
||||
@@ -450,8 +465,9 @@ def test_network_router_port_enable_integration(game_and_agent: Tuple[PrimaiteGa
|
||||
|
||||
# 2: Use action to enable port
|
||||
action = (
|
||||
"NETWORK_PORT_ENABLE",
|
||||
"network_port_enable",
|
||||
{
|
||||
"type": "network_port_enable",
|
||||
"target_nodename": "router", # router
|
||||
"port_id": 1, # port 1
|
||||
},
|
||||
@@ -480,7 +496,10 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P
|
||||
assert browser.health_state_visible == SoftwareHealthState.UNUSED
|
||||
|
||||
# 2: Scan and check that the visible state is now correct
|
||||
action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0})
|
||||
action = (
|
||||
"node_application_scan",
|
||||
{"type": "node_application_scan", "node_name": "client_1", "application_name": "WebBrowser"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert browser.health_state_actual == SoftwareHealthState.GOOD
|
||||
@@ -491,7 +510,10 @@ def test_node_application_scan_integration(game_and_agent: Tuple[PrimaiteGame, P
|
||||
assert browser.health_state_visible == SoftwareHealthState.GOOD
|
||||
|
||||
# 4: Scan and check that the visible state is now correct
|
||||
action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0})
|
||||
action = (
|
||||
"node_application_scan",
|
||||
{"type": "node_application_scan", "node_name": "client_1", "application_name": "WebBrowser"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert browser.health_state_actual == SoftwareHealthState.COMPROMISED
|
||||
@@ -512,7 +534,10 @@ def test_node_application_fix_integration(game_and_agent: Tuple[PrimaiteGame, Pr
|
||||
browser.health_state_actual = SoftwareHealthState.COMPROMISED
|
||||
|
||||
# 2: Apply a fix action
|
||||
action = ("NODE_APPLICATION_FIX", {"node_id": 0, "application_id": 0})
|
||||
action = (
|
||||
"node_application_fix",
|
||||
{"type": "node_application_fix", "node_name": "client_1", "application_name": "WebBrowser"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -520,7 +545,7 @@ def test_node_application_fix_integration(game_and_agent: Tuple[PrimaiteGame, Pr
|
||||
assert browser.health_state_actual == SoftwareHealthState.FIXING
|
||||
|
||||
# 4: perform a few do-nothing steps and check that the application is now in the good state
|
||||
action = ("DONOTHING", {})
|
||||
action = ("do_nothing", {})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert browser.health_state_actual == SoftwareHealthState.GOOD
|
||||
@@ -538,7 +563,10 @@ def test_node_application_close_integration(game_and_agent: Tuple[PrimaiteGame,
|
||||
assert browser.operating_state == ApplicationOperatingState.RUNNING
|
||||
|
||||
# 2: Apply a close action
|
||||
action = ("NODE_APPLICATION_CLOSE", {"node_id": 0, "application_id": 0})
|
||||
action = (
|
||||
"node_application_close",
|
||||
{"type": "node_application_close", "node_name": "client_1", "application_name": "WebBrowser"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -549,7 +577,7 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl
|
||||
"""Test that the NodeApplicationInstallAction and NodeApplicationRemoveAction can form a request and that
|
||||
it is accepted by the simulation.
|
||||
|
||||
When you initiate a install action, the Application will be installed and configured on the node.
|
||||
When you initiate an install action, the Application will be installed and configured on the node.
|
||||
The remove action will uninstall the application from the node."""
|
||||
game, agent = game_and_agent
|
||||
|
||||
@@ -557,13 +585,19 @@ def test_node_application_install_and_uninstall_integration(game_and_agent: Tupl
|
||||
|
||||
assert client_1.software_manager.software.get("DoSBot") is None
|
||||
|
||||
action = ("NODE_APPLICATION_INSTALL", {"node_id": 0, "application_name": "DoSBot"})
|
||||
action = (
|
||||
"node_application_install",
|
||||
{"type": "node_application_install", "node_name": "client_1", "application_name": "DoSBot"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
assert client_1.software_manager.software.get("DoSBot") is not None
|
||||
|
||||
action = ("NODE_APPLICATION_REMOVE", {"node_id": 0, "application_name": "DoSBot"})
|
||||
action = (
|
||||
"node_application_remove",
|
||||
{"type": "node_application_remove", "node_name": "client_1", "application_name": "DoSBot"},
|
||||
)
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
|
||||
@@ -656,9 +690,9 @@ def test_firewall_acl_add_remove_rule_integration():
|
||||
assert firewall.external_outbound_acl.acl[1].action.name == "DENY"
|
||||
assert firewall.external_outbound_acl.acl[1].src_ip_address == IPv4Address("192.168.20.10")
|
||||
assert firewall.external_outbound_acl.acl[1].dst_ip_address == IPv4Address("192.168.0.10")
|
||||
assert firewall.external_outbound_acl.acl[1].dst_port is None
|
||||
assert firewall.external_outbound_acl.acl[1].src_port is None
|
||||
assert firewall.external_outbound_acl.acl[1].protocol is None
|
||||
assert firewall.external_outbound_acl.acl[1].dst_port == PORT_LOOKUP["NONE"]
|
||||
assert firewall.external_outbound_acl.acl[1].src_port == PORT_LOOKUP["NONE"]
|
||||
assert firewall.external_outbound_acl.acl[1].protocol == PROTOCOL_LOOKUP["NONE"]
|
||||
|
||||
env.step(12) # Remove ACL rule from External Outbound
|
||||
assert firewall.external_outbound_acl.num_rules == 1
|
||||
|
||||
@@ -18,12 +18,14 @@ from tests import TEST_ASSETS_ROOT
|
||||
from tests.conftest import ControlledAgent
|
||||
|
||||
|
||||
def test_WebpageUnavailablePenalty(game_and_agent):
|
||||
def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, ControlledAgent]):
|
||||
"""Test that we get the right reward for failing to fetch a website."""
|
||||
# set up the scenario, configure the web browser to the correct url
|
||||
game, agent = game_and_agent
|
||||
agent: ControlledAgent
|
||||
comp = WebpageUnavailablePenalty(node_hostname="client_1")
|
||||
schema = WebpageUnavailablePenalty.ConfigSchema(node_hostname="client_1", sticky=True)
|
||||
comp = WebpageUnavailablePenalty(config=schema)
|
||||
|
||||
client_1 = game.simulation.network.get_node_by_hostname("client_1")
|
||||
browser: WebBrowser = client_1.software_manager.software.get("WebBrowser")
|
||||
browser.run()
|
||||
@@ -31,7 +33,7 @@ def test_WebpageUnavailablePenalty(game_and_agent):
|
||||
agent.reward_function.register_component(comp, 0.7)
|
||||
|
||||
# Check that before trying to fetch the webpage, the reward is 0.0
|
||||
agent.store_action(("DONOTHING", {}))
|
||||
agent.store_action(("do_nothing", {}))
|
||||
game.step()
|
||||
assert agent.reward_function.current_reward == 0.0
|
||||
|
||||
@@ -53,7 +55,7 @@ def test_WebpageUnavailablePenalty(game_and_agent):
|
||||
assert agent.reward_function.current_reward == -0.7
|
||||
|
||||
|
||||
def test_uc2_rewards(game_and_agent):
|
||||
def test_uc2_rewards(game_and_agent: tuple[PrimaiteGame, ControlledAgent]):
|
||||
"""Test that the reward component correctly applies a penalty when the selected client cannot reach the database."""
|
||||
game, agent = game_and_agent
|
||||
agent: ControlledAgent
|
||||
@@ -74,7 +76,8 @@ def test_uc2_rewards(game_and_agent):
|
||||
ACLAction.PERMIT, src_port=PORT_LOOKUP["POSTGRES_SERVER"], dst_port=PORT_LOOKUP["POSTGRES_SERVER"], position=2
|
||||
)
|
||||
|
||||
comp = GreenAdminDatabaseUnreachablePenalty("client_1")
|
||||
schema = GreenAdminDatabaseUnreachablePenalty.ConfigSchema(node_hostname="client_1", sticky=True)
|
||||
comp = GreenAdminDatabaseUnreachablePenalty(config=schema)
|
||||
|
||||
request = ["network", "node", "client_1", "application", "DatabaseClient", "execute"]
|
||||
response = game.simulation.apply_request(request)
|
||||
@@ -139,17 +142,19 @@ def test_action_penalty_loads_from_config():
|
||||
act_penalty_obj = comp[0]
|
||||
if act_penalty_obj is None:
|
||||
pytest.fail("Action penalty reward component was not added to the agent from config.")
|
||||
assert act_penalty_obj.action_penalty == -0.75
|
||||
assert act_penalty_obj.do_nothing_penalty == 0.125
|
||||
assert act_penalty_obj.config.action_penalty == -0.75
|
||||
assert act_penalty_obj.config.do_nothing_penalty == 0.125
|
||||
|
||||
|
||||
def test_action_penalty():
|
||||
"""Test that the action penalty is correctly applied when agent performs any action"""
|
||||
|
||||
# Create an ActionPenalty Reward
|
||||
Penalty = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125)
|
||||
schema = ActionPenalty.ConfigSchema(action_penalty=-0.75, do_nothing_penalty=0.125)
|
||||
# Penalty = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125)
|
||||
Penalty = ActionPenalty(config=schema)
|
||||
|
||||
# Assert that penalty is applied if action isn't DONOTHING
|
||||
# Assert that penalty is applied if action isn't do_nothing
|
||||
reward_value = Penalty.calculate(
|
||||
state={},
|
||||
last_action_response=AgentHistoryItem(
|
||||
@@ -163,12 +168,12 @@ def test_action_penalty():
|
||||
|
||||
assert reward_value == -0.75
|
||||
|
||||
# Assert that no penalty applied for a DONOTHING action
|
||||
# Assert that no penalty applied for a do_nothing action
|
||||
reward_value = Penalty.calculate(
|
||||
state={},
|
||||
last_action_response=AgentHistoryItem(
|
||||
timestep=0,
|
||||
action="DONOTHING",
|
||||
action="do_nothing",
|
||||
parameters={},
|
||||
request=["do_nothing"],
|
||||
response=RequestResponse.from_bool(True),
|
||||
@@ -178,15 +183,16 @@ def test_action_penalty():
|
||||
assert reward_value == 0.125
|
||||
|
||||
|
||||
def test_action_penalty_e2e(game_and_agent):
|
||||
def test_action_penalty_e2e(game_and_agent: tuple[PrimaiteGame, ControlledAgent]):
|
||||
"""Test that we get the right reward for doing actions to fetch a website."""
|
||||
game, agent = game_and_agent
|
||||
agent: ControlledAgent
|
||||
comp = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125)
|
||||
schema = ActionPenalty.ConfigSchema(action_penalty=-0.75, do_nothing_penalty=0.125)
|
||||
comp = ActionPenalty(config=schema)
|
||||
|
||||
agent.reward_function.register_component(comp, 1.0)
|
||||
|
||||
action = ("DONOTHING", {})
|
||||
action = ("do_nothing", {})
|
||||
agent.store_action(action)
|
||||
game.step()
|
||||
assert agent.reward_function.current_reward == 0.125
|
||||
|
||||
@@ -3,9 +3,11 @@ from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from primaite.game.agent.actions import (
|
||||
from primaite.game.agent.actions import ( # DoNothingAction,; NodeServiceDisableAction,; NodeServiceEnableAction,; NodeServicePauseAction,; NodeServiceRestartAction,; NodeServiceResumeAction,; NodeServiceScanAction,; NodeServiceStartAction,; NodeServiceStopAction,
|
||||
ActionManager,
|
||||
DoNothingAction,
|
||||
)
|
||||
from primaite.game.agent.actions.manager import DoNothingAction
|
||||
from primaite.game.agent.actions.service import (
|
||||
NodeServiceDisableAction,
|
||||
NodeServiceEnableAction,
|
||||
NodeServicePauseAction,
|
||||
@@ -18,7 +20,7 @@ from primaite.game.agent.actions import (
|
||||
|
||||
|
||||
def test_do_nothing_action_form_request():
|
||||
"""Test that the DoNothingAction can form a request and that it is correct."""
|
||||
"""Test that the do_nothingAction can form a request and that it is correct."""
|
||||
manager = Mock()
|
||||
|
||||
action = DoNothingAction(manager=manager)
|
||||
|
||||
@@ -28,9 +28,9 @@ def test_probabilistic_agent():
|
||||
|
||||
action_space_cfg = {
|
||||
"action_list": [
|
||||
{"type": "DONOTHING"},
|
||||
{"type": "NODE_APPLICATION_EXECUTE"},
|
||||
{"type": "NODE_FILE_DELETE"},
|
||||
{"type": "do_nothing"},
|
||||
{"type": "node_application_execute"},
|
||||
{"type": "node_file_delete"},
|
||||
],
|
||||
"nodes": [
|
||||
{
|
||||
@@ -48,9 +48,9 @@ def test_probabilistic_agent():
|
||||
"protocols": ["TCP", "UDP", "ICMP"],
|
||||
"ports": ["HTTP", "DNS", "ARP"],
|
||||
"act_map": {
|
||||
0: {"action": "DONOTHING", "options": {}},
|
||||
1: {"action": "NODE_APPLICATION_EXECUTE", "options": {"node_id": 0, "application_id": 0}},
|
||||
2: {"action": "NODE_FILE_DELETE", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
|
||||
0: {"action": "do_nothing", "options": {}},
|
||||
1: {"action": "node_application_execute", "options": {"node_id": 0, "application_id": 0}},
|
||||
2: {"action": "node_file_delete", "options": {"node_id": 0, "folder_id": 0, "file_id": 0}},
|
||||
},
|
||||
"options": {},
|
||||
}
|
||||
@@ -80,11 +80,11 @@ def test_probabilistic_agent():
|
||||
node_file_delete_count = 0
|
||||
for _ in range(N_TRIALS):
|
||||
a = pa.get_action(0)
|
||||
if a == ("DONOTHING", {}):
|
||||
if a == ("do_nothing", {}):
|
||||
do_nothing_count += 1
|
||||
elif a == ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}):
|
||||
elif a == ("node_application_execute", {"node_name": "client_1", "application_name": "WebBrowser"}):
|
||||
node_application_execute_count += 1
|
||||
elif a == ("NODE_FILE_DELETE", {"node_id": 0, "folder_id": 0, "file_id": 0}):
|
||||
elif a == ("node_file_delete", {"node_name": "client_1", "folder_name": "downloads", "file_name": "cat.png"}):
|
||||
node_file_delete_count += 1
|
||||
else:
|
||||
raise AssertionError("Probabilistic agent produced an unexpected action.")
|
||||
|
||||
@@ -11,7 +11,12 @@ from primaite.interface.request import RequestResponse
|
||||
|
||||
class TestWebServer404PenaltySticky:
|
||||
def test_non_sticky(self):
|
||||
reward = WebServer404Penalty("computer", "WebService", sticky=False)
|
||||
schema = WebServer404Penalty.ConfigSchema(
|
||||
node_hostname="computer",
|
||||
service_name="WebService",
|
||||
sticky=False,
|
||||
)
|
||||
reward = WebServer404Penalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
codes = []
|
||||
@@ -38,7 +43,12 @@ class TestWebServer404PenaltySticky:
|
||||
assert reward.calculate(state, last_action_response) == -1.0
|
||||
|
||||
def test_sticky(self):
|
||||
reward = WebServer404Penalty("computer", "WebService", sticky=True)
|
||||
schema = WebServer404Penalty.ConfigSchema(
|
||||
node_hostname="computer",
|
||||
service_name="WebService",
|
||||
sticky=True,
|
||||
)
|
||||
reward = WebServer404Penalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
codes = []
|
||||
@@ -67,10 +77,11 @@ class TestWebServer404PenaltySticky:
|
||||
|
||||
class TestWebpageUnavailabilitySticky:
|
||||
def test_non_sticky(self):
|
||||
reward = WebpageUnavailablePenalty("computer", sticky=False)
|
||||
schema = WebpageUnavailablePenalty.ConfigSchema(node_hostname="computer", sticky=False)
|
||||
reward = WebpageUnavailablePenalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "do_nothing", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
browser_history = []
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
|
||||
@@ -93,7 +104,7 @@ class TestWebpageUnavailabilitySticky:
|
||||
|
||||
# THE IMPORTANT BIT
|
||||
# agent did nothing, because reward is not sticky, it goes back to 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
browser_history = []
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
|
||||
@@ -127,10 +138,11 @@ class TestWebpageUnavailabilitySticky:
|
||||
assert reward.calculate(state, last_action_response) == -1.0
|
||||
|
||||
def test_sticky(self):
|
||||
reward = WebpageUnavailablePenalty("computer", sticky=True)
|
||||
schema = WebpageUnavailablePenalty.ConfigSchema(node_hostname="computer", sticky=True)
|
||||
reward = WebpageUnavailablePenalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
browser_history = []
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
|
||||
@@ -153,7 +165,7 @@ class TestWebpageUnavailabilitySticky:
|
||||
|
||||
# THE IMPORTANT BIT
|
||||
# agent did nothing, because reward is sticky, it stays at 1.0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
|
||||
last_action_response = AgentHistoryItem(
|
||||
@@ -188,10 +200,14 @@ class TestWebpageUnavailabilitySticky:
|
||||
|
||||
class TestGreenAdminDatabaseUnreachableSticky:
|
||||
def test_non_sticky(self):
|
||||
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=False)
|
||||
schema = GreenAdminDatabaseUnreachablePenalty.ConfigSchema(
|
||||
node_hostname="computer",
|
||||
sticky=False,
|
||||
)
|
||||
reward = GreenAdminDatabaseUnreachablePenalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
|
||||
last_action_response = AgentHistoryItem(
|
||||
@@ -212,9 +228,8 @@ class TestGreenAdminDatabaseUnreachableSticky:
|
||||
|
||||
# THE IMPORTANT BIT
|
||||
# agent did nothing, because reward is not sticky, it goes back to 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
browser_history = []
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
|
||||
last_action_response = AgentHistoryItem(
|
||||
timestep=0, action=action, parameters=params, request=request, response=response
|
||||
@@ -244,10 +259,14 @@ class TestGreenAdminDatabaseUnreachableSticky:
|
||||
assert reward.calculate(state, last_action_response) == -1.0
|
||||
|
||||
def test_sticky(self):
|
||||
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=True)
|
||||
schema = GreenAdminDatabaseUnreachablePenalty.ConfigSchema(
|
||||
node_hostname="computer",
|
||||
sticky=True,
|
||||
)
|
||||
reward = GreenAdminDatabaseUnreachablePenalty(config=schema)
|
||||
|
||||
# no response codes yet, reward is 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
|
||||
last_action_response = AgentHistoryItem(
|
||||
@@ -268,7 +287,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
|
||||
|
||||
# THE IMPORTANT BIT
|
||||
# agent did nothing, because reward is not sticky, it goes back to 0
|
||||
action, params, request = "DO_NOTHING", {}, ["DONOTHING"]
|
||||
action, params, request = "DO_NOTHING", {}, ["do_nothing"]
|
||||
response = RequestResponse(status="success", data={})
|
||||
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
|
||||
last_action_response = AgentHistoryItem(
|
||||
|
||||
Reference in New Issue
Block a user