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.
"""
@abstractmethod
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager)
self.shape: Dict[str, int] = {"node_id": num_nodes}
self.verb: str # define but don't initialise: defends against children classes not defining this
config: "NodeAbstractAction.ConfigSchema"
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration schema for NodeAbstractAction."""
verb: str = "Node_Abstract_Action"
def form_request(self, node_id: int) -> RequestFormat:
"""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):
"""Action which shuts down a node."""
def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
super().__init__(manager=manager, num_nodes=num_nodes)
self.verb: str = "shutdown"
config: "NodeShutdownAction.ConfigSchema"
class ConfigSchema(NodeAbstractAction.ConfigSchema):
"""Configuration Schema for NodeShutdownAction."""
verb: str = "shutdown"
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."""
class ConfigSchema(AbstractAction.ConfigSchema):
"""Configuration Schema for DoNothingAction."""
"""Configuration Schema for do_nothingAction."""
type: str = "do_nothing"
@@ -44,6 +44,7 @@ class ActionManager:
def __init__(
self,
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
act_map: Optional[
Dict[int, Dict]
] = 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()}
# 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))])
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]:
"""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
def form_request(cls, config: ConfigSchema) -> RequestFormat:
"""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"):

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

@@ -447,7 +447,7 @@ class SharedReward(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:
"""
@@ -455,9 +455,9 @@ class ActionPenalty(AbstractReward):
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
: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
"""
self.action_penalty = action_penalty
@@ -473,7 +473,7 @@ class ActionPenalty(AbstractReward):
:return: Reward value
:rtype: float
"""
if last_action_response.action == "DONOTHING":
if last_action_response.action == "do_nothing":
return self.do_nothing_penalty
else:
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:
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.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)
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]
"""
if timestep < self.next_execution_timestep:
return "DONOTHING", {}
return "do_nothing", {}
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",
"| node_id | node name |\n",
"|---------|------------------|\n",
"| 1 | domain_controller|\n",
"| 2 | web_server |\n",
"| 3 | database_server |\n",
"| 4 | backup_server |\n",
"| 5 | security_suite |\n",
"| 6 | client_1 |\n",
"| 7 | client_2 |\n",
"| 0 | domain_controller|\n",
"| 1 | web_server |\n",
"| 2 | database_server |\n",
"| 3 | backup_server |\n",
"| 4 | security_suite |\n",
"| 5 | client_1 |\n",
"| 6 | client_2 |\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",
"\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.
action = ("DONOTHING", {})
action = ("do_nothing", {})
agent.store_action(action)
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
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)
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
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)
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()
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)
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)
for i in range(100):
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)
for i in range(100):
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
@@ -40,11 +40,11 @@ def test_rng_seed_unset(create_env):
env.reset()
for i in range(100):
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()
for i in range(100):
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

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]
node_obj.operating_state = NodeOperatingState.ON
if act_type == "DONOTHING":
if act_type == "do_nothing":
assert mask[action_num]
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]):
"""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
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)
# Check that before trying to fetch the webpage, the reward is 0.0
agent.store_action(("DONOTHING", {}))
agent.store_action(("do_nothing", {}))
game.step()
assert agent.reward_function.current_reward == 0.0
@@ -149,7 +149,7 @@ def test_action_penalty():
# Create an ActionPenalty Reward
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(
state={},
last_action_response=AgentHistoryItem(
@@ -163,12 +163,12 @@ def test_action_penalty():
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(
state={},
last_action_response=AgentHistoryItem(
timestep=0,
action="DONOTHING",
action="do_nothing",
parameters={},
request=["do_nothing"],
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)
action = ("DONOTHING", {})
action = ("do_nothing", {})
agent.store_action(action)
game.step()
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 (
ActionManager,
DoNothingAction,
do_nothingAction,
NodeServiceDisableAction,
NodeServiceEnableAction,
NodeServicePauseAction,
@@ -18,10 +18,10 @@ from primaite.game.agent.actions import (
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()
action = DoNothingAction(manager=manager)
action = do_nothingAction(manager=manager)
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)
# 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={})
browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -93,7 +93,7 @@ class TestWebpageUnavailabilitySticky:
# THE IMPORTANT BIT
# 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={})
browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -130,7 +130,7 @@ class TestWebpageUnavailabilitySticky:
reward = WebpageUnavailablePenalty("computer", sticky=True)
# 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={})
browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
@@ -153,7 +153,7 @@ class TestWebpageUnavailabilitySticky:
# THE IMPORTANT BIT
# 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={})
state = {"network": {"nodes": {"computer": {"applications": {"WebBrowser": {"history": browser_history}}}}}}
last_action_response = AgentHistoryItem(
@@ -191,7 +191,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=False)
# 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={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem(
@@ -212,7 +212,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# THE IMPORTANT BIT
# 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={})
browser_history = []
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -247,7 +247,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
reward = GreenAdminDatabaseUnreachablePenalty("computer", sticky=True)
# 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={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem(
@@ -268,7 +268,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# THE IMPORTANT BIT
# 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={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
last_action_response = AgentHistoryItem(