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

#2869 - Starter changes in refactor of agent classes for refactor to become extensible. Identifiers added to classes and beginning of the inclusion of a ConfigSchema to base AbstractAgentClass

Browse Source
This commit is contained in:
Charlie Crane
2024-11-20 17:19:35 +00:00
parent 6844bd692a
commit a3dc616126
7 changed files with 174 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
src/primaite/game/agent/interface.py
View File

@@ -1,7 +1,9 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
"""Interface for agents.""" """Interface for agents."""
from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union
from gymnasium.core import ActType, ObsType from gymnasium.core import ActType, ObsType
from pydantic import BaseModel, model_validator from pydantic import BaseModel, model_validator
@@ -69,7 +71,7 @@ class AgentSettings(BaseModel):
"""Settings for configuring the operation of an agent.""" """Settings for configuring the operation of an agent."""
start_settings: Optional[AgentStartSettings] = None start_settings: Optional[AgentStartSettings] = None
"Configuration for when an agent begins performing it's actions" "Configuration for when an agent begins performing it's actions."
flatten_obs: bool = True flatten_obs: bool = True
"Whether to flatten the observation space before passing it to the agent. True by default." "Whether to flatten the observation space before passing it to the agent. True by default."
action_masking: bool = False action_masking: bool = False
@@ -90,38 +92,78 @@ class AgentSettings(BaseModel):
return cls(**config) return cls(**config)
class AbstractAgent(ABC): class AbstractAgent(BaseModel, ABC, identifier="Abstract_Agent"):
"""Base class for scripted and RL agents.""" """Base class for scripted and RL agents."""
def __init__( _registry: ClassVar[Dict[str, Type[AbstractAgent]]] = {}
self,
agent_name: Optional[str],
action_space: Optional[ActionManager],
observation_space: Optional[ObservationManager],
reward_function: Optional[RewardFunction],
agent_settings: Optional[AgentSettings] = None,
) -> None:
"""
Initialize an agent.
config: "AbstractAgent.ConfigSchema"
action_manager: Optional[ActionManager]
observation_manager: Optional[ObservationManager]
reward_function: Optional[RewardFunction]
class ConfigSchema(BaseModel):
"""
Configuration Schema for AbstractAgents.
:param type: Type of agent being generated.
:type type: str
:param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes. :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes.
:type agent_name: Optional[str] :type agent_name: str
:param action_space: Action space for the agent.
:type action_space: Optional[ActionManager]
:param observation_space: Observation space for the agent. :param observation_space: Observation space for the agent.
:type observation_space: Optional[ObservationSpace] :type observation_space: Optional[ObservationSpace]
:param reward_function: Reward function for the agent. :param reward_function: Reward function for the agent.
:type reward_function: Optional[RewardFunction] :type reward_function: Optional[RewardFunction]
:param agent_settings: Configurable Options for Abstracted Agents :param agent_settings: Configurable Options for Abstracted Agents.
:type agent_settings: Optional[AgentSettings] :type agent_settings: Optional[AgentSettings]
""" """
self.agent_name: str = agent_name or "unnamed_agent"
self.action_manager: Optional[ActionManager] = action_space type: str
self.observation_manager: Optional[ObservationManager] = observation_space agent_name: ClassVar[str]
self.reward_function: Optional[RewardFunction] = reward_function agent_settings = Optional[AgentSettings] = None
self.agent_settings = agent_settings or AgentSettings() history: List[AgentHistoryItem] = []
self.history: List[AgentHistoryItem] = [] logger: AgentLog = AgentLog(agent_name)
self.logger = AgentLog(agent_name)
# 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:
# """
# Initialize an agent.
# :param agent_name: Unique string identifier for the agent, for reporting and multi-agent purposes.
# :type agent_name: Optional[str]
# :param action_space: Action space for the agent.
# :type action_space: Optional[ActionManager]
# :param observation_space: Observation space for the agent.
# :type observation_space: Optional[ObservationSpace]
# :param reward_function: Reward function for the agent.
# :type reward_function: Optional[RewardFunction]
# :param agent_settings: Configurable Options for Abstracted Agents
# :type agent_settings: Optional[AgentSettings]
# """
# self.agent_name: str = agent_name or "unnamed_agent"
# self.action_manager: Optional[ActionManager] = action_space
# self.observation_manager: Optional[ObservationManager] = observation_space
# self.reward_function: Optional[RewardFunction] = reward_function
# self.agent_settings = agent_settings or AgentSettings()
# self.history: List[AgentHistoryItem] = []
# self.logger = AgentLog(agent_name)
def __init_subclass__(cls, identifier: str, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
if identifier in cls._registry:
raise ValueError(f"Cannot create a new agent under reserved name {identifier}")
cls._registry[identifier] = cls
@classmethod
def from_config(cls, config: Dict) -> "AbstractAgent":
"""Creates an agent component from a configuration dictionary."""
return cls(config=cls.ConfigSchema(**config))
def update_observation(self, state: Dict) -> ObsType: def update_observation(self, state: Dict) -> ObsType:
""" """
@@ -130,7 +172,7 @@ class AbstractAgent(ABC):
state : dict state directly from simulation.describe_state state : dict state directly from simulation.describe_state
output : dict state according to CAOS. output : dict state according to CAOS.
""" """
return self.observation_manager.update(state) return self.config.observation_manager.update(state)
def update_reward(self, state: Dict) -> float: def update_reward(self, state: Dict) -> float:
""" """
@@ -141,7 +183,7 @@ class AbstractAgent(ABC):
:return: Reward from the state. :return: Reward from the state.
:rtype: float :rtype: float
""" """
return self.reward_function.update(state=state, last_action_response=self.history[-1]) return self.config.reward_function.update(state=state, last_action_response=self.history[-1])
@abstractmethod @abstractmethod
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
@@ -165,14 +207,14 @@ class AbstractAgent(ABC):
# this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator. # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator.
# therefore the execution definition needs to be a mapping from CAOS into SIMULATOR # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR
"""Format action into format expected by the simulator, and apply execution definition if applicable.""" """Format action into format expected by the simulator, and apply execution definition if applicable."""
request = self.action_manager.form_request(action_identifier=action, action_options=options) request = self.config.action_manager.form_request(action_identifier=action, action_options=options)
return request return request
def process_action_response( def process_action_response(
self, timestep: int, action: str, parameters: Dict[str, Any], request: RequestFormat, response: RequestResponse self, timestep: int, action: str, parameters: Dict[str, Any], request: RequestFormat, response: RequestResponse
) -> None: ) -> None:
"""Process the response from the most recent action.""" """Process the response from the most recent action."""
self.history.append( self.config.history.append(
AgentHistoryItem( AgentHistoryItem(
timestep=timestep, action=action, parameters=parameters, request=request, response=response timestep=timestep, action=action, parameters=parameters, request=request, response=response
) )
@@ -180,10 +222,10 @@ class AbstractAgent(ABC):
def save_reward_to_history(self) -> None: def save_reward_to_history(self) -> None:
"""Update the most recent history item with the reward value.""" """Update the most recent history item with the reward value."""
self.history[-1].reward = self.reward_function.current_reward self.config.history[-1].reward = self.config.reward_function.current_reward
class AbstractScriptedAgent(AbstractAgent): class AbstractScriptedAgent(AbstractAgent, identifier="Abstract_Scripted_Agent"):
"""Base class for actors which generate their own behaviour.""" """Base class for actors which generate their own behaviour."""
@abstractmethod @abstractmethod
@@ -192,26 +234,34 @@ class AbstractScriptedAgent(AbstractAgent):
return super().get_action(obs=obs, timestep=timestep) return super().get_action(obs=obs, timestep=timestep)
class ProxyAgent(AbstractAgent): class ProxyAgent(AbstractAgent, identifier="Proxy_Agent"):
"""Agent that sends observations to an RL model and receives actions from that model.""" """Agent that sends observations to an RL model and receives actions from that model."""
def __init__( class ConfigSchema(AbstractAgent.ConfigSchema):
self, """Configuration Schema for Proxy Agent."""
agent_name: Optional[str],
action_space: Optional[ActionManager], agent_settings = Union[AgentSettings | None] = None
observation_space: Optional[ObservationManager], most_reason_action: ActType
reward_function: Optional[RewardFunction], flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False
agent_settings: Optional[AgentSettings] = None, action_masking: bool = agent_settings.action_masking if agent_settings else False
) -> None:
super().__init__( # def __init__(
agent_name=agent_name, # self,
action_space=action_space, # agent_name: Optional[str],
observation_space=observation_space, # action_space: Optional[ActionManager],
reward_function=reward_function, # observation_space: Optional[ObservationManager],
) # reward_function: Optional[RewardFunction],
self.most_recent_action: ActType # agent_settings: Optional[AgentSettings] = None,
self.flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False # ) -> None:
self.action_masking: bool = agent_settings.action_masking if agent_settings else False # 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]: def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
""" """
@@ -224,7 +274,7 @@ class ProxyAgent(AbstractAgent):
:return: Action to be taken in CAOS format. :return: Action to be taken in CAOS format.
:rtype: Tuple[str, Dict] :rtype: Tuple[str, Dict]
""" """
return self.action_manager.get_action(self.most_recent_action) return self.config.action_manager.get_action(self.most_recent_action)
def store_action(self, action: ActType): def store_action(self, action: ActType):
""" """
@@ -232,4 +282,4 @@ class ProxyAgent(AbstractAgent):
The environment is responsible for calling this method when it receives an action from the agent policy. The environment is responsible for calling this method when it receives an action from the agent policy.
""" """
self.most_recent_action = action self.config.most_recent_action = action

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

@@ -7,12 +7,22 @@ from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent from primaite.game.agent.interface import AbstractScriptedAgent
class DataManipulationAgent(AbstractScriptedAgent): class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation_Agent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack.""" """Agent that uses a DataManipulationBot to perform an SQL injection attack."""
next_execution_timestep: int = 0 next_execution_timestep: int = 0
starting_node_idx: int = 0 starting_node_idx: int = 0
config: "DataManipulationAgent.ConfigSchema"
class ConfigSchema(AbstractScriptedAgent.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): def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.setup_agent() self.setup_agent()
@@ -38,12 +48,15 @@ class DataManipulationAgent(AbstractScriptedAgent):
:rtype: Tuple[str, Dict] :rtype: Tuple[str, Dict]
""" """
if timestep < self.next_execution_timestep: if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do NOTHING") self.logger.debug(msg="Performing do nothing action")
return "DONOTHING", {} return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency) self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
self.logger.info(msg="Performing a data manipulation attack!") self.logger.info(msg="Performing a data manipulation attack!")
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0} return "node_application_execute", {
"node_name": self.config.starting_node_name,
"application_name": self.config.starting_application_name,
}
def setup_agent(self) -> None: def setup_agent(self) -> None:
"""Set the next execution timestep when the episode resets.""" """Set the next execution timestep when the episode resets."""

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

@@ -12,7 +12,7 @@ from primaite.game.agent.observations.observation_manager import ObservationMana
from primaite.game.agent.rewards import RewardFunction from primaite.game.agent.rewards import RewardFunction
class ProbabilisticAgent(AbstractScriptedAgent): class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action.""" """Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
class Settings(pydantic.BaseModel): class Settings(pydantic.BaseModel):

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

@@ -11,7 +11,7 @@ from primaite.game.agent.observations.observation_manager import ObservationMana
from primaite.game.agent.rewards import RewardFunction from primaite.game.agent.rewards import RewardFunction
class RandomAgent(AbstractScriptedAgent): class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"):
"""Agent that ignores its observation and acts completely at random.""" """Agent that ignores its observation and acts completely at random."""
def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]: def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
@@ -27,7 +27,7 @@ class RandomAgent(AbstractScriptedAgent):
return self.action_manager.get_action(self.action_manager.space.sample()) return self.action_manager.get_action(self.action_manager.space.sample())
class PeriodicAgent(AbstractScriptedAgent): class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"):
"""Agent that does nothing most of the time, but executes application at regular intervals (with variance).""" """Agent that does nothing most of the time, but executes application at regular intervals (with variance)."""
class Settings(BaseModel): class Settings(BaseModel):

55
src/primaite/game/agent/scripted_agents/tap001.py
View File

@@ -1,4 +1,6 @@
# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from __future__ import annotations
import random import random
from typing import Dict, Tuple from typing import Dict, Tuple
@@ -7,20 +9,27 @@ from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent from primaite.game.agent.interface import AbstractScriptedAgent
class TAP001(AbstractScriptedAgent): class TAP001(AbstractScriptedAgent, identifier="TAP001"):
""" """
TAP001 | Mobile Malware -- Ransomware Variant. TAP001 | Mobile Malware -- Ransomware Variant.
Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application) Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application)
""" """
def __init__(self, *args, **kwargs): # TODO: Link with DataManipulationAgent via a parent "TAP" agent class.
super().__init__(*args, **kwargs)
self.setup_agent()
next_execution_timestep: int = 0 config: "TAP001.ConfigSchema"
starting_node_idx: int = 0
installed: bool = False class ConfigSchema(AbstractScriptedAgent.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()
def _set_next_execution_timestep(self, timestep: int) -> None: def _set_next_execution_timestep(self, timestep: int) -> None:
"""Set the next execution timestep with a configured random variance. """Set the next execution timestep with a configured random variance.
@@ -28,9 +37,9 @@ class TAP001(AbstractScriptedAgent):
:param timestep: The timestep to add variance to. :param timestep: The timestep to add variance to.
""" """
random_timestep_increment = random.randint( random_timestep_increment = random.randint(
-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance -self.config.agent_settings.start_settings.variance, self.config.agent_settings.start_settings.variance
) )
self.next_execution_timestep = timestep + random_timestep_increment self.config.next_execution_timestep = timestep + random_timestep_increment
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]: def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
"""Waits until a specific timestep, then attempts to execute the ransomware application. """Waits until a specific timestep, then attempts to execute the ransomware application.
@@ -45,28 +54,28 @@ class TAP001(AbstractScriptedAgent):
:return: Action formatted in CAOS format :return: Action formatted in CAOS format
:rtype: Tuple[str, Dict] :rtype: Tuple[str, Dict]
""" """
if timestep < self.next_execution_timestep: if timestep < self.config.next_execution_timestep:
return "DONOTHING", {} return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency) self._set_next_execution_timestep(timestep + self.config.agent_settings.start_settings.frequency)
if not self.installed: if not self.config.installed:
self.installed = True self.config.installed = True
return "NODE_APPLICATION_INSTALL", { return "node_application_install", {
"node_id": self.starting_node_idx, "node_name": self.config.starting_node_name,
"application_name": "RansomwareScript", "application_name": "RansomwareScript",
} }
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0} return "node_application_execute", {"node_name": self.config.starting_node_name, "application_id": 0}
def setup_agent(self) -> None: def setup_agent(self) -> None:
"""Set the next execution timestep when the episode resets.""" """Set the next execution timestep when the episode resets."""
self._select_start_node() 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)
for n, act in self.action_manager.action_map.items(): for n, act in self.config.action_manager.action_map.items():
if not act[0] == "NODE_APPLICATION_INSTALL": if not act[0] == "node_application_install":
continue continue
if act[1]["node_id"] == self.starting_node_idx: if act[1]["node_name"] == self.config.starting_node_name:
self.ip_address = act[1]["ip_address"] self.ip_address = act[1]["ip_address"]
return return
raise RuntimeError("TAP001 agent could not find database server ip address in action map") raise RuntimeError("TAP001 agent could not find database server ip address in action map")
@@ -74,5 +83,7 @@ class TAP001(AbstractScriptedAgent):
def _select_start_node(self) -> None: 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.""" """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 # 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) 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.starting_node_idx = random.randint(0, num_nodes - 1)
self.logger.debug(f"Selected Starting node ID: {self.starting_node_idx}")

20
src/primaite/game/game.py
View File

@@ -547,6 +547,26 @@ class PrimaiteGame:
reward_function = RewardFunction.from_config(reward_function_cfg) reward_function = RewardFunction.from_config(reward_function_cfg)
# CREATE AGENT # CREATE AGENT
# TODO: MAKE THIS BIT WORK AND NOT THE IF/ELSE CHAIN OF HORRORS
# Pass through:
# config
# action manager
# observation_manager
# reward_function
new_agent_cfg = {
"action_manager": action_space,
"agent_name": agent_cfg["ref"],
"observation_manager": obs_space,
"agent_settings": agent_cfg.get("agent_settings", {}),
"reward_function": reward_function,
}
new_agent_cfg = agent_cfg["settings"]
# new_agent_cfg.update{}
new_agent = AbstractAgent._registry[agent_cfg["type"]].from_config(config=new_agent_cfg)
if agent_type == "ProbabilisticAgent": if agent_type == "ProbabilisticAgent":
# TODO: implement non-random agents and fix this parsing # TODO: implement non-random agents and fix this parsing
settings = agent_cfg.get("agent_settings", {}) settings = agent_cfg.get("agent_settings", {})

2
tests/conftest.py
View File

@@ -264,7 +264,7 @@ def example_network() -> Network:
return network return network
class ControlledAgent(AbstractAgent): class ControlledAgent(AbstractAgent, identifier="Controlled_Agent"):
"""Agent that can be controlled by the tests.""" """Agent that can be controlled by the tests."""
def __init__( def __init__(