From 4481d073e4ea34eb47bcded32e9d3a24c81454e7 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 20 Jan 2025 08:35:11 +0000 Subject: [PATCH] Fix action config schemas and formrequest method for dos bot action --- src/primaite/game/agent/actions/network.py | 6 +++--- src/primaite/game/agent/actions/software.py | 2 +- .../game_layer/actions/test_c2_suite_actions.py | 3 ++- tests/integration_tests/game_layer/test_actions.py | 4 ++-- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/primaite/game/agent/actions/network.py b/src/primaite/game/agent/actions/network.py index 7f1e069a..d244fb74 100644 --- a/src/primaite/game/agent/actions/network.py +++ b/src/primaite/game/agent/actions/network.py @@ -17,20 +17,20 @@ class NetworkPortAbstractAction(AbstractAction, identifier="network_port_abstrac """Base configuration schema for NetworkPort actions.""" target_nodename: str - port_id: int + port_num: int verb: ClassVar[str] @classmethod 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_id is None: + if config.target_nodename is None or config.port_num is None: return ["do_nothing"] return [ "network", "node", config.target_nodename, "network_interface", - config.port_id, + config.port_num, config.verb, ] diff --git a/src/primaite/game/agent/actions/software.py b/src/primaite/game/agent/actions/software.py index 23fbd70d..da751a15 100644 --- a/src/primaite/game/agent/actions/software.py +++ b/src/primaite/game/agent/actions/software.py @@ -62,7 +62,7 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): max_sessions: Optional[int] = None @classmethod - def form_request(config: ConfigSchema) -> RequestFormat: + def form_request(cls, config: ConfigSchema) -> RequestFormat: """Return the action formatted as a request that can be ingested by the simulation.""" data = dict( target_ip_address=config.target_ip_address, diff --git a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py index c52c5761..59eb8a60 100644 --- a/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py +++ b/tests/integration_tests/game_layer/actions/test_c2_suite_actions.py @@ -121,7 +121,8 @@ def test_c2_server_ransomware(game_and_agent_fixture: Tuple[PrimaiteGame, ProxyA "c2_server_ransomware_configure", { "node_name": "client_1", - "config": {"server_ip_address": "10.0.2.3", "payload": "ENCRYPT"}, + "server_ip_address": "10.0.2.3", + "payload": "ENCRYPT", }, ) agent.store_action(action) diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index cf8a33ce..80f359da 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -424,7 +424,7 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG "network_port_disable", { "target_nodename": "router", # router - "port_id": 1, # port 1 + "port_num": 2, # port 1 }, ) agent.store_action(action) @@ -456,7 +456,7 @@ def test_network_router_port_enable_integration(game_and_agent: Tuple[PrimaiteGa "network_port_enable", { "target_nodename": "router", # router - "port_id": 1, # port 1 + "port_num": 2, # port 1 }, ) agent.store_action(action)