From 48af0229637726c9fc953ecf54b2329947151a1a Mon Sep 17 00:00:00 2001 From: Jake Walker Date: Tue, 21 Nov 2023 13:41:38 +0000 Subject: [PATCH] Run agent at configured timesteps --- src/primaite/game/agent/interface.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index 94878947..d2479b38 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -1,4 +1,5 @@ """Interface for agents.""" +import random from abc import ABC, abstractmethod from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, TypeAlias, Union @@ -178,10 +179,13 @@ 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) + self.next_execution_timestep = self.agent_settings.start_settings.start_step + # get node ids that are part of the agent's observation space node_ids: List[str] = [n.where[-1] for n in self.observation_space.obs.nodes] # get all nodes from their ids @@ -207,10 +211,19 @@ class DataManipulationAgent(AbstractScriptedAgent): """ # TODO: Move this to the appropriate place # return self.action_space.get_action(self.action_space.space.sample()) + + timestep = self.action_space.session.step_counter + + if timestep < self.next_execution_timestep: + return "DONOTHING", {"dummy": 0} + + var = random.randint(-self.agent_settings.start_settings.variance, self.agent_settings.start_settings.variance) + self.next_execution_timestep = timestep + self.agent_settings.start_settings.frequency + var + for bot in self.data_manipulation_bots: bot.execute() - return ("DONOTHING", {"dummy": 0}) + return "DONOTHING", {"dummy": 0} class AbstractGATEAgent(AbstractAgent):