oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2869 - Refactor agent and action config system

Browse Source
This commit is contained in:
Marek Wolan
2025-01-16 15:17:42 +00:00
parent f8fb052dad
commit 504f4bd134
12 changed files with 158 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/primaite/game/agent/scripted_agents/abstract_tap.py
View File

@@ -3,27 +3,36 @@ from __future__ import annotations
import random
from abc import abstractmethod
from typing import Dict, Optional, Tuple
from typing import Dict, List, Optional, Tuple
from gymnasium.core import ObsType
from pydantic import Field
from primaite.game.agent.interface import AbstractScriptedAgent
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "AbstractTAPAgent"
class AbstractTAPAgent(AbstractScriptedAgent, identifier="AbstractTAP"):
class AbstractTAPAgent(PeriodicAgent, identifier="AbstractTAP"):
"""Base class for TAP agents to inherit from."""
config: "AbstractTAPAgent.ConfigSchema" = Field(default_factory=lambda: AbstractTAPAgent.ConfigSchema())
next_execution_timestep: int = 0
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
possible_starting_nodes: List[str] = Field(default_factory=list)
class ConfigSchema(PeriodicAgent.ConfigSchema):
"""Configuration schema for Abstract TAP agents."""
type: str = "AbstractTAP"
starting_node_name: Optional[str] = None
agent_settings: AbstractTAPAgent.AgentSettingsSchema = Field(
default_factory=lambda: AbstractTAPAgent.AgentSettingsSchema()
)
starting_node: Optional[str] = None
@abstractmethod
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
@@ -40,13 +49,13 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="AbstractTAP"):
:param timestep: The timestep to add variance to.
"""
random_timestep_increment = random.randint(-self.config.variance, self.config.variance)
random_timestep_increment = random.randint(
-self.config.agent_settings.variance, self.config.agent_settings.variance
)
self.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.action_manager.node_names)
starting_node_idx = random.randint(0, num_nodes - 1)
self.config.starting_node_name = self.action_manager.node_names[starting_node_idx]
self.logger.debug(f"Selected starting node: {self.config.starting_node_name}")
self.starting_node = random.choice(self.config.agent_settings.possible_starting_nodes)
self.logger.debug(f"Selected starting node: {self.starting_node}")

24
src/primaite/game/agent/scripted_agents/data_manipulation_bot.py
View File

@@ -1,6 +1,5 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import random
from typing import Dict, List, Tuple
from typing import Dict, Tuple
from gymnasium.core import ObsType
from pydantic import Field
@@ -13,21 +12,24 @@ __all__ = "DataManipulationAgent"
class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
class AgentSettingsSchema(PeriodicAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
target_application: str = "DataManipulationBot"
class ConfigSchema(PeriodicAgent.ConfigSchema):
"""Configuration Schema for DataManipulationAgent."""
type: str = "RedDatabaseCorruptingAgent"
starting_application_name: str = "DataManipulationBot"
possible_start_nodes: List[str]
agent_settings: "DataManipulationAgent.AgentSettingsSchema" = Field(
default_factory=lambda: DataManipulationAgent.AgentSettingsSchema()
)
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)
self._set_next_execution_timestep(timestep=self.config.agent_settings.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.
@@ -43,9 +45,11 @@ class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgen
self.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
self._set_next_execution_timestep(timestep=timestep + self.config.frequency, variance=self.config.variance)
self._set_next_execution_timestep(
timestep=timestep + self.config.agent_settings.frequency, variance=self.config.agent_settings.variance
)
self.logger.info(msg="Performing a data manipulation attack!")
return "node_application_execute", {
"node_name": self.start_node,
"application_name": self.config.starting_application_name,
"application_name": self.config.agent_settings.target_application,
}

16
src/primaite/game/agent/scripted_agents/probabilistic_agent.py
View File

@@ -19,10 +19,8 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
config: "ProbabilisticAgent.ConfigSchema" = Field(default_factory=lambda: ProbabilisticAgent.ConfigSchema())
rng: Generator = np.random.default_rng(np.random.randint(0, 65535))
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration schema for Probabilistic Agent."""
type: str = "ProbabilisticAgent"
class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
action_probabilities: Dict[int, float] = None
"""Probability to perform each action in the action map. The sum of probabilities should sum to 1."""
@@ -46,10 +44,18 @@ class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent")
)
return v
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration schema for Probabilistic Agent."""
type: str = "ProbabilisticAgent"
agent_settings: "ProbabilisticAgent.AgentSettingsSchema" = Field(
default_factory=lambda: ProbabilisticAgent.AgentSettingsSchema()
)
@property
def probabilities(self) -> Dict[str, int]:
"""Convenience method to view the probabilities of the Agent."""
return np.asarray(list(self.config.action_probabilities.values()))
return np.asarray(list(self.config.agent_settings.action_probabilities.values()))
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
"""

38
src/primaite/game/agent/scripted_agents/random_agent.py
View File

@@ -1,9 +1,10 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
import random
from typing import Dict, Tuple
from functools import cached_property
from typing import Dict, List, Tuple
from gymnasium.core import ObsType
from pydantic import Field, model_validator
from pydantic import computed_field, Field, model_validator
from primaite.game.agent.interface import AbstractScriptedAgent
@@ -38,17 +39,17 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
config: "PeriodicAgent.ConfigSchema" = Field(default_factory=lambda: PeriodicAgent.ConfigSchema())
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Periodic Agent."""
class AgentSettingsSchema(AbstractScriptedAgent.AgentSettingsSchema):
"""Schema for the `agent_settings` part of the agent config."""
type: str = "PeriodicAgent"
"""Name of the agent."""
start_step: int = 5
"The timestep at which an agent begins performing it's actions"
frequency: int = 5
"The number of timesteps to wait between performing actions"
variance: int = 0
"The amount the frequency can randomly change to"
possible_start_nodes: List[str]
target_application: str
@model_validator(mode="after")
def check_variance_lt_frequency(self) -> "PeriodicAgent.ConfigSchema":
@@ -66,6 +67,15 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
)
return self
class ConfigSchema(AbstractScriptedAgent.ConfigSchema):
"""Configuration Schema for Periodic Agent."""
type: str = "PeriodicAgent"
"""Name of the agent."""
agent_settings: "PeriodicAgent.AgentSettingsSchema" = Field(
default_factory=lambda: PeriodicAgent.AgentSettingsSchema()
)
max_executions: int = 999999
"Maximum number of times the agent can execute its action."
num_executions: int = 0
@@ -73,6 +83,12 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
next_execution_timestep: int = 0
"""Timestep of the next action execution by the agent."""
@computed_field
@cached_property
def start_node(self) -> str:
"""On instantiation, randomly select a start node."""
return random.choice(self.config.agent_settings.possible_start_nodes)
def _set_next_execution_timestep(self, timestep: int, variance: int) -> None:
"""Set the next execution timestep with a configured random variance.
@@ -88,8 +104,12 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="PeriodicAgent"):
"""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.max_executions:
self.num_executions += 1
self._set_next_execution_timestep(timestep + self.config.frequency, self.config.variance)
self.target_node = self.action_manager.node_names[0]
return "node_application_execute", {"node_name": self.target_node, "application_name": 0}
self._set_next_execution_timestep(
timestep + self.config.agent_settings.frequency, self.config.agent_settings.variance
)
return "node_application_execute", {
"node_name": self.start_node,
"application_name": self.config.agent_settings.target_application,
}
return "do_nothing", {}