#2869 - Commit before switching branches. Changes to make pydantic happy with AgentLog

src/primaite/game/agent/scripted_agents/abstract_tap.py

@@ -15,6 +15,7 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"):
     class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
         """Configuration schema for Abstract TAP agents."""
 
+        agent_name: str = "Abstract_TAP"
         starting_node_name: str
         next_execution_timestep: int
 
@@ -32,3 +33,11 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"):
             -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance
         )
         self.config.next_execution_timestep = timestep + random_timestep_increment
+
+    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."""
+        # 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)
+        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}")

src/primaite/game/agent/scripted_agents/data_manipulation_bot.py

@@ -16,6 +16,11 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="Data_Manipulation_Agen
         """Configuration Schema for DataManipulationAgent."""
 
         starting_application_name: str
+        agent_name: str = "Data_Manipulation_Agent"
+
+    def __init__(self) -> None:
+        """Meh."""
+        self.setup_agent()
 
     @property
     def next_execution_timestep(self) -> int:
@@ -52,10 +57,3 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="Data_Manipulation_Agen
         """Set the next execution timestep when the episode resets."""
         self._select_start_node()
         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."""
-        # we are assuming that every node in the node manager has a data manipulation application at idx 0
-        num_nodes = len(self.action_manager.node_names)
-        self.starting_node_idx = random.randint(0, num_nodes - 1)
-        self.config.logger.debug(msg=f"Select Start Node ID: {self.starting_node_idx}")

src/primaite/game/agent/scripted_agents/probabilistic_agent.py

@@ -18,6 +18,7 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"
     class ConfigSchema(pydantic.BaseModel):
         """Config schema for Probabilistic agent settings."""
 
+        agent_name: str = "Probabilistic_Agent"
         model_config = pydantic.ConfigDict(extra="forbid")
         """Strict validation."""
         action_probabilities: Dict[int, float]

src/primaite/game/agent/scripted_agents/random_agent.py

@@ -10,7 +10,12 @@ from primaite.game.agent.interface import AbstractScriptedAgent
 class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"):
     """Agent that ignores its observation and acts completely at random."""
 
-    def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
+    class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
+        """Configuration Schema for Random Agents."""
+
+        agent_name = "Random_Agent"
+
+    def get_action(self) -> Tuple[str, Dict]:
         """Sample the action space randomly.
 
         :param obs: Current observation for this agent, not used in RandomAgent
@@ -31,6 +36,8 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"):
     class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
         """Configuration Schema for Periodic Agent."""
 
+        agent_name = "Periodic_Agent"
+        """Name of the agent."""
         start_step: int = 20
         "The timestep at which an agent begins performing it's actions."
         start_variance: int = 5
@@ -69,9 +76,9 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"):
 
     def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
         """Do nothing, unless the current timestep is the next execution timestep, in which case do the action."""
-        if timestep == self.next_execution_timestep and self.num_executions < self.settings.max_executions:
+        if timestep == self.next_execution_timestep and self.num_executions < self.config.max_executions:
             self.num_executions += 1
-            self._set_next_execution_timestep(timestep + self.settings.frequency, self.settings.variance)
+            self._set_next_execution_timestep(timestep + self.config.frequency, self.config.variance)
             return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
 
         return "DONOTHING", {}

src/primaite/game/agent/scripted_agents/tap001.py

@@ -19,20 +19,19 @@ class TAP001(AbstractTAPAgent, identifier="TAP001"):
     class ConfigSchema(AbstractTAPAgent.ConfigSchema):
         """Configuration Schema for TAP001 Agent."""
 
+        agent_name: str = "TAP001"
         installed: bool = False
 
+    def __init__(self) -> None:
+        """___init___ bruv. Restecpa."""
+        super().__init__()
+        self.setup_agent()
+
     @property
     def starting_node_name(self) -> str:
         """Node that TAP001 starts from."""
         return self.config.starting_node_name
 
-    @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
-
     def get_action(self, timestep: int) -> Tuple[str, Dict]:
         """Waits until a specific timestep, then attempts to execute the ransomware application.
 
@@ -72,11 +71,3 @@ class TAP001(AbstractTAPAgent, identifier="TAP001"):
                 self.ip_address = act[1]["ip_address"]
                 return
         raise RuntimeError("TAP001 agent could not find database server ip address in action map")
-
-    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."""
-        # 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)
-        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}")