58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
import random
|
|
from typing import Dict, List, Tuple
|
|
|
|
from gymnasium.core import ObsType
|
|
|
|
from primaite.game.agent.interface import AbstractScriptedAgent
|
|
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
|
|
|
|
|
|
class DataManipulationAgent(AbstractScriptedAgent):
|
|
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
|
|
|
|
data_manipulation_bots: List["DataManipulationBot"] = []
|
|
next_execution_timestep: int = 0
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
print("red start step: ", self.agent_settings.start_settings.start_step)
|
|
|
|
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
|
|
|
|
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.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance
|
|
)
|
|
self.next_execution_timestep = timestep + random_timestep_increment
|
|
print("next execution red step: ", self.next_execution_timestep)
|
|
|
|
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
|
|
"""Randomly sample an action from the action space.
|
|
|
|
:param obs: _description_
|
|
:type obs: ObsType
|
|
:param reward: _description_, defaults to None
|
|
:type reward: float, optional
|
|
:return: _description_
|
|
:rtype: Tuple[str, Dict]
|
|
"""
|
|
current_timestep = self.action_manager.game.step_counter
|
|
|
|
if current_timestep < self.next_execution_timestep:
|
|
print("red agent doing nothing")
|
|
return "DONOTHING", {"dummy": 0}
|
|
|
|
self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency)
|
|
|
|
print("red agent doing an execute")
|
|
return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
|
|
|
|
def reset_agent_for_episode(self) -> None:
|
|
"""Set the next execution timestep when the episode resets."""
|
|
super().reset_agent_for_episode()
|
|
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
|