From 18a665e56246b376bade49ae786ca41ed8d5836c Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 20 Jan 2025 14:07:51 +0000 Subject: [PATCH] Update actions and agents to get all tests passing post-refactor --- src/primaite/game/agent/actions/software.py | 1 + .../scripted_agents/probabilistic_agent.py | 2 +- src/primaite/game/game.py | 27 ------------------- .../configs/firewall_actions_network.yaml | 16 +++++------ .../game_layer/test_actions.py | 16 +++++------ 5 files changed, 18 insertions(+), 44 deletions(-) diff --git a/src/primaite/game/agent/actions/software.py b/src/primaite/game/agent/actions/software.py index aeb7c582..12462ede 100644 --- a/src/primaite/game/agent/actions/software.py +++ b/src/primaite/game/agent/actions/software.py @@ -84,6 +84,7 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"): dos_intensity=config.dos_intensity, 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] diff --git a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py index 959eaadc..de643ed8 100644 --- a/src/primaite/game/agent/scripted_agents/probabilistic_agent.py +++ b/src/primaite/game/agent/scripted_agents/probabilistic_agent.py @@ -16,7 +16,7 @@ __all__ = "ProbabilisticAgent" class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): """Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" - rng: Generator = np.random.default_rng(np.random.randint(0, 65535)) + rng: Generator = Field(default_factory=lambda: np.random.default_rng(np.random.randint(0, 65535))) class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema): """Schema for the `agent_settings` part of the agent config.""" diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 198e8750..8bc37597 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -476,33 +476,6 @@ class PrimaiteGame: if isinstance(new_agent, ProxyAgent): game.rl_agents[agent_cfg["ref"]] = new_agent - # agent_name = agent_cfg["ref"] # noqa: F841 - # agent_type = agent_cfg["type"] - # action_space_cfg = agent_cfg["action_space"] - # observation_space_cfg = agent_cfg["observation_space"] - # reward_function_cfg = agent_cfg["reward_function"] - # agent_settings = agent_cfg["agent_settings"] - - # agent_config = { - # "type": agent_type, - # "action_space": action_space_cfg, - # "observation_space": observation_space_cfg, - # "reward_function": reward_function_cfg, - # "agent_settings": agent_settings, - # "game": game, - # } - - # # CREATE AGENT - # if agent_type in AbstractAgent._registry: - # new_agent = AbstractAgent._registry[agent_cfg["type"]].from_config(config=agent_config) - # # If blue agent is created, add to game.rl_agents - # if agent_type == "ProxyAgent": - # game.rl_agents[agent_cfg["ref"]] = new_agent - # else: - # msg = f"Configuration error: {agent_type} is not a valid agent type." - # _LOGGER.error(msg) - # raise ValueError(msg) - # Validate that if any agents are sharing rewards, they aren't forming an infinite loop. game.setup_reward_sharing() diff --git a/tests/assets/configs/firewall_actions_network.yaml b/tests/assets/configs/firewall_actions_network.yaml index ff8e784d..6b454a12 100644 --- a/tests/assets/configs/firewall_actions_network.yaml +++ b/tests/assets/configs/firewall_actions_network.yaml @@ -110,9 +110,9 @@ agents: permission: PERMIT src_ip: 192.168.0.10 dst_ip: ALL - src_port: 80 - dst_port: HTTP - protocol_name: TCP + src_port: ALL + dst_port: ALL + protocol_name: ALL src_wildcard: NONE dst_wildcard: NONE 2: @@ -131,7 +131,7 @@ agents: position: 1 permission: DENY src_ip: 192.168.0.10 # client 1 - dest_ip: ALL + dst_ip: ALL src_port: ARP dst_port: DNS protocol_name: ICMP @@ -153,7 +153,7 @@ agents: position: 1 permission: DENY src_ip: 192.168.10.10 # dmz_server - dest_ip: 192.168.0.10 # client_1 + dst_ip: 192.168.0.10 # client_1 src_port: HTTP dst_port: HTTP protocol_name: UDP @@ -175,7 +175,7 @@ agents: position: 2 permission: DENY src_ip: 192.168.10.10 # dmz_server - dest_ip: 192.168.0.10 # client_1 + dst_ip: 192.168.0.10 # client_1 src_port: HTTP dst_port: HTTP protocol_name: TCP @@ -197,7 +197,7 @@ agents: position: 10 permission: DENY src_ip: 192.168.20.10 # external_computer - dest_ip: 192.168.10.10 # dmz + dst_ip: 192.168.10.10 # dmz src_port: POSTGRES_SERVER dst_port: POSTGRES_SERVER protocol_name: ICMP @@ -219,7 +219,7 @@ agents: position: 1 permission: DENY src_ip: 192.168.20.10 # external_computer - dest_ip: 192.168.0.10 # client_1 + dst_ip: 192.168.0.10 # client_1 src_port: NONE dst_port: NONE protocol_name: none diff --git a/tests/integration_tests/game_layer/test_actions.py b/tests/integration_tests/game_layer/test_actions.py index cef41e1b..800549bc 100644 --- a/tests/integration_tests/game_layer/test_actions.py +++ b/tests/integration_tests/game_layer/test_actions.py @@ -127,12 +127,12 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox "position": 4, "permission": "DENY", "src_ip": "10.0.1.2", - "src_wildcard": 0, - "src_port": "HTTP", + "src_wildcard": "NONE", + "src_port": "ALL", "dst_ip": "10.0.2.3", - "dst_wildcard": 0, - "dst_port": "HTTP", - "protocol_name": "udp", + "dst_wildcard": "NONE", + "dst_port": "ALL", + "protocol_name": "icmp", }, ) agent.store_action(action) @@ -155,7 +155,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox "permission": "DENY", # DENY "src_ip": "10.0.2.2", # 10.0.2.2 (server_1) "src_wildcard": 0, - "source_port": "ALL", # ALL + "src_port": "ALL", # ALL "dst_ip": "10.0.2.3", # 10.0.2.3 (server_2) "dst_wildcard": 0, "dst_port": "ALL", # ALL @@ -424,7 +424,7 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG "network_port_disable", { "target_nodename": "router", # router - "port_num": 2, # port 1 + "port_num": 1, # 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_num": 2, # port 1 + "port_num": 1, # port 1 }, ) agent.store_action(action)