Make data manipulation agent inherit from periodic agent & fix it a bit

This commit is contained in:
Marek Wolan
2025-01-14 09:57:35 +00:00
parent a447c5f43c
commit e7cfeeafc0
3 changed files with 23 additions and 24 deletions

View File

@@ -150,7 +150,9 @@ agents:
reward_components:
- type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self)
agent_settings:
possible_start_nodes: [client_1, client_2]
starting_application_name: DataManipulationBot
start_step: 25
frequency: 20
variance: 5

View File

@@ -149,7 +149,9 @@ agents:
reward_components:
- type: DUMMY
agent_settings: # options specific to this particular agent type, basically args of __init__(self)
agent_settings:
possible_start_nodes: [client_1, client_2]
starting_application_name: DataManipulationBot
start_step: 25
frequency: 20
variance: 5

View File

@@ -1,29 +1,33 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from typing import Dict, Optional, Tuple
import random
from typing import Dict, List, Tuple
from gymnasium.core import ObsType
from pydantic import Field
from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent
from primaite.game.agent.scripted_agents.random_agent import PeriodicAgent
__all__ = "DataManipulationAgent"
class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingAgent"):
class DataManipulationAgent(PeriodicAgent, identifier="RedDatabaseCorruptingAgent"):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
config: "DataManipulationAgent.ConfigSchema" = Field(default_factory=lambda: DataManipulationAgent.ConfigSchema())
class ConfigSchema(AbstractTAPAgent.ConfigSchema):
class ConfigSchema(PeriodicAgent.ConfigSchema):
"""Configuration Schema for DataManipulationAgent."""
type: str = "RedDatabaseCorruptingAgent"
starting_application_name: Optional[str] = None
starting_application_name: str = "DataManipulationBot"
possible_start_nodes: List[str]
@property
def starting_node_name(self) -> str:
"""Returns the agents starting node name."""
return self.config.starting_node_name
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)
def get_action(self, obs: ObsType, timestep: int) -> Tuple[str, Dict]:
"""Waits until a specific timestep, then attempts to execute its data manipulation application.
@@ -35,22 +39,13 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA
:return: Action formatted in CAOS format
:rtype: Tuple[str, Dict]
"""
if self.starting_node_name or self.config is None:
self.setup_agent()
self.get_action(obs=obs, timestep=timestep)
if timestep < self.next_execution_timestep:
self.logger.debug(msg="Performing do nothing action")
return "do_nothing", {}
self._set_next_execution_timestep(timestep + self.config.frequency)
self._set_next_execution_timestep(timestep=timestep + self.config.frequency, variance=self.config.variance)
self.logger.info(msg="Performing a data manipulation attack!")
return "node_application_execute", {
"node_name": self.config.starting_node_name,
"node_name": self.start_node,
"application_name": self.config.starting_application_name,
}
def setup_agent(self) -> None:
"""Set the next execution timestep when the episode resets."""
self._select_start_node()
self._set_next_execution_timestep(self.config.start_step)