#2869 - Agents Refactor

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__(
+    @property
-    #     self,
+    def most_recent_action(self) -> ActType:
-    #     agent_name: Optional[str],
+        """Convenience method to access the agents most recent action."""
-    #     action_space: Optional[ActionManager],
+        return self.config.most_recent_action
-    #     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
 
     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

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

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."""

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):
+    config: "PeriodicAgent.ConfigSchema"
-        """Configuration values for when an agent starts performing actions."""
+
+    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__(
+    @property
-        self,
+    def num_executions(self) -> int:
-        agent_name: str,
+        """Convenience method for accessing num_executions from config."""
-        action_space: ActionManager,
+        return self.config.num_executions
-        observation_space: ObservationManager,
+
-        reward_function: RewardFunction,
+    @property
-        settings: Optional[Settings] = None,
+    def next_execution_timestep(self) -> int:
-    ) -> None:
+        """Convenience method for accessing next_execution_timestep from config."""
-        """Initialise PeriodicAgent."""
+        return self.config.next_execution_timestep
-        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
 
     def _set_next_execution_timestep(self, timestep: int, variance: int) -> None:
         """Set the next execution timestep with a configured random variance.

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.scripted_agents.abstract_tap import AbstractTAPAgent
-
-from primaite.game.agent.interface import AbstractScriptedAgent
 
 
-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):
+    @property
-    #     super().__init__(*args, **kwargs)
+    def starting_node_name(self) -> str:
-    #     self.setup_agent()
+        """Node that TAP001 starts from."""
+        return self.config.starting_node_name
 
-    def _set_next_execution_timestep(self, timestep: int) -> None:
+    @classmethod
-        """Set the next execution timestep with a configured random variance.
+    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.
+    def get_action(self, timestep: int) -> Tuple[str, Dict]:
-        """
-        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]:
         """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?
+        starting_node_idx = random.randint(0, num_nodes - 1)
-        self.starting_node_idx = random.randint(0, num_nodes - 1)
+        self.starting_node_name = self.config.action_manager.node_names[starting_node_idx]
-        self.logger.debug(f"Selected Starting node ID: {self.starting_node_idx}")
+        self.config.logger.debug(f"Selected Starting node ID: {self.starting_node_name}")

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

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__(
+    config: "ControlledAgent.ConfigSchema"
-        self,
+
-        agent_name: str,
+    class ConfigSchema(AbstractAgent.ConfigSchema):
-        action_space: ActionManager,
+        """Configuration Schema for Abstract Agent used in tests."""
-        observation_space: ObservationManager,
+
-        reward_function: RewardFunction,
+        most_recent_action: Tuple[str, Dict]
-    ) -> 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]
 
     def get_action(self, obs: None, timestep: int = 0) -> Tuple[str, Dict]:
         """Return the agent's most recent action, formatted in CAOS format."""