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

#2869 - Agents Refactor

Browse Source
This commit is contained in:
Charlie Crane
2024-11-27 15:29:51 +00:00
parent 7435a4dee8
commit 917386d638
7 changed files with 95 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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

34
src/primaite/game/agent/scripted_agents/abstract_tap.py Normal file
View File

@@ -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

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

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

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

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

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

@@ -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}")

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

@@ -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

20
tests/conftest.py
View File

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