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:
Marek Wolan
2024-02-20 20:22:20 +00:00
13 changed files with 105 additions and 31 deletions

View File

@@ -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)

View File

@@ -1,6 +1,6 @@
"""PrimAITE game - Encapsulates the simulation and agents."""
from ipaddress import IPv4Address
from typing import Dict, List
from typing import Dict, List, Tuple
from pydantic import BaseModel, ConfigDict
@@ -152,8 +152,14 @@ class PrimaiteGame:
agent.update_reward(state)
agent.reward_function.total_reward += agent.reward_function.current_reward
def apply_agent_actions(self) -> None:
"""Apply all actions to simulation as requests."""
def apply_agent_actions(self) -> Dict[str, Tuple[str, Dict]]:
"""
Apply all actions to simulation as requests.
:return: A recap of each agent's actions, in CAOS format.
:rtype: Dict[str, Tuple[str, Dict]]
"""
agent_actions = {}
for agent in self.agents:
obs = agent.observation_manager.current_observation