From e7cfeeafc04c3862e6baac46c59030090e7079fb Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 14 Jan 2025 09:57:35 +0000 Subject: [PATCH] Make data manipulation agent inherit from periodic agent & fix it a bit --- .../_package_data/data_manipulation.yaml | 4 +- .../_package_data/data_manipulation_marl.yaml | 4 +- .../scripted_agents/data_manipulation_bot.py | 39 ++++++++----------- 3 files changed, 23 insertions(+), 24 deletions(-) diff --git a/src/primaite/config/_package_data/data_manipulation.yaml b/src/primaite/config/_package_data/data_manipulation.yaml index d604192e..58986bce 100644 --- a/src/primaite/config/_package_data/data_manipulation.yaml +++ b/src/primaite/config/_package_data/data_manipulation.yaml @@ -150,7 +150,9 @@ agents: reward_components: - type: DUMMY - agent_settings: # options specific to this particular agent type, basically args of __init__(self) + agent_settings: + possible_start_nodes: [client_1, client_2] + starting_application_name: DataManipulationBot start_step: 25 frequency: 20 variance: 5 diff --git a/src/primaite/config/_package_data/data_manipulation_marl.yaml b/src/primaite/config/_package_data/data_manipulation_marl.yaml index 00a34403..c4a3b562 100644 --- a/src/primaite/config/_package_data/data_manipulation_marl.yaml +++ b/src/primaite/config/_package_data/data_manipulation_marl.yaml @@ -149,7 +149,9 @@ agents: reward_components: - type: DUMMY - agent_settings: # options specific to this particular agent type, basically args of __init__(self) + agent_settings: + possible_start_nodes: [client_1, client_2] + starting_application_name: DataManipulationBot start_step: 25 frequency: 20 variance: 5 diff --git a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py index b9d57a8b..8fe0690b 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py @@ -1,29 +1,33 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK -from typing import Dict, Optional, Tuple +import random +from typing import Dict, List, Tuple from gymnasium.core import ObsType from pydantic import Field -from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent +from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent __all__ = "DataManipulationAgent" -class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingAgent"): +class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"): """Agent that uses a DataManipulationBot to perform an SQL injection attack.""" - config: "DataManipulationAgent.ConfigSchema" = Field(default_factory=lambda: DataManipulationAgent.ConfigSchema()) - - class ConfigSchema(AbstractTAPAgent.ConfigSchema): + class ConfigSchema(PeriodicAgent.ConfigSchema): """Configuration Schema for DataManipulationAgent.""" type: str = "RedDatabaseCorruptingAgent" - starting_application_name: Optional[str] = None + starting_application_name: str = "DataManipulationBot" + possible_start_nodes: List[str] - @property - def starting_node_name(self) -> str: - """Returns the agents starting node name.""" - return self.config.starting_node_name + config: "DataManipulationAgent.ConfigSchema" = Field(default_factory=lambda: DataManipulationAgent.ConfigSchema()) + + start_node: str + + def __init__(self, **kwargs): + kwargs["start_node"] = random.choice(kwargs["config"].possible_start_nodes) + super().__init__(**kwargs) + self._set_next_execution_timestep(timestep=self.config.start_step, variance=0) def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: """Waits until a specific timestep, then attempts to execute its data manipulation application. @@ -35,22 +39,13 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA :return: Action formatted in CAOS format :rtype: Tuple[str, Dict] """ - if self.starting_node_name or self.config is None: - self.setup_agent() - self.get_action(obs=obs, timestep=timestep) - if timestep < self.next_execution_timestep: self.logger.debug(msg="Performing do nothing action") return "do_nothing", {} - self._set_next_execution_timestep(timestep + self.config.frequency) + self._set_next_execution_timestep(timestep=timestep + self.config.frequency, variance=self.config.variance) self.logger.info(msg="Performing a data manipulation attack!") return "node_application_execute", { - "node_name": self.config.starting_node_name, + "node_name": self.start_node, "application_name": self.config.starting_application_name, } - - def setup_agent(self) -> None: - """Set the next execution timestep when the episode resets.""" - self._select_start_node() - self._set_next_execution_timestep(self.config.start_step)