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

#2912 - Replace DONOTHING reference with do_nothing, tweaks following milpac actions

Browse Source
This commit is contained in:
Charlie Crane
2024-12-12 16:08:11 +00:00
parent e40fd053f7
commit 47ed585ee2
16 changed files with 58 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/primaite/_legacy/actions.py
View File

@@ -455,11 +455,12 @@ class NodeAbstractAction(AbstractAction):
Any action which applies to a node and uses node_id as its only parameter can inherit from this base class. Any action which applies to a node and uses node_id as its only parameter can inherit from this base class.
""" """
@abstractmethod config: "NodeAbstractAction.ConfigSchema"
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager) class ConfigSchema(AbstractAction.ConfigSchema):
self.shape: Dict[str, int] = {"node_id": num_nodes} """Configuration schema for NodeAbstractAction."""
self.verb: str # define but don't initialise: defends against children classes not defining this
verb: str = "Node_Abstract_Action"
def form_request(self, node_id: int) -> RequestFormat: def form_request(self, node_id: int) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
@@ -478,9 +479,12 @@ class NodeOSScanAction(NodeAbstractAction):
class NodeShutdownAction(NodeAbstractAction): class NodeShutdownAction(NodeAbstractAction):
"""Action which shuts down a node.""" """Action which shuts down a node."""
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None: config: "NodeShutdownAction.ConfigSchema"
super().__init__(manager=manager, num_nodes=num_nodes)
self.verb: str = "shutdown" class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration Schema for NodeShutdownAction."""
verb: str = "shutdown"
class NodeStartupAction(NodeAbstractAction): class NodeStartupAction(NodeAbstractAction):

5
src/primaite/game/agent/actions/manager.py
View File

@@ -28,7 +28,7 @@ class DoNothingAction(AbstractAction, identifier="do_nothing"):
"""Do Nothing Action.""" """Do Nothing Action."""
class ConfigSchema(AbstractAction.ConfigSchema): class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for DoNothingAction.""" """Configuration Schema for do_nothingAction."""
type: str = "do_nothing" type: str = "do_nothing"
@@ -44,6 +44,7 @@ class ActionManager:
def __init__( def __init__(
self, self,
actions: List[Dict], # stores list of actions available to agent actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
act_map: Optional[ act_map: Optional[
Dict[int, Dict] Dict[int, Dict]
] = None, # allows restricting set of possible actions - TODO: Refactor to be a list? ] = None, # allows restricting set of possible actions - TODO: Refactor to be a list?
@@ -79,6 +80,8 @@ class ActionManager:
self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()} self.action_map = {i: (a["action"], a["options"]) for i, a in act_map.items()}
# make sure all numbers between 0 and N are represented as dict keys in action map # make sure all numbers between 0 and N are represented as dict keys in action map
assert all([i in self.action_map.keys() for i in range(len(self.action_map))]) assert all([i in self.action_map.keys() for i in range(len(self.action_map))])
self.node_names: List[str] = [n["node_name"] for n in nodes]
"""List of node names in this action space. The list order is the mapping between node index and node name."""
def get_action(self, action: int) -> Tuple[str, Dict]: def get_action(self, action: int) -> Tuple[str, Dict]:
"""Produce action in CAOS format.""" """Produce action in CAOS format."""

3
src/primaite/game/agent/actions/node.py
View File

@@ -34,7 +34,8 @@ class NodeAbstractAction(AbstractAction, identifier="node_abstract"):
@classmethod @classmethod
def form_request(cls, config: ConfigSchema) -> RequestFormat: def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation.""" """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
return ["network", "node", config.node_name, cls.config.verb] print(config)
return ["network", "node", config.node_name, config.verb]
class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"): class NodeOSScanAction(NodeAbstractAction, identifier="node_os_scan"):

8
src/primaite/game/agent/rewards.py
View File

@@ -447,7 +447,7 @@ class SharedReward(AbstractReward):
class ActionPenalty(AbstractReward): class ActionPenalty(AbstractReward):
"""Apply a negative reward when taking any action except DONOTHING.""" """Apply a negative reward when taking any action except do_nothing."""
def __init__(self, action_penalty: float, do_nothing_penalty: float) -> None: def __init__(self, action_penalty: float, do_nothing_penalty: float) -> None:
""" """
@@ -455,9 +455,9 @@ class ActionPenalty(AbstractReward):
Reward or penalise agents for doing nothing or taking actions. Reward or penalise agents for doing nothing or taking actions.
:param action_penalty: Reward to give agents for taking any action except DONOTHING :param action_penalty: Reward to give agents for taking any action except do_nothing
:type action_penalty: float :type action_penalty: float
:param do_nothing_penalty: Reward to give agent for taking the DONOTHING action :param do_nothing_penalty: Reward to give agent for taking the do_nothing action
:type do_nothing_penalty: float :type do_nothing_penalty: float
""" """
self.action_penalty = action_penalty self.action_penalty = action_penalty
@@ -473,7 +473,7 @@ class ActionPenalty(AbstractReward):
:return: Reward value :return: Reward value
:rtype: float :rtype: float
""" """
if last_action_response.action == "DONOTHING": if last_action_response.action == "do_nothing":
return self.do_nothing_penalty return self.do_nothing_penalty
else: else:
return self.action_penalty return self.action_penalty

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

@@ -39,7 +39,7 @@ class DataManipulationAgent(AbstractScriptedAgent):
""" """
if timestep < self.next_execution_timestep: if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do NOTHING") self.logger.debug(msg="Performing do NOTHING")
return "DONOTHING", {} return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency) self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
self.logger.info(msg="Performing a data manipulation attack!") self.logger.info(msg="Performing a data manipulation attack!")

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

@@ -81,4 +81,4 @@ class PeriodicAgent(AbstractScriptedAgent):
self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance) self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance)
return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0} return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
return "DONOTHING", {} return "do_nothing", {}

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

@@ -46,7 +46,7 @@ class TAP001(AbstractScriptedAgent):
:rtype: Tuple[str, Dict] :rtype: Tuple[str, Dict]
""" """
if timestep < self.next_execution_timestep: if timestep < self.next_execution_timestep:
return "DONOTHING", {} return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency) self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)

14
src/primaite/notebooks/Data-Manipulation-E2E-Demonstration.ipynb
View File

@@ -165,13 +165,13 @@
"\n", "\n",
"| node_id | node name |\n", "| node_id | node name |\n",
"|---------|------------------|\n", "|---------|------------------|\n",
"| 1 | domain_controller|\n", "| 0 | domain_controller|\n",
"| 2 | web_server |\n", "| 1 | web_server |\n",
"| 3 | database_server |\n", "| 2 | database_server |\n",
"| 4 | backup_server |\n", "| 3 | backup_server |\n",
"| 5 | security_suite |\n", "| 4 | security_suite |\n",
"| 6 | client_1 |\n", "| 5 | client_1 |\n",
"| 7 | client_2 |\n", "| 6 | client_2 |\n",
"\n", "\n",
"Service 1 on node 2 (web_server) corresponds to the Web Server service. Other services are only there for padding to ensure that each node's observation space has the same shape. They are filled with zeroes.\n", "Service 1 on node 2 (web_server) corresponds to the Web Server service. Other services are only there for padding to ensure that each node's observation space has the same shape. They are filled with zeroes.\n",
"\n", "\n",

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

@@ -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. # Stepping a few timesteps to allow for the RansowmareScript to finish installing.
action = ("DONOTHING", {}) action = ("do_nothing", {})
agent.store_action(action) agent.store_action(action)
game.step() game.step()
game.step() game.step()

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

@@ -36,7 +36,7 @@ def test_node_startup_shutdown(game_and_agent_fixture: Tuple[PrimaiteGame, Proxy
assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN assert client_1.operating_state == NodeOperatingState.SHUTTING_DOWN
for i in range(client_1.shut_down_duration + 1): 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) agent.store_action(action)
game.step() 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 assert client_1.operating_state == NodeOperatingState.BOOTING
for i in range(client_1.start_up_duration + 1): 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) agent.store_action(action)
game.step() 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() client_1.power_off()
for i in range(client_1.shut_down_duration + 1): 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) agent.store_action(action)
game.step() game.step()

8
tests/integration_tests/game_layer/test_RNG_seed.py
View File

@@ -24,12 +24,12 @@ def test_rng_seed_set(create_env):
env.reset(seed=3) env.reset(seed=3)
for i in range(100): for i in range(100):
env.step(0) 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) env.reset(seed=3)
for i in range(100): for i in range(100):
env.step(0) 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 assert a == b
@@ -40,11 +40,11 @@ def test_rng_seed_unset(create_env):
env.reset() env.reset()
for i in range(100): for i in range(100):
env.step(0) 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() env.reset()
for i in range(100): for i in range(100):
env.step(0) 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 assert a != b

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

@@ -91,7 +91,7 @@ def test_mask_contents_correct():
assert mask[action_num] assert mask[action_num]
node_obj.operating_state = NodeOperatingState.ON node_obj.operating_state = NodeOperatingState.ON
if act_type == "DONOTHING": if act_type == "do_nothing":
assert mask[action_num] assert mask[action_num]
if act_type == "NODE_SERVICE_DISABLE": if act_type == "NODE_SERVICE_DISABLE":

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

@@ -32,7 +32,7 @@ FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.
def test_do_nothing_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): 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 game, agent = game_and_agent
action = ("do_nothing", {}) action = ("do_nothing", {})

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

@@ -31,7 +31,7 @@ def test_WebpageUnavailablePenalty(game_and_agent):
agent.reward_function.register_component(comp, 0.7) agent.reward_function.register_component(comp, 0.7)
# Check that before trying to fetch the webpage, the reward is 0.0 # Check that before trying to fetch the webpage, the reward is 0.0
agent.store_action(("DONOTHING", {})) agent.store_action(("do_nothing", {}))
game.step() game.step()
assert agent.reward_function.current_reward == 0.0 assert agent.reward_function.current_reward == 0.0
@@ -149,7 +149,7 @@ def test_action_penalty():
# Create an ActionPenalty Reward # Create an ActionPenalty Reward
Penalty = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125) Penalty = ActionPenalty(action_penalty=-0.75, do_nothing_penalty=0.125)
# 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( reward_value = Penalty.calculate(
state={}, state={},
last_action_response=AgentHistoryItem( last_action_response=AgentHistoryItem(
@@ -163,12 +163,12 @@ def test_action_penalty():
assert reward_value == -0.75 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( reward_value = Penalty.calculate(
state={}, state={},
last_action_response=AgentHistoryItem( last_action_response=AgentHistoryItem(
timestep=0, timestep=0,
action="DONOTHING", action="do_nothing",
parameters={}, parameters={},
request=["do_nothing"], request=["do_nothing"],
response=RequestResponse.from_bool(True), response=RequestResponse.from_bool(True),
@@ -186,7 +186,7 @@ def test_action_penalty_e2e(game_and_agent):
agent.reward_function.register_component(comp, 1.0) agent.reward_function.register_component(comp, 1.0)
action = ("DONOTHING", {}) action = ("do_nothing", {})
agent.store_action(action) agent.store_action(action)
game.step() game.step()
assert agent.reward_function.current_reward == 0.125 assert agent.reward_function.current_reward == 0.125

6
tests/unit_tests/_primaite/_game/_agent/test_actions.py
View File

@@ -5,7 +5,7 @@ import pytest
from primaite.game.agent.actions import ( from primaite.game.agent.actions import (
ActionManager, ActionManager,
DoNothingAction, do_nothingAction,
NodeServiceDisableAction, NodeServiceDisableAction,
NodeServiceEnableAction, NodeServiceEnableAction,
NodeServicePauseAction, NodeServicePauseAction,
@@ -18,10 +18,10 @@ from primaite.game.agent.actions import (
def test_do_nothing_action_form_request(): 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() manager = Mock()
action = DoNothingAction(manager=manager) action = do_nothingAction(manager=manager)
request = action.form_request() request = action.form_request()

16
tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py
View File

@@ -70,7 +70,7 @@ class TestWebpageUnavailabilitySticky:
reward = WebpageUnavailablePenalty("computer", sticky=False) reward = WebpageUnavailablePenalty("computer", sticky=False)
# no response codes yet, reward is 0 # 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={}) response = RequestResponse(status="success", data={})
browser_history = [] browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -93,7 +93,7 @@ class TestWebpageUnavailabilitySticky:
# THE IMPORTANT BIT # THE IMPORTANT BIT
# agent did nothing, because reward is not sticky, it goes back to 0 # 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={}) response = RequestResponse(status="success", data={})
browser_history = [] browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -130,7 +130,7 @@ class TestWebpageUnavailabilitySticky:
reward = WebpageUnavailablePenalty("computer", sticky=True) reward = WebpageUnavailablePenalty("computer", sticky=True)
# no response codes yet, reward is 0 # 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={}) response = RequestResponse(status="success", data={})
browser_history = [] browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -153,7 +153,7 @@ class TestWebpageUnavailabilitySticky:
# THE IMPORTANT BIT # THE IMPORTANT BIT
# agent did nothing, because reward is sticky, it stays at 1.0 # 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={}) response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
last_action_response = AgentHistoryItem( last_action_response = AgentHistoryItem(
@@ -191,7 +191,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=False) reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=False)
# no response codes yet, reward is 0 # 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={}) response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem( last_action_response = AgentHistoryItem(
@@ -212,7 +212,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# THE IMPORTANT BIT # THE IMPORTANT BIT
# agent did nothing, because reward is not sticky, it goes back to 0 # 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={}) response = RequestResponse(status="success", data={})
browser_history = [] browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -247,7 +247,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=True) reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=True)
# no response codes yet, reward is 0 # 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={}) response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem( last_action_response = AgentHistoryItem(
@@ -268,7 +268,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# THE IMPORTANT BIT # THE IMPORTANT BIT
# agent did nothing, because reward is not sticky, it goes back to 0 # 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={}) response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}} state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem( last_action_response = AgentHistoryItem(