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

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
class DataManipulationAgent(AbstractScriptedAgent):
class DataManipulationAgent(AbstractScriptedAgent, identifier="Data_Manipulation_Agent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
next_execution_timestep: 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):
super().__init__(*args, **kwargs)
self.setup_agent()
@@ -38,12 +48,15 @@ class DataManipulationAgent(AbstractScriptedAgent):
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do NOTHING")
return "DONOTHING", {}
self.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
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:
"""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
class ProbabilisticAgent(AbstractScriptedAgent):
class ProbabilisticAgent(AbstractScriptedAgent, identifier="Probabilistic_Agent"):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
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
class RandomAgent(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]:
@@ -27,7 +27,7 @@ class RandomAgent(AbstractScriptedAgent):
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)."""
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
from __future__ import annotations
import random
from typing import Dict, Tuple
@@ -7,20 +9,27 @@ from gymnasium.core import ObsType
from primaite.game.agent.interface import AbstractScriptedAgent
class TAP001(AbstractScriptedAgent):
class TAP001(AbstractScriptedAgent, identifier="TAP001"):
"""
TAP001 | Mobile Malware -- Ransomware Variant.
Scripted Red Agent. Capable of one action; launching the kill-chain (Ransomware Application)
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.setup_agent()
# TODO: Link with DataManipulationAgent via a parent "TAP" agent class.
next_execution_timestep: int = 0
starting_node_idx: int = 0
installed: bool = False
config: "TAP001.ConfigSchema"
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:
"""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.
"""
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]:
"""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
:rtype: Tuple[str, Dict]
"""
if timestep < self.next_execution_timestep:
return "DONOTHING", {}
if timestep < self.config.next_execution_timestep:
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:
self.installed = True
return "NODE_APPLICATION_INSTALL", {
"node_id": self.starting_node_idx,
if not self.config.installed:
self.config.installed = True
return "node_application_install", {
"node_name": self.config.starting_node_name,
"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:
"""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)
for n, act in self.action_manager.action_map.items():
if not act[0] == "NODE_APPLICATION_INSTALL":
self._set_next_execution_timestep(self.config.agent_settings.start_settings.start_step)
for n, act in self.config.action_manager.action_map.items():
if not act[0] == "node_application_install":
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"]
return
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:
"""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)
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}")