diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 962e13f7..402c7ce2 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -222,6 +222,8 @@ class AbstractScriptedAgent(AbstractAgent, identifier="Abstract_Scripted_Agent") class ProxyAgent(AbstractAgent, identifier="Proxy_Agent"): """Agent that sends observations to an RL model and receives actions from that model.""" + config: "ProxyAgent.ConfigSchema" + class ConfigSchema(AbstractAgent.ConfigSchema): """Configuration Schema for Proxy Agent.""" @@ -230,23 +232,10 @@ class ProxyAgent(AbstractAgent, identifier="Proxy_Agent"): flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False action_masking: bool = agent_settings.action_masking if agent_settings else False - # def __init__( - # self, - # agent_name: Optional[str], - # action_space: Optional[ActionManager], - # observation_space: Optional[ObservationManager], - # reward_function: Optional[RewardFunction], - # agent_settings: Optional[AgentSettings] = None, - # ) -> None: - # super().__init__( - # agent_name=agent_name, - # action_space=action_space, - # observation_space=observation_space, - # reward_function=reward_function, - # ) - # self.most_recent_action: ActType - # self.flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False - # self.action_masking: bool = agent_settings.action_masking if agent_settings else False + @property + def most_recent_action(self) -> ActType: + """Convenience method to access the agents most recent action.""" + return self.config.most_recent_action def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: """ @@ -267,4 +256,4 @@ class ProxyAgent(AbstractAgent, identifier="Proxy_Agent"): The environment is responsible for calling this method when it receives an action from the agent policy. """ - self.config.most_recent_action = action + self.most_recent_action = action diff --git a/src/primaite/game/agent/scripted_agents/abstract_tap.py b/src/primaite/game/agent/scripted_agents/abstract_tap.py new file mode 100644 index 00000000..2523f9f7 --- /dev/null +++ b/src/primaite/game/agent/scripted_agents/abstract_tap.py @@ -0,0 +1,34 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +from __future__ import annotations + +import random +from abc import abstractmethod + +from primaite.game.agent.interface import AbstractScriptedAgent + + +class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"): + """Base class for TAP agents to inherit from.""" + + config: "AbstractTAPAgent.ConfigSchema" + + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """Configuration schema for Abstract TAP agents.""" + + starting_node_name: str + next_execution_timestep: int + + @abstractmethod + def setup_agent(self) -> None: + """Set up agent.""" + pass + + def _set_next_execution_timestep(self, timestep: int) -> None: + """Set the next execution timestep with a configured random variance. + + :param timestep: The timestep to add variance to. + """ + random_timestep_increment = random.randint( + -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance + ) + self.config.next_execution_timestep = timestep + random_timestep_increment 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 2f49decd..b375da66 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py @@ -4,46 +4,29 @@ from typing import Dict, Tuple from gymnasium.core import ObsType -from primaite.game.agent.interface import AbstractScriptedAgent +from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent -class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation_Agent"): +class DataManipulationAgent(AbstractTAPAgent, identifier="Data_Manipulation_Agent"): """Agent that uses a DataManipulationBot to perform an SQL injection attack.""" config: "DataManipulationAgent.ConfigSchema" - class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + class ConfigSchema(AbstractTAPAgent.ConfigSchema): """Configuration Schema for DataManipulationAgent.""" - # TODO: Could be worth moving this to a "AbstractTAPAgent" - starting_node_name: str starting_application_name: str - next_execution_timestep: int - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.setup_agent() @property - def next_execution_timestep(self): + def next_execution_timestep(self) -> int: """Returns the agents next execution timestep.""" return self.config.next_execution_timestep @property - def starting_node_name(self): + def starting_node_name(self) -> str: """Returns the agents starting node name.""" return self.config.starting_node_name - def _set_next_execution_timestep(self, timestep: int) -> None: - """Set the next execution timestep with a configured random variance. - - :param timestep: The timestep to add variance to. - """ - random_timestep_increment = random.randint( - -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance - ) - self.next_execution_timestep = timestep + random_timestep_increment - def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: """Waits until a specific timestep, then attempts to execute its data manipulation application. @@ -68,7 +51,7 @@ class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation def setup_agent(self) -> None: """Set the next execution timestep when the episode resets.""" self._select_start_node() - self._set_next_execution_timestep(self.agent_settings.start_settings.start_step) + self._set_next_execution_timestep(self.config.agent_settings.start_settings.start_step) def _select_start_node(self) -> None: """Set the starting starting node of the agent to be a random node from this agent's action manager.""" diff --git a/src/primaite/game/agent/scripted_agents/random_agent.py b/src/primaite/game/agent/scripted_agents/random_agent.py index 99b8a1e9..a9082eda 100644 --- a/src/primaite/game/agent/scripted_agents/random_agent.py +++ b/src/primaite/game/agent/scripted_agents/random_agent.py @@ -1,14 +1,10 @@ # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK import random -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple from gymnasium.core import ObsType -from pydantic import BaseModel -from primaite.game.agent.actions import ActionManager from primaite.game.agent.interface import AbstractScriptedAgent -from primaite.game.agent.observations.observation_manager import ObservationManager -from primaite.game.agent.rewards import RewardFunction class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"): @@ -30,8 +26,10 @@ class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"): class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"): """Agent that does nothing most of the time, but executes application at regular intervals (with variance).""" - class Settings(BaseModel): - """Configuration values for when an agent starts performing actions.""" + config: "PeriodicAgent.ConfigSchema" + + class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + """Configuration Schema for Periodic Agent.""" start_step: int = 20 "The timestep at which an agent begins performing it's actions." @@ -43,25 +41,20 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"): "The amount the frequency can randomly change to." max_executions: int = 999999 "Maximum number of times the agent can execute its action." + num_executions: int = 0 + """Number of times the agent has executed an action.""" + next_execution_timestep: int = 0 + """Timestep of the next action execution by the agent.""" - def __init__( - self, - agent_name: str, - action_space: ActionManager, - observation_space: ObservationManager, - reward_function: RewardFunction, - settings: Optional[Settings] = None, - ) -> None: - """Initialise PeriodicAgent.""" - super().__init__( - agent_name=agent_name, - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - ) - self.settings = settings or PeriodicAgent.Settings() - self._set_next_execution_timestep(timestep=self.settings.start_step, variance=self.settings.start_variance) - self.num_executions = 0 + @property + def num_executions(self) -> int: + """Convenience method for accessing num_executions from config.""" + return self.config.num_executions + + @property + def next_execution_timestep(self) -> int: + """Convenience method for accessing next_execution_timestep from config.""" + return self.config.next_execution_timestep def _set_next_execution_timestep(self, timestep: int, variance: int) -> None: """Set the next execution timestep with a configured random variance. diff --git a/src/primaite/game/agent/scripted_agents/tap001.py b/src/primaite/game/agent/scripted_agents/tap001.py index 7365fd88..d3a82bbe 100644 --- a/src/primaite/game/agent/scripted_agents/tap001.py +++ b/src/primaite/game/agent/scripted_agents/tap001.py @@ -4,51 +4,41 @@ from __future__ import annotations import random from typing import Dict, Tuple -from gymnasium.core import ObsType - -from primaite.game.agent.interface import AbstractScriptedAgent +from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent -class TAP001(AbstractScriptedAgent, identifier="TAP001"): +class TAP001(AbstractTAPAgent, identifier="TAP001"): """ TAP001 | Mobile Malware -- Ransomware Variant. Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application) """ - # TODO: Link with DataManipulationAgent class via a parent "TAP" agent class. - config: "TAP001.ConfigSchema" - class ConfigSchema(AbstractScriptedAgent.ConfigSchema): + class ConfigSchema(AbstractTAPAgent.ConfigSchema): """Configuration Schema for TAP001 Agent.""" - starting_node_name: str - next_execution_timestep: int = 0 installed: bool = False - # def __init__(self, *args, **kwargs): - # super().__init__(*args, **kwargs) - # self.setup_agent() + @property + def starting_node_name(self) -> str: + """Node that TAP001 starts from.""" + return self.config.starting_node_name - def _set_next_execution_timestep(self, timestep: int) -> None: - """Set the next execution timestep with a configured random variance. + @classmethod + def from_config(cls, config: Dict) -> TAP001: + """Override the base from_config method to ensure successful agent setup.""" + obj: TAP001 = cls(config=cls.ConfigSchema(**config)) + obj.setup_agent() + return obj - :param timestep: The timestep to add variance to. - """ - random_timestep_increment = random.randint( - -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance - ) - self.config.next_execution_timestep = timestep + random_timestep_increment - - def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: + def get_action(self, timestep: int) -> Tuple[str, Dict]: """Waits until a specific timestep, then attempts to execute the ransomware application. This application acts a wrapper around the kill-chain, similar to green-analyst and the previous UC2 data manipulation bot. - :param obs: Current observation for this agent. - :type obs: ObsType :param timestep: The current simulation timestep, used for scheduling actions :type timestep: int :return: Action formatted in CAOS format @@ -62,11 +52,14 @@ class TAP001(AbstractScriptedAgent, identifier="TAP001"): if not self.config.installed: self.config.installed = True return "node_application_install", { - "node_name": self.config.starting_node_name, + "node_name": self.starting_node_name, "application_name": "RansomwareScript", } - return "node_application_execute", {"node_name": self.config.starting_node_name, "application_id": 0} + return "node_application_execute", { + "node_name": self.starting_node_name, + "application_name": "RansomwareScript", + } def setup_agent(self) -> None: """Set the next execution timestep when the episode resets.""" @@ -75,7 +68,7 @@ class TAP001(AbstractScriptedAgent, identifier="TAP001"): for n, act in self.config.action_manager.action_map.items(): if not act[0] == "node_application_install": continue - if act[1]["node_name"] == self.config.starting_node_name: + if act[1]["node_name"] == self.starting_node_name: self.ip_address = act[1]["ip_address"] return raise RuntimeError("TAP001 agent could not find database server ip address in action map") @@ -84,6 +77,6 @@ class TAP001(AbstractScriptedAgent, identifier="TAP001"): """Set the starting starting node of the agent to be a random node from this agent's action manager.""" # we are assuming that every node in the node manager has a data manipulation application at idx 0 num_nodes = len(self.config.action_manager.node_names) - # TODO: Change this to something? - self.starting_node_idx = random.randint(0, num_nodes - 1) - self.logger.debug(f"Selected Starting node ID: {self.starting_node_idx}") + starting_node_idx = random.randint(0, num_nodes - 1) + self.starting_node_name = self.config.action_manager.node_names[starting_node_idx] + self.config.logger.debug(f"Selected Starting node ID: {self.starting_node_name}") diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 03f3feec..79587e47 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -8,7 +8,7 @@ from pydantic import BaseModel, ConfigDict from primaite import DEFAULT_BANDWIDTH, getLogger from primaite.game.agent.actions import ActionManager -from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent +from primaite.game.agent.interface import AbstractAgent, ProxyAgent from primaite.game.agent.observations.observation_manager import ObservationManager from primaite.game.agent.rewards import RewardFunction, SharedReward from primaite.game.science import graph_has_cycle, topological_sort diff --git a/tests/conftest.py b/tests/conftest.py index efdb515e..b24c4c76 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -267,20 +267,12 @@ def example_network() -> Network: class ControlledAgent(AbstractAgent, identifier="Controlled_Agent"): """Agent that can be controlled by the tests.""" - def __init__( - self, - agent_name: str, - action_space: ActionManager, - observation_space: ObservationManager, - reward_function: RewardFunction, - ) -> None: - super().__init__( - agent_name=agent_name, - action_space=action_space, - observation_space=observation_space, - reward_function=reward_function, - ) - self.most_recent_action: Tuple[str, Dict] + config: "ControlledAgent.ConfigSchema" + + class ConfigSchema(AbstractAgent.ConfigSchema): + """Configuration Schema for Abstract Agent used in tests.""" + + most_recent_action: Tuple[str, Dict] def get_action(self, obs: None, timestep: int = 0) -> Tuple[str, Dict]: """Return the agent's most recent action, formatted in CAOS format."""