From daa34385e550dc43e984706faa2df024e74d1ad7 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 9 Jan 2024 14:53:15 +0000 Subject: [PATCH 1/4] Add agent reset for episodes --- src/primaite/game/agent/data_manipulation_bot.py | 9 +++++++++ src/primaite/game/agent/interface.py | 4 ++++ src/primaite/game/game.py | 1 + 3 files changed, 14 insertions(+) diff --git a/src/primaite/game/agent/data_manipulation_bot.py b/src/primaite/game/agent/data_manipulation_bot.py index 8237ce06..3b558087 100644 --- a/src/primaite/game/agent/data_manipulation_bot.py +++ b/src/primaite/game/agent/data_manipulation_bot.py @@ -15,6 +15,7 @@ class DataManipulationAgent(AbstractScriptedAgent): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + print("red start step: ", self.agent_settings.start_settings.start_step) self._set_next_execution_timestep(self.agent_settings.start_settings.start_step) @@ -27,6 +28,7 @@ class DataManipulationAgent(AbstractScriptedAgent): -self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance ) self.next_execution_timestep = timestep + random_timestep_increment + print("next execution red step: ", self.next_execution_timestep) def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: """Randomly sample an action from the action space. @@ -41,8 +43,15 @@ class DataManipulationAgent(AbstractScriptedAgent): current_timestep = self.action_manager.game.step_counter if current_timestep < self.next_execution_timestep: + print("red agent doing nothing") return "DONOTHING", {"dummy": 0} self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency) + print("red agent doing an execute") return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0} + + def reset_agent_for_episode(self) -> None: + """Set the next execution timestep when the episode resets.""" + super().reset_agent_for_episode() + self._set_next_execution_timestep(self.agent_settings.start_settings.start_step) diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 8657fc45..8b6dd6d4 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -135,6 +135,10 @@ class AbstractAgent(ABC): request = self.action_manager.form_request(action_identifier=action, action_options=options) return request + def reset_agent_for_episode(self) -> None: + """Agent reset logic should go here.""" + pass + class AbstractScriptedAgent(AbstractAgent): """Base class for actors which generate their own behaviour.""" diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 586bca79..08098754 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -162,6 +162,7 @@ class PrimaiteGame: self.simulation.reset_component_for_episode(episode=self.episode_counter) for agent in self.agents: agent.reward_function.total_reward = 0.0 + agent.reset_agent_for_episode() def close(self) -> None: """Close the game, this will close the simulation.""" From e73783f6fa2dc03c5c3274f923ae42460690d561 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 9 Jan 2024 17:10:12 +0000 Subject: [PATCH 2/4] Fixed issue where data manipulation was always executing --- .../services/red_services/data_manipulation_bot.py | 10 ++++++++-- src/primaite/simulator/system/software.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py index 44a56cf1..fcd9a3cc 100644 --- a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py +++ b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py @@ -73,7 +73,7 @@ class DataManipulationBot(DatabaseClient): def _init_request_manager(self) -> RequestManager: rm = super()._init_request_manager() - rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.run())) + rm.add_request(name="execute", request_type=RequestType(func=lambda request, context: self.attack())) return rm @@ -169,6 +169,12 @@ class DataManipulationBot(DatabaseClient): Calls the parent classes execute method before starting the application loop. """ super().run() + + def attack(self): + """Perform the attack steps after opening the application.""" + if not self._can_perform_action(): + _LOGGER.debug("Data manipulation application attempted to execute but it cannot perform actions right now.") + self.run() self._application_loop() def _application_loop(self): @@ -199,4 +205,4 @@ class DataManipulationBot(DatabaseClient): :param timestep: The timestep value to update the bot's state. """ - self._application_loop() + pass diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 048e6fec..ca667f46 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -278,7 +278,7 @@ class IOSoftware(Software): Returns true if the software can perform actions. """ - if self.software_manager and self.software_manager.node.operating_state is NodeOperatingState.OFF: + if self.software_manager and self.software_manager.node.operating_state is not NodeOperatingState.ON: _LOGGER.debug(f"{self.name} Error: {self.software_manager.node.hostname} is not online.") return False return True From b7cc940e9dd1485b4aada9d0f208bb17b3c4bee4 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 10 Jan 2024 09:07:51 +0000 Subject: [PATCH 3/4] Remove temporary print statements --- src/primaite/game/agent/data_manipulation_bot.py | 5 ----- src/primaite/game/agent/rewards.py | 1 - 2 files changed, 6 deletions(-) diff --git a/src/primaite/game/agent/data_manipulation_bot.py b/src/primaite/game/agent/data_manipulation_bot.py index 3b558087..7ad45518 100644 --- a/src/primaite/game/agent/data_manipulation_bot.py +++ b/src/primaite/game/agent/data_manipulation_bot.py @@ -15,8 +15,6 @@ class DataManipulationAgent(AbstractScriptedAgent): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - print("red start step: ", self.agent_settings.start_settings.start_step) - self._set_next_execution_timestep(self.agent_settings.start_settings.start_step) def _set_next_execution_timestep(self, timestep: int) -> None: @@ -28,7 +26,6 @@ class DataManipulationAgent(AbstractScriptedAgent): -self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance ) self.next_execution_timestep = timestep + random_timestep_increment - print("next execution red step: ", self.next_execution_timestep) def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]: """Randomly sample an action from the action space. @@ -43,12 +40,10 @@ class DataManipulationAgent(AbstractScriptedAgent): current_timestep = self.action_manager.game.step_counter if current_timestep < self.next_execution_timestep: - print("red agent doing nothing") return "DONOTHING", {"dummy": 0} self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency) - print("red agent doing an execute") return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0} def reset_agent_for_episode(self) -> None: diff --git a/src/primaite/game/agent/rewards.py b/src/primaite/game/agent/rewards.py index 9b3dfb80..cb8f8cb1 100644 --- a/src/primaite/game/agent/rewards.py +++ b/src/primaite/game/agent/rewards.py @@ -181,7 +181,6 @@ class WebServer404Penalty(AbstractReward): """ web_service_state = access_from_nested_dict(state, self.location_in_state) if web_service_state is NOT_PRESENT_IN_STATE: - print("error getting web service state") return 0.0 most_recent_return_code = web_service_state["last_response_status_code"] # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan. From b6e414bd705615ea2bad75f1376f6c825d4e7004 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Wed, 10 Jan 2024 09:14:07 +0000 Subject: [PATCH 4/4] Update changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c712ef66..7a0ef4c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] - Made packet capture and system logging optional (off by default). To turn on, change the io_settings.save_pcap_logs and io_settings.save_sys_logs settings in the config. - Made observation space flattening optional (on by default). To turn off for an agent, change the agent_settings.flatten_obs setting in the config. - +- Fixed an issue where the data manipulation attack was triggered at episode start. ### Added - Network Hardware - Added base hardware module with NIC, SwitchPort, Node, and Link. Nodes have