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

Update actions and agents to get all tests passing post-refactor

Browse Source
This commit is contained in:
Marek Wolan
2025-01-20 14:07:51 +00:00
parent c521528a0a
commit 18a665e562
5 changed files with 18 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -84,6 +84,7 @@ class ConfigureDoSBotAction(AbstractAction, identifier="configure_dos_bot"):
dos_intensity=config.dos_intensity, dos_intensity=config.dos_intensity,
max_sessions=config.max_sessions, 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", "DoSBot", "configure", data]

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

@@ -16,7 +16,7 @@ __all__ = "ProbabilisticAgent"
class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"): class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" """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): class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config.""" """Schema for the `agent_settings` part of the agent config."""

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

@@ -476,33 +476,6 @@ class PrimaiteGame:
if isinstance(new_agent, ProxyAgent): if isinstance(new_agent, ProxyAgent):
game.rl_agents[agent_cfg["ref"]] = new_agent 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. # Validate that if any agents are sharing rewards, they aren't forming an infinite loop.
game.setup_reward_sharing() game.setup_reward_sharing()

16
tests/assets/configs/firewall_actions_network.yaml
View File

@@ -110,9 +110,9 @@ agents:
permission: PERMIT permission: PERMIT
src_ip: 192.168.0.10 src_ip: 192.168.0.10
dst_ip: ALL dst_ip: ALL
src_port: 80 src_port: ALL
dst_port: HTTP dst_port: ALL
protocol_name: TCP protocol_name: ALL
src_wildcard: NONE src_wildcard: NONE
dst_wildcard: NONE dst_wildcard: NONE
2: 2:
@@ -131,7 +131,7 @@ agents:
position: 1 position: 1
permission: DENY permission: DENY
src_ip: 192.168.0.10 # client 1 src_ip: 192.168.0.10 # client 1
dest_ip: ALL dst_ip: ALL
src_port: ARP src_port: ARP
dst_port: DNS dst_port: DNS
protocol_name: ICMP protocol_name: ICMP
@@ -153,7 +153,7 @@ agents:
position: 1 position: 1
permission: DENY permission: DENY
src_ip: 192.168.10.10 # dmz_server 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 src_port: HTTP
dst_port: HTTP dst_port: HTTP
protocol_name: UDP protocol_name: UDP
@@ -175,7 +175,7 @@ agents:
position: 2 position: 2
permission: DENY permission: DENY
src_ip: 192.168.10.10 # dmz_server 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 src_port: HTTP
dst_port: HTTP dst_port: HTTP
protocol_name: TCP protocol_name: TCP
@@ -197,7 +197,7 @@ agents:
position: 10 position: 10
permission: DENY permission: DENY
src_ip: 192.168.20.10 # external_computer 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 src_port: POSTGRES_SERVER
dst_port: POSTGRES_SERVER dst_port: POSTGRES_SERVER
protocol_name: ICMP protocol_name: ICMP
@@ -219,7 +219,7 @@ agents:
position: 1 position: 1
permission: DENY permission: DENY
src_ip: 192.168.20.10 # external_computer 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 src_port: NONE
dst_port: NONE dst_port: NONE
protocol_name: none protocol_name: none

16
tests/integration_tests/game_layer/test_actions.py
View File

@@ -127,12 +127,12 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
"position": 4, "position": 4,
"permission": "DENY", "permission": "DENY",
"src_ip": "10.0.1.2", "src_ip": "10.0.1.2",
"src_wildcard": 0, "src_wildcard": "NONE",
"src_port": "HTTP", "src_port": "ALL",
"dst_ip": "10.0.2.3", "dst_ip": "10.0.2.3",
"dst_wildcard": 0, "dst_wildcard": "NONE",
"dst_port": "HTTP", "dst_port": "ALL",
"protocol_name": "udp", "protocol_name": "icmp",
}, },
) )
agent.store_action(action) agent.store_action(action)
@@ -155,7 +155,7 @@ def test_router_acl_addrule_integration(game_and_agent: Tuple[PrimaiteGame, Prox
"permission": "DENY", # DENY "permission": "DENY", # DENY
"src_ip": "10.0.2.2", # 10.0.2.2 (server_1) "src_ip": "10.0.2.2", # 10.0.2.2 (server_1)
"src_wildcard": 0, "src_wildcard": 0,
"source_port": "ALL", # ALL "src_port": "ALL", # ALL
"dst_ip": "10.0.2.3", # 10.0.2.3 (server_2) "dst_ip": "10.0.2.3", # 10.0.2.3 (server_2)
"dst_wildcard": 0, "dst_wildcard": 0,
"dst_port": "ALL", # ALL "dst_port": "ALL", # ALL
@@ -424,7 +424,7 @@ def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteG
"network_port_disable", "network_port_disable",
{ {
"target_nodename": "router", # router "target_nodename": "router", # router
"port_num": 2, # port 1 "port_num": 1, # port 1
}, },
) )
agent.store_action(action) agent.store_action(action)
@@ -456,7 +456,7 @@ def test_network_router_port_enable_integration(game_and_agent: Tuple[PrimaiteGa
"network_port_enable", "network_port_enable",
{ {
"target_nodename": "router", # router "target_nodename": "router", # router
"port_num": 2, # port 1 "port_num": 1, # port 1
}, },
) )
agent.store_action(action) agent.store_action(action)