oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Get the db admin green agent working

Browse Source
This commit is contained in:
Marek Wolan
2024-02-27 13:30:16 +00:00
parent c54f82fb1b
commit af8ca82fcb
16 changed files with 386 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
src/primaite/game/agent/actions.py
View File

@@ -607,7 +607,6 @@ class ActionManager:
def __init__(
self,
game: "PrimaiteGame", # reference to game for information lookup
actions: List[Dict], # stores list of actions available to agent
nodes: List[Dict], # extra configuration for each node
max_folders_per_node: int = 2, # allows calculating shape
@@ -618,7 +617,7 @@ class ActionManager:
max_acl_rules: int = 10, # allows calculating shape
protocols: List[str] = ["TCP", "UDP", "ICMP"], # allow mapping index to protocol
ports: List[str] = ["HTTP", "DNS", "ARP", "FTP", "NTP"], # allow mapping index to port
ip_address_list: Optional[List[str]] = None, # to allow us to map an index to an ip address.
ip_address_list: List[str] = [], # to allow us to map an index to an ip address.
act_map: Optional[Dict[int, Dict]] = None, # allows restricting set of possible actions
) -> None:
"""Init method for ActionManager.
@@ -649,7 +648,6 @@ class ActionManager:
:param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
:type act_map: Optional[Dict[int, Dict]]
"""
self.game: "PrimaiteGame" = game
self.node_names: List[str] = [n["node_name"] for n in nodes]
"""List of node names in this action space. The list order is the mapping between node index and node name."""
self.application_names: List[List[str]] = []
@@ -707,25 +705,7 @@ class ActionManager:
self.protocols: List[str] = protocols
self.ports: List[str] = ports
self.ip_address_list: List[str]
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
if ip_address_list is not None:
self.ip_address_list = ip_address_list
else:
self.ip_address_list = []
for node_name in self.node_names:
node_obj = self.game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
self.ip_address_list.append(nic_obj.ip_address)
self.ip_address_list: List[str] = ip_address_list
# action_args are settings which are applied to the action space as a whole.
global_action_args = {
@@ -958,6 +938,12 @@ class ActionManager:
:return: The constructed ActionManager.
:rtype: ActionManager
"""
# If the user has provided a list of IP addresses, use that. Otherwise, generate a list of IP addresses from
# the nodes in the simulation.
# TODO: refactor. Options:
# 1: This should be pulled out into it's own function for clarity
# 2: The simulation itself should be able to provide a list of IP addresses with its API, rather than having to
# go through the nodes here.
ip_address_order = cfg["options"].pop("ip_address_order", {})
ip_address_list = []
for entry in ip_address_order:
@@ -967,13 +953,22 @@ class ActionManager:
ip_address = node_obj.network_interface[nic_num].ip_address
ip_address_list.append(ip_address)
if not ip_address_list:
node_names = [n["node_name"] for n in cfg.get("nodes", {})]
for node_name in node_names:
node_obj = game.simulation.network.get_node_by_hostname(node_name)
if node_obj is None:
continue
network_interfaces = node_obj.network_interfaces
for nic_uuid, nic_obj in network_interfaces.items():
ip_address_list.append(nic_obj.ip_address)
obj = cls(
game=game,
actions=cfg["action_list"],
**cfg["options"],
protocols=game.options.protocols,
ports=game.options.ports,
ip_address_list=ip_address_list or None,
ip_address_list=ip_address_list,
act_map=cfg.get("action_map"),
)

10
src/primaite/game/agent/data_manipulation_bot.py
View File

@@ -1,5 +1,5 @@
import random
from typing import Dict, Tuple
from typing import Dict, Optional, Tuple
from gymnasium.core import ObsType
@@ -26,7 +26,7 @@ class DataManipulationAgent(AbstractScriptedAgent):
)
self.next_execution_timestep = timestep + random_timestep_increment
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
"""Randomly sample an action from the action space.
:param obs: _description_
@@ -36,12 +36,10 @@ class DataManipulationAgent(AbstractScriptedAgent):
:return: _description_
:rtype: Tuple[str, Dict]
"""
current_timestep = self.action_manager.game.step_counter
if current_timestep < self.next_execution_timestep:
if timestep < self.next_execution_timestep:
return "DONOTHING", {"dummy": 0}
self._set_next_execution_timestep(current_timestep + self.agent_settings.start_settings.frequency)
self._set_next_execution_timestep(timestep + self.agent_settings.start_settings.frequency)
return "NODE_APPLICATION_EXECUTE", {"node_id": self.starting_node_idx, "application_id": 0}

10
src/primaite/game/agent/interface.py
View File

@@ -112,7 +112,7 @@ class AbstractAgent(ABC):
return self.reward_function.update(state)
@abstractmethod
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
"""
Return an action to be taken in the environment.
@@ -122,6 +122,8 @@ class AbstractAgent(ABC):
:type obs: ObsType
:param reward: Reward from the previous action, defaults to None TODO: should this parameter even be accepted?
:type reward: float, optional
:param timestep: The current timestep in the simulation, used for non-RL agents. Optional
:type timestep: int
:return: Action to be taken in the environment.
:rtype: Tuple[str, Dict]
"""
@@ -144,13 +146,13 @@ class AbstractAgent(ABC):
class AbstractScriptedAgent(AbstractAgent):
"""Base class for actors which generate their own behaviour."""
...
pass
class RandomAgent(AbstractScriptedAgent):
"""Agent that ignores its observation and acts completely at random."""
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
"""Randomly sample an action from the action space.
:param obs: _description_
@@ -183,7 +185,7 @@ class ProxyAgent(AbstractAgent):
self.most_recent_action: ActType
self.flatten_obs: bool = agent_settings.flatten_obs if agent_settings else False
def get_action(self, obs: ObsType, reward: float = 0.0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
"""
Return the agent's most recent action, formatted in CAOS format.

43
src/primaite/game/agent/scripted_agents.py
View File

@@ -11,30 +11,39 @@ from primaite.game.agent.observations import ObservationManager
from primaite.game.agent.rewards import RewardFunction
class GreenUC2Agent(AbstractScriptedAgent):
"""Scripted agent which attempts to send web requests to a target node."""
class ProbabilisticAgent(AbstractScriptedAgent):
"""Scripted agent which randomly samples its action space with prescribed probabilities for each action."""
class Settings(pydantic.BaseModel):
"""Config schema for Probabilistic agent settings."""
class GreenUC2AgentSettings(pydantic.BaseModel):
model_config = pydantic.ConfigDict(extra="forbid")
"""Strict validation."""
action_probabilities: Dict[int, float]
"""Probability to perform each action in the action map. The sum of probabilities should sum to 1."""
random_seed: Optional[int] = None
"""Random seed. If set, each episode the agent will choose the same random sequence of actions."""
# TODO: give the option to still set a random seed, but have it vary each episode in a predictable way
# for example if the user sets seed 123, have it be 123 + episode_num, so that each ep it's the next seed.
@pydantic.field_validator("action_probabilities", mode="after")
@classmethod
def probabilities_sum_to_one(cls, v: Dict[int, float]) -> Dict[int, float]:
"""Make sure the probabilities sum to 1."""
if not abs(sum(v.values()) - 1) < 1e-6:
raise ValueError(f"Green action probabilities must sum to 1")
raise ValueError("Green action probabilities must sum to 1")
return v
@pydantic.field_validator("action_probabilities", mode="after")
@classmethod
def action_map_covered_correctly(cls, v: Dict[int, float]) -> Dict[int, float]:
"""Ensure that the keys of the probability dictionary cover all integers from 0 to N."""
if not all((i in v) for i in range(len(v))):
raise ValueError(
"Green action probabilities must be defined as a mapping where the keys are consecutive integers "
"from 0 to N."
)
return v
def __init__(
self,
@@ -52,23 +61,27 @@ class GreenUC2Agent(AbstractScriptedAgent):
# If seed not specified, set it to None so that numpy chooses a random one.
settings.setdefault("random_seed")
self.settings = GreenUC2Agent.GreenUC2AgentSettings(settings)
self.settings = ProbabilisticAgent.Settings(**settings)
self.rng = np.random.default_rng(self.settings.random_seed)
# convert probabilities from
self.probabilities = np.array[self.settings.action_probabilities.values()]
self.probabilities = np.asarray(list(self.settings.action_probabilities.values()))
super().__init__(agent_name, action_space, observation_space, reward_function)
def get_action(self, obs: ObsType, reward: float = 0) -> Tuple[str, Dict]:
def get_action(self, obs: ObsType, reward: float = 0.0, timestep: Optional[int] = None) -> Tuple[str, Dict]:
"""
Choose a random action from the action space.
The probability of each action is given by the corresponding index in ``self.probabilities``.
:param obs: Current observation of the simulation
:type obs: ObsType
:param reward: Reward for the last step, not used for scripted agents, defaults to 0
:type reward: float, optional
:return: Action to be taken in CAOS format.
:rtype: Tuple[str, Dict]
"""
choice = self.rng.choice(len(self.action_manager.action_map), p=self.probabilities)
return self.action_manager.get_action(choice)
raise NotImplementedError
class RedDatabaseCorruptingAgent(AbstractScriptedAgent):
"""Scripted agent which attempts to corrupt the database of the target node."""
raise NotImplementedError