PrimAITE/src/primaite/game/agent/data_manipulation_bot.py

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)

        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

    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:
            return "DONOTHING", {"dummy": 0}

        self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency)

        return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}