Merged PR 278: Enable the red agent to vary its start node
## Summary - Made the data manipulation red agent be able to choose between the two clients to start operating on - changed the attacker name in the config to 'data_manipulation_attacker' to because it is no longer tied to any client - Updated the documentation notebook accordingly. - Fixed a bug where the database client made a new connection every time it sent a SQL query (it tries to reuse its most recent one instead) - Fixed a bug where link loads were not being cleared between episodes (?) **warning** - the green agents are not working properly after reset right now, but I'm gonna fix this in the next ticket where I refactor episode reset. ## Test process - unit tests pass - UC2 notebook passes with both clients. (currently this doesn't work after an episode reset, but the very next thing I'm gonna work on is refactoring the reset, so I don't want to waste time fixing this.) ## Checklist - [x] PR is linked to a **work item** - [x] **acceptance criteria** of linked ticket are met - [x] performed **self-review** of the code - [x] written **tests** for any new functionality added with this PR - [ ] updated the **documentation** if this PR changes or adds functionality - [ ] written/updated **design docs** if this PR implements new functionality - [x] updated the **change log** - [x] ran **pre-commit** checks for code style - [n] attended to any **TO-DOs** left in the code Related work items: #2232
This commit is contained in:
@@ -1,21 +1,20 @@
|
||||
import random
|
||||
from typing import Dict, List, Tuple
|
||||
from typing import Dict, Tuple
|
||||
|
||||
from gymnasium.core import ObsType
|
||||
|
||||
from primaite.game.agent.interface import AbstractScriptedAgent
|
||||
from primaite.simulator.system.applications.red_applications.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
|
||||
starting_node_idx: int = 0
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
|
||||
self.reset_agent_for_episode()
|
||||
|
||||
def _set_next_execution_timestep(self, timestep: int) -> None:
|
||||
"""Set the next execution timestep with a configured random variance.
|
||||
@@ -44,9 +43,16 @@ class DataManipulationAgent(AbstractScriptedAgent):
|
||||
|
||||
self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency)
|
||||
|
||||
return "NODE_APPLICATION_EXECUTE", {"node_id": 0, "application_id": 0}
|
||||
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "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._select_start_node()
|
||||
self._set_next_execution_timestep(self.agent_settings.start_settings.start_step)
|
||||
|
||||
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)
|
||||
self.starting_node_idx = random.randint(0, num_nodes - 1)
|
||||
|
||||
Reference in New Issue
Block a user