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

Update configs to new action naming schema and remove redundant agent config

Browse Source
This commit is contained in:
Marek Wolan
2025-01-16 15:18:13 +00:00
parent 504f4bd134
commit 03cab0fcec
59 changed files with 3291 additions and 3466 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
tests/integration_tests/game_layer/actions/test_application_request_permission.py
View File

@@ -33,22 +33,22 @@ def test_application_cannot_perform_actions_unless_running(game_and_agent_fixtur
browser.close()
assert browser.operating_state == ApplicationOperatingState.CLOSED
action = ("NODE_APPLICATION_SCAN", {"node_id": 0, "application_id": 0})
action = ("node_application_scan", {"node_id": 0, "application_id": 0})
agent.store_action(action)
game.step()
assert browser.operating_state == ApplicationOperatingState.CLOSED
action = ("NODE_APPLICATION_CLOSE", {"node_id": 0, "application_id": 0})
action = ("node_application_close", {"node_id": 0, "application_id": 0})
agent.store_action(action)
game.step()
assert browser.operating_state == ApplicationOperatingState.CLOSED
action = ("NODE_APPLICATION_FIX", {"node_id": 0, "application_id": 0})
action = ("node_application_fix", {"node_id": 0, "application_id": 0})
agent.store_action(action)
game.step()
assert browser.operating_state == ApplicationOperatingState.CLOSED
action = ("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0})
action = ("node_application_execute", {"node_id": 0, "application_id": 0})
agent.store_action(action)
game.step()
assert browser.operating_state == ApplicationOperatingState.CLOSED

12
tests/integration_tests/game_layer/actions/test_c2_suite_actions.py
View File

@@ -46,7 +46,7 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
action = (
"NODE_APPLICATION_INSTALL",
"node_application_install",
{"node_id": 1, "application_name": "C2Beacon"},
)
agent.store_action(action)
@@ -54,7 +54,7 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen
assert agent.history[-1].response.status == "success"
action = (
"CONFIGURE_C2_BEACON",
"configure_c2_beacon",
{
"node_id": 1,
"config": {
@@ -70,7 +70,7 @@ def test_c2_beacon_default(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgen
assert agent.history[-1].response.status == "success"
action = (
"NODE_APPLICATION_EXECUTE",
"node_application_execute",
{"node_id": 1, "application_id": 0},
)
agent.store_action(action)
@@ -122,7 +122,7 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA
assert agent.history[-1].response.status == "success"
action = (
"C2_SERVER_RANSOMWARE_CONFIGURE",
"c2_server_ransomware_configure",
{
"node_id": 0,
"config": {"server_ip_address": "10.0.2.3", "payload": "ENCRYPT"},
@@ -141,7 +141,7 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA
game.step()
action = (
"C2_SERVER_RANSOMWARE_LAUNCH",
"c2_server_ransomware_launch",
{
"node_id": 0,
},
@@ -181,7 +181,7 @@ def test_c2_server_data_exfiltration(game_and_agent_fixture: Tuple[PrimaiteGame,
# C2 Action: Data exfiltrate.
action = (
"C2_SERVER_DATA_EXFILTRATE",
"c2_server_data_exfiltrate",
{
"node_id": 0,
"target_file_name": "database.db",

24
tests/integration_tests/game_layer/actions/test_configure_actions.py
View File

@@ -27,7 +27,7 @@ class TestConfigureDatabaseAction:
def test_configure_ip_password(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")
@@ -61,7 +61,7 @@ class TestConfigureDatabaseAction:
db_client: DatabaseClient = client_1.software_manager.software["DatabaseClient"]
action = (
"CONFIGURE_DATABASE_CLIENT",
"configure_database_client",
{
"node_id": 0,
"config": {
@@ -78,7 +78,7 @@ class TestConfigureDatabaseAction:
def test_configure_password(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")
@@ -87,7 +87,7 @@ class TestConfigureDatabaseAction:
old_ip = db_client.server_ip_address
action = (
"CONFIGURE_DATABASE_CLIENT",
"configure_database_client",
{
"node_id": 0,
"config": {
@@ -120,7 +120,7 @@ class TestConfigureRansomwareScriptAction:
def test_configure_ip_password(self, game_and_agent, config):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction(
agent.action_manager.actions["c2_server_ransomware_configure"] = ConfigureRansomwareScriptAction(
agent.action_manager
)
@@ -134,7 +134,7 @@ class TestConfigureRansomwareScriptAction:
old_payload = ransomware_script.payload
action = (
"CONFIGURE_RANSOMWARE_SCRIPT",
"c2_server_ransomware_configure",
{"node_id": 0, "config": config},
)
agent.store_action(action)
@@ -151,7 +151,7 @@ class TestConfigureRansomwareScriptAction:
def test_invalid_config(self, game_and_agent):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_RANSOMWARE_SCRIPT"] = ConfigureRansomwareScriptAction(
agent.action_manager.actions["c2_server_ransomware_configure"] = ConfigureRansomwareScriptAction(
agent.action_manager
)
@@ -160,7 +160,7 @@ class TestConfigureRansomwareScriptAction:
client_1.software_manager.install(RansomwareScript)
ransomware_script: RansomwareScript = client_1.software_manager.software["RansomwareScript"]
action = (
"CONFIGURE_RANSOMWARE_SCRIPT",
"c2_server_ransomware_configure",
{
"node_id": 0,
"config": {"server_password": "admin123", "bad_option": 70},
@@ -172,17 +172,17 @@ class TestConfigureRansomwareScriptAction:
class TestConfigureDoSBot:
def test_configure_DoSBot(self, game_and_agent):
def test_configure_dos_bot(self, game_and_agent):
game, agent = game_and_agent
agent: ControlledAgent
agent.action_manager.actions["CONFIGURE_DOSBOT"] = ConfigureDoSBotAction(agent.action_manager)
agent.action_manager.actions["configure_dos_bot"] = ConfigureDoSBotAction(agent.action_manager)
client_1 = game.simulation.network.get_node_by_hostname("client_1")
client_1.software_manager.install(DoSBot)
dos_bot: DoSBot = client_1.software_manager.software["DoSBot"]
action = (
"CONFIGURE_DOSBOT",
"configure_dos_bot",
{
"node_id": 0,
"config": {
@@ -239,7 +239,7 @@ class TestConfigureYAML:
assert db_client.server_password == "correct_password"
assert db_client.connect()
def test_configure_ransomware_script(self):
def test_c2_server_ransomware_configure(self):
env = PrimaiteGymEnv(env_config=APP_CONFIG_YAML)
client_2 = env.game.simulation.network.get_node_by_hostname("client_2")
assert client_2.software_manager.software.get("RansomwareScript") is None

14
tests/integration_tests/game_layer/actions/test_file_request_permission.py
View File

@@ -33,7 +33,7 @@ def test_create_file(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
assert client_1.file_system.get_file(folder_name=random_folder, file_name=random_file) is None
action = (
"NODE_FILE_CREATE",
"node_file_create",
{"node_id": 0, "folder_name": random_folder, "file_name": random_file},
)
agent.store_action(action)
@@ -51,7 +51,7 @@ def test_file_delete_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge
assert file.deleted is False
action = (
"NODE_FILE_DELETE",
"node_file_delete",
{"node_id": 0, "folder_id": 0, "file_id": 0},
)
agent.store_action(action)
@@ -72,7 +72,7 @@ def test_file_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent
assert file.visible_health_status == FileSystemItemHealthStatus.GOOD
action = (
"NODE_FILE_SCAN",
"node_file_scan",
{"node_id": 0, "folder_id": 0, "file_id": 0},
)
agent.store_action(action)
@@ -93,7 +93,7 @@ def test_file_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge
assert file.health_status == FileSystemItemHealthStatus.CORRUPT
action = (
"NODE_FILE_REPAIR",
"node_file_repair",
{"node_id": 0, "folder_id": 0, "file_id": 0},
)
agent.store_action(action)
@@ -113,7 +113,7 @@ def test_file_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg
assert file.health_status == FileSystemItemHealthStatus.CORRUPT
action = (
"NODE_FILE_RESTORE",
"node_file_restore",
{"node_id": 0, "folder_id": 0, "file_id": 0},
)
agent.store_action(action)
@@ -132,7 +132,7 @@ def test_file_corrupt_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAg
assert file.health_status == FileSystemItemHealthStatus.GOOD
action = (
"NODE_FILE_CORRUPT",
"node_file_corrupt",
{"node_id": 0, "folder_id": 0, "file_id": 0},
)
agent.store_action(action)
@@ -150,7 +150,7 @@ def test_file_access_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge
assert file.num_access == 0
action = (
"NODE_FILE_ACCESS",
"node_file_access",
{"node_id": 0, "folder_name": file.folder_name, "file_name": file.name},
)
agent.store_action(action)

8
tests/integration_tests/game_layer/actions/test_folder_request_permission.py
View File

@@ -32,7 +32,7 @@ def test_create_folder(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
assert client_1.file_system.get_folder(folder_name=random_folder) is None
action = (
"NODE_FOLDER_CREATE",
"node_folder_create",
{
"node_id": 0,
"folder_name": random_folder,
@@ -60,7 +60,7 @@ def test_folder_scan_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAge
assert folder.visible_health_status == FileSystemItemHealthStatus.GOOD
action = (
"NODE_FOLDER_SCAN",
"node_folder_scan",
{
"node_id": 0, # client_1,
"folder_id": 0, # downloads
@@ -87,7 +87,7 @@ def test_folder_repair_action(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA
assert folder.health_status == FileSystemItemHealthStatus.CORRUPT
action = (
"NODE_FOLDER_REPAIR",
"node_folder_repair",
{
"node_id": 0, # client_1,
"folder_id": 0, # downloads
@@ -111,7 +111,7 @@ def test_folder_restore_action(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert folder.health_status == FileSystemItemHealthStatus.CORRUPT
action = (
"NODE_FOLDER_RESTORE",
"node_folder_restore",
{
"node_id": 0, # client_1,
"folder_id": 0, # downloads

8
tests/integration_tests/game_layer/actions/test_nic_request_permission.py
View File

@@ -29,7 +29,7 @@ def test_nic_cannot_be_turned_off_if_not_on(game_and_agent_fixture: Tuple[Primai
assert nic.enabled is False
action = (
"HOST_NIC_DISABLE",
"host_nic_disable",
{
"node_id": 0, # client_1
"nic_id": 0, # the only nic (eth-1)
@@ -50,7 +50,7 @@ def test_nic_cannot_be_turned_on_if_already_on(game_and_agent_fixture: Tuple[Pri
assert nic.enabled
action = (
"HOST_NIC_ENABLE",
"host_nic_enable",
{
"node_id": 0, # client_1
"nic_id": 0, # the only nic (eth-1)
@@ -71,7 +71,7 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr
assert nic.enabled
action = (
"HOST_NIC_DISABLE",
"host_nic_disable",
{
"node_id": 0, # client_1
"nic_id": 0, # the only nic (eth-1)
@@ -83,7 +83,7 @@ def test_that_a_nic_can_be_enabled_and_disabled(game_and_agent_fixture: Tuple[Pr
assert nic.enabled is False
action = (
"HOST_NIC_ENABLE",
"host_nic_enable",
{
"node_id": 0, # client_1
"nic_id": 0, # the only nic (eth-1)

8
tests/integration_tests/game_layer/actions/test_node_request_permission.py
View File

@@ -29,7 +29,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert client_1.operating_state == NodeOperatingState.ON
# turn it off
action = ("NODE_SHUTDOWN", {"node_id": 0})
action = ("node_shutdown", {"node_id": 0})
agent.store_action(action)
game.step()
@@ -43,7 +43,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert client_1.operating_state == NodeOperatingState.OFF
# turn it on
action = ("NODE_STARTUP", {"node_id": 0})
action = ("node_startup", {"node_id": 0})
agent.store_action(action)
game.step()
@@ -65,7 +65,7 @@ def test_node_cannot_be_started_up_if_node_is_already_on(game_and_agent_fixture:
assert client_1.operating_state == NodeOperatingState.ON
# turn it on
action = ("NODE_STARTUP", {"node_id": 0})
action = ("node_startup", {"node_id": 0})
agent.store_action(action)
game.step()
@@ -87,7 +87,7 @@ def test_node_cannot_be_shut_down_if_node_is_already_off(game_and_agent_fixture:
assert client_1.operating_state == NodeOperatingState.OFF
# turn it ff
action = ("NODE_SHUTDOWN", {"node_id": 0})
action = ("node_shutdown", {"node_id": 0})
agent.store_action(action)
game.step()

18
tests/integration_tests/game_layer/actions/test_service_request_permission.py
View File

@@ -31,7 +31,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
dns_server.pause()
assert dns_server.operating_state == ServiceOperatingState.PAUSED
action = ("NODE_SERVICE_START", {"node_id": 1, "service_id": 0})
action = ("node_service_start", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.PAUSED
@@ -40,7 +40,7 @@ def test_service_start(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent]):
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_START", {"node_id": 1, "service_id": 0})
action = ("node_service_start", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
@@ -54,7 +54,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent])
server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
dns_server = server_1.software_manager.software.get("DNSServer")
action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0})
action = ("node_service_resume", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.RUNNING
@@ -63,7 +63,7 @@ def test_service_resume(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyAgent])
assert dns_server.operating_state == ServiceOperatingState.PAUSED
action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0})
action = ("node_service_resume", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
@@ -80,27 +80,27 @@ def test_service_cannot_perform_actions_unless_running(game_and_agent_fixture: T
dns_server.stop()
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_SCAN", {"node_id": 1, "service_id": 0})
action = ("node_service_scan", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_PAUSE", {"node_id": 1, "service_id": 0})
action = ("node_service_pause", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_RESUME", {"node_id": 1, "service_id": 0})
action = ("node_service_resume", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_RESTART", {"node_id": 1, "service_id": 0})
action = ("node_service_restart", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.STOPPED
action = ("NODE_SERVICE_FIX", {"node_id": 1, "service_id": 0})
action = ("node_service_fix", {"node_id": 1, "service_id": 0})
agent.store_action(action)
game.step()
assert dns_server.operating_state == ServiceOperatingState.STOPPED

6
tests/integration_tests/game_layer/actions/test_terminal_actions.py
View File

@@ -100,7 +100,7 @@ def test_remote_login_change_password(game_and_agent_fixture: Tuple[PrimaiteGame
server_1_um.add_user("user123", "password", is_admin=True)
action = (
"NODE_ACCOUNTS_CHANGE_PASSWORD",
"node_accounts_change_password",
{
"node_id": 1, # server_1
"username": "user123",
@@ -139,7 +139,7 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam
# Change password
action = (
"NODE_ACCOUNTS_CHANGE_PASSWORD",
"node_accounts_change_password",
{
"node_id": 1, # server_1
"username": "user123",
@@ -152,7 +152,7 @@ def test_change_password_logs_out_user(game_and_agent_fixture: Tuple[PrimaiteGam
# Assert that the user cannot execute an action
action = (
"NODE_SEND_REMOTE_COMMAND",
"node_send_remote_command",
{
"node_id": 0,
"remote_ip": str(server_1.network_interface[1].ip_address),

65
tests/integration_tests/game_layer/test_action_mask.py
View File

@@ -15,7 +15,6 @@ def test_mask_contents_correct():
net = sim.network
mask = game.action_mask("defender")
agent = env.agent
node_list = agent.action_manager.node_names
action_map = agent.action_manager.action_map
# CHECK NIC ENABLE/DISABLE ACTIONS
@@ -23,8 +22,8 @@ def test_mask_contents_correct():
mask = game.action_mask("defender")
act_type, act_params = action
if act_type == "NODE_NIC_ENABLE":
node_name = node_list[act_params["node_id"]]
if act_type == "node_nic_enable":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
nic_obj = node_obj.network_interface[act_params["nic_id"] + 1]
assert nic_obj.enabled
@@ -34,8 +33,8 @@ def test_mask_contents_correct():
assert mask[action_num]
nic_obj.enable()
if act_type == "NODE_NIC_DISABLE":
node_name = node_list[act_params["node_id"]]
if act_type == "node_nic_disable":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
nic_obj = node_obj.network_interface[act_params["nic_id"] + 1]
assert nic_obj.enabled
@@ -45,14 +44,14 @@ def test_mask_contents_correct():
assert not mask[action_num]
nic_obj.enable()
if act_type == "ROUTER_ACL_ADDRULE":
if act_type == "router_acl_add_rule":
assert mask[action_num]
if act_type == "ROUTER_ACL_REMOVERULE":
if act_type == "router_acl_remove_rule":
assert mask[action_num]
if act_type == "NODE_RESET":
node_name = node_list[act_params["node_id"]]
if act_type == "node_reset":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
assert node_obj.operating_state is NodeOperatingState.ON
assert mask[action_num]
@@ -61,8 +60,8 @@ def test_mask_contents_correct():
assert not mask[action_num]
node_obj.operating_state = NodeOperatingState.ON
if act_type == "NODE_SHUTDOWN":
node_name = node_list[act_params["node_id"]]
if act_type == "node_shutdown":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
assert node_obj.operating_state is NodeOperatingState.ON
assert mask[action_num]
@@ -71,8 +70,8 @@ def test_mask_contents_correct():
assert not mask[action_num]
node_obj.operating_state = NodeOperatingState.ON
if act_type == "NODE_OS_SCAN":
node_name = node_list[act_params["node_id"]]
if act_type == "node_os_scan":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
assert node_obj.operating_state is NodeOperatingState.ON
assert mask[action_num]
@@ -81,8 +80,8 @@ def test_mask_contents_correct():
assert not mask[action_num]
node_obj.operating_state = NodeOperatingState.ON
if act_type == "NODE_STARTUP":
node_name = node_list[act_params["node_id"]]
if act_type == "node_startup":
node_name = act_params["node_name"]
node_obj = net.get_node_by_hostname(node_name)
assert node_obj.operating_state is NodeOperatingState.ON
assert not mask[action_num]
@@ -94,12 +93,12 @@ def test_mask_contents_correct():
if act_type == "do_nothing":
assert mask[action_num]
if act_type == "NODE_SERVICE_DISABLE":
if act_type == "node_service_disable":
assert mask[action_num]
if act_type in ["NODE_SERVICE_SCAN", "NODE_SERVICE_STOP", "NODE_SERVICE_PAUSE"]:
node_name = node_list[act_params["node_id"]]
service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
if act_type in ["node_service_scan", "node_service_stop", "node_service_pause"]:
node_name = act_params["node_name"]
service_name = act_params["service_name"]
node_obj = net.get_node_by_hostname(node_name)
service_obj = node_obj.software_manager.software.get(service_name)
assert service_obj.operating_state is ServiceOperatingState.RUNNING
@@ -109,9 +108,9 @@ def test_mask_contents_correct():
assert not mask[action_num]
service_obj.operating_state = ServiceOperatingState.RUNNING
if act_type == "NODE_SERVICE_RESUME":
node_name = node_list[act_params["node_id"]]
service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
if act_type == "node_service_resume":
node_name = act_params["node_name"]
service_name = act_params["service_name"]
node_obj = net.get_node_by_hostname(node_name)
service_obj = node_obj.software_manager.software.get(service_name)
assert service_obj.operating_state is ServiceOperatingState.RUNNING
@@ -121,9 +120,9 @@ def test_mask_contents_correct():
assert mask[action_num]
service_obj.operating_state = ServiceOperatingState.RUNNING
if act_type == "NODE_SERVICE_START":
node_name = node_list[act_params["node_id"]]
service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
if act_type == "node_service_start":
node_name = act_params["node_name"]
service_name = act_params["service_name"]
node_obj = net.get_node_by_hostname(node_name)
service_obj = node_obj.software_manager.software.get(service_name)
assert service_obj.operating_state is ServiceOperatingState.RUNNING
@@ -133,9 +132,9 @@ def test_mask_contents_correct():
assert mask[action_num]
service_obj.operating_state = ServiceOperatingState.RUNNING
if act_type == "NODE_SERVICE_ENABLE":
node_name = node_list[act_params["node_id"]]
service_name = agent.action_manager.service_names[act_params["node_id"]][act_params["service_id"]]
if act_type == "node_service_enable":
node_name = act_params["node_name"]
service_name = act_params["service_name"]
node_obj = net.get_node_by_hostname(node_name)
service_obj = node_obj.software_manager.software.get(service_name)
assert service_obj.operating_state is ServiceOperatingState.RUNNING
@@ -145,12 +144,10 @@ def test_mask_contents_correct():
assert mask[action_num]
service_obj.operating_state = ServiceOperatingState.RUNNING
if act_type in ["NODE_FILE_SCAN", "NODE_FILE_CHECKHASH", "NODE_FILE_DELETE"]:
node_name = node_list[act_params["node_id"]]
folder_name = agent.action_manager.get_folder_name_by_idx(act_params["node_id"], act_params["folder_id"])
file_name = agent.action_manager.get_file_name_by_idx(
act_params["node_id"], act_params["folder_id"], act_params["file_id"]
)
if act_type in ["node_file_scan", "node_file_checkhash", "node_file_delete"]:
node_name = act_params["node_name"]
folder_name = act_params["folder_name"]
file_name = act_params["file_name"]
node_obj = net.get_node_by_hostname(node_name)
file_obj = node_obj.file_system.get_file(folder_name, file_name, include_deleted=True)
assert not file_obj.deleted

12
tests/integration_tests/game_layer/test_rewards.py
View File

@@ -38,7 +38,7 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle
assert agent.reward_function.current_reward == 0.0
# Check that successfully fetching the webpage yields a reward of 0.7
agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
agent.store_action(("node_application_execute", {"node_id": 0, "application_id": 0}))
game.step()
assert agent.reward_function.current_reward == 0.7
@@ -50,7 +50,7 @@ def test_WebpageUnavailablePenalty(game_and_agent: tuple[PrimaiteGame, Controlle
src_port=PORT_LOOKUP["HTTP"],
dst_port=PORT_LOOKUP["HTTP"],
)
agent.store_action(("NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}))
agent.store_action(("node_application_execute", {"node_id": 0, "application_id": 0}))
game.step()
assert agent.reward_function.current_reward == -0.7
@@ -83,7 +83,7 @@ def test_uc2_rewards(game_and_agent: tuple[PrimaiteGame, ControlledAgent]):
response = game.simulation.apply_request(request)
state = game.get_sim_state()
ahi = AgentHistoryItem(
timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response
timestep=0, action="node_application_execute", parameters={}, request=request, response=response
)
reward_value = comp.calculate(state, last_action_response=ahi)
assert reward_value == 1.0
@@ -94,7 +94,7 @@ def test_uc2_rewards(game_and_agent: tuple[PrimaiteGame, ControlledAgent]):
response = game.simulation.apply_request(request)
state = game.get_sim_state()
ahi = AgentHistoryItem(
timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=request, response=response
timestep=0, action="node_application_execute", parameters={}, request=request, response=response
)
reward_value = comp.calculate(
state,
@@ -159,7 +159,7 @@ def test_action_penalty():
state={},
last_action_response=AgentHistoryItem(
timestep=0,
action="NODE_APPLICATION_EXECUTE",
action="node_application_execute",
parameters={"node_id": 0, "application_id": 1},
request=["execute"],
response=RequestResponse.from_bool(True),
@@ -197,7 +197,7 @@ def test_action_penalty_e2e(game_and_agent: tuple[PrimaiteGame, ControlledAgent]
game.step()
assert agent.reward_function.current_reward == 0.125
action = ("NODE_FILE_SCAN", {"node_id": 0, "folder_id": 0, "file_id": 0})
action = ("node_file_scan", {"node_id": 0, "folder_id": 0, "file_id": 0})
agent.store_action(action)
game.step()
assert agent.reward_function.current_reward == -0.75