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

#2417 Move classes to correct files

Browse Source
This commit is contained in:
Marek Wolan
2024-03-29 14:14:03 +00:00
parent 9123aff592
commit 22e1dfea2f
10 changed files with 1332 additions and 1913 deletions
Download Patch File Download Diff File Expand all files Collapse all files

187
src/primaite/game/agent/observations/acl_observation.py Normal file
View File

@@ -0,0 +1,187 @@
from __future__ import annotations
from ipaddress import IPv4Address
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
class ACLObservation(AbstractObservation, identifier="ACL"):
"""ACL observation, provides information about access control lists within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for ACLObservation."""
ip_list: Optional[List[IPv4Address]] = None
"""List of IP addresses."""
wildcard_list: Optional[List[str]] = None
"""List of wildcard strings."""
port_list: Optional[List[int]] = None
"""List of port numbers."""
protocol_list: Optional[List[str]] = None
"""List of protocol names."""
num_rules: Optional[int] = None
"""Number of ACL rules."""
def __init__(
self,
where: WhereType,
num_rules: int,
ip_list: List[IPv4Address],
wildcard_list: List[str],
port_list: List[int],
protocol_list: List[str],
) -> None:
"""
Initialize an ACL observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this ACL.
:type where: WhereType
:param num_rules: Number of ACL rules.
:type num_rules: int
:param ip_list: List of IP addresses.
:type ip_list: List[IPv4Address]
:param wildcard_list: List of wildcard strings.
:type wildcard_list: List[str]
:param port_list: List of port numbers.
:type port_list: List[int]
:param protocol_list: List of protocol names.
:type protocol_list: List[str]
"""
self.where = where
self.num_rules: int = num_rules
self.ip_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(ip_list)}
self.wildcard_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(wildcard_list)}
self.port_to_id: Dict[int, int] = {i + 2: p for i, p in enumerate(port_list)}
self.protocol_to_id: Dict[str, int] = {i + 2: p for i, p in enumerate(protocol_list)}
self.default_observation: Dict = {
i
+ 1: {
"position": i,
"permission": 0,
"source_ip_id": 0,
"source_wildcard_id": 0,
"source_port_id": 0,
"dest_ip_id": 0,
"dest_wildcard_id": 0,
"dest_port_id": 0,
"protocol_id": 0,
}
for i in range(self.num_rules)
}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing ACL rules.
:rtype: ObsType
"""
acl_state: Dict = access_from_nested_dict(state, self.where)
if acl_state is NOT_PRESENT_IN_STATE:
return self.default_observation
obs = {}
acl_items = dict(acl_state.items())
i = 1 # don't show rule 0 for compatibility reasons.
while i < self.num_rules + 1:
rule_state = acl_items[i]
if rule_state is None:
obs[i] = {
"position": i - 1,
"permission": 0,
"source_ip_id": 0,
"source_wildcard_id": 0,
"source_port_id": 0,
"dest_ip_id": 0,
"dest_wildcard_id": 0,
"dest_port_id": 0,
"protocol_id": 0,
}
else:
src_ip = rule_state["src_ip_address"]
src_node_id = self.ip_to_id.get(src_ip, 1)
dst_ip = rule_state["dst_ip_address"]
dst_node_ip = self.ip_to_id.get(dst_ip, 1)
src_wildcard = rule_state["source_wildcard_id"]
src_wildcard_id = self.wildcard_to_id.get(src_wildcard, 1)
dst_wildcard = rule_state["dest_wildcard_id"]
dst_wildcard_id = self.wildcard_to_id.get(dst_wildcard, 1)
src_port = rule_state["source_port_id"]
src_port_id = self.port_to_id.get(src_port, 1)
dst_port = rule_state["dest_port_id"]
dst_port_id = self.port_to_id.get(dst_port, 1)
protocol = rule_state["protocol"]
protocol_id = self.protocol_to_id.get(protocol, 1)
obs[i] = {
"position": i - 1,
"permission": rule_state["action"],
"source_ip_id": src_node_id,
"source_wildcard_id": src_wildcard_id,
"source_port_id": src_port_id,
"dest_ip_id": dst_node_ip,
"dest_wildcard_id": dst_wildcard_id,
"dest_port_id": dst_port_id,
"protocol_id": protocol_id,
}
i += 1
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for ACL rules.
:rtype: spaces.Space
"""
return spaces.Dict(
{
i
+ 1: spaces.Dict(
{
"position": spaces.Discrete(self.num_rules),
"permission": spaces.Discrete(3),
# adding two to lengths is to account for reserved values 0 (unused) and 1 (any)
"source_ip_id": spaces.Discrete(len(self.ip_to_id) + 2),
"source_wildcard_id": spaces.Discrete(len(self.wildcard_to_id) + 2),
"source_port_id": spaces.Discrete(len(self.port_to_id) + 2),
"dest_ip_id": spaces.Discrete(len(self.ip_to_id) + 2),
"dest_wildcard_id": spaces.Discrete(len(self.wildcard_to_id) + 2),
"dest_port_id": spaces.Discrete(len(self.port_to_id) + 2),
"protocol_id": spaces.Discrete(len(self.protocol_to_id) + 2),
}
)
for i in range(self.num_rules)
}
)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ACLObservation:
"""
Create an ACL observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the ACL observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this ACL's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed ACL observation instance.
:rtype: ACLObservation
"""
return cls(
where=parent_where + ["acl", "acl"],
num_rules=config.num_rules,
ip_list=config.ip_list,
wildcard_list=config.wildcard_list,
port_list=config.port_list,
protocol_list=config.protocol_list,
)

138
src/primaite/game/agent/observations/agent_observations.py
View File

@@ -1,138 +0,0 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
from gymnasium import spaces
from primaite.game.agent.observations.node_observations import NodeObservation
from primaite.game.agent.observations.observations import (
AbstractObservation,
AclObservation,
ICSObservation,
LinkObservation,
NullObservation,
)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class UC2BlueObservation(AbstractObservation):
"""Container for all observations used by the blue agent in UC2.
TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler
for the purpose of compiling several observation components.
"""
def __init__(
self,
nodes: List[NodeObservation],
links: List[LinkObservation],
acl: AclObservation,
ics: ICSObservation,
where: Optional[List[str]] = None,
) -> None:
"""Initialise UC2 blue observation.
:param nodes: List of node observations
:type nodes: List[NodeObservation]
:param links: List of link observations
:type links: List[LinkObservation]
:param acl: The Access Control List observation
:type acl: AclObservation
:param ics: The ICS observation
:type ics: ICSObservation
:param where: Where in the simulation state dict to find information. Not used in this particular observation
because it only compiles other observations and doesn't contribute any new information, defaults to None
:type where: Optional[List[str]], optional
"""
super().__init__()
self.where: Optional[Tuple[str]] = where
self.nodes: List[NodeObservation] = nodes
self.links: List[LinkObservation] = links
self.acl: AclObservation = acl
self.ics: ICSObservation = ics
self.default_observation: Dict = {
"NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
"LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)},
"ACL": self.acl.default_observation,
"ICS": self.ics.default_observation,
}
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict
:return: Observation
:rtype: Dict
"""
if self.where is None:
return self.default_observation
obs = {}
obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)}
obs["ACL"] = self.acl.observe(state)
obs["ICS"] = self.ics.observe(state)
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Space
:rtype: spaces.Space
"""
return spaces.Dict(
{
"NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
"LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}),
"ACL": self.acl.space,
"ICS": self.ics.space,
}
)
@classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation":
"""Create UC2 blue observation from a config.
:param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes,
links, ACL and ICS observations.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:return: Constructed UC2 blue observation
:rtype: UC2BlueObservation
"""
node_configs = config["nodes"]
num_services_per_node = config["num_services_per_node"]
num_folders_per_node = config["num_folders_per_node"]
num_files_per_folder = config["num_files_per_folder"]
num_nics_per_node = config["num_nics_per_node"]
nodes = [
NodeObservation.from_config(
config=n,
game=game,
num_services_per_node=num_services_per_node,
num_folders_per_node=num_folders_per_node,
num_files_per_folder=num_files_per_folder,
num_nics_per_node=num_nics_per_node,
)
for n in node_configs
]
link_configs = config["links"]
links = [LinkObservation.from_config(config=link, game=game) for link in link_configs]
acl_config = config["acl"]
acl = AclObservation.from_config(config=acl_config, game=game)
ics_config = config["ics"]
ics = ICSObservation.from_config(config=ics_config, game=game)
new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"])
return new

207
src/primaite/game/agent/observations/file_system_observations.py
View File

@@ -1,107 +1,130 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING from __future__ import annotations
from typing import Dict, Iterable, List, Optional
from gymnasium import spaces from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class FileObservation(AbstractObservation, identifier="FILE"):
"""File observation, provides status information about a file within the simulation environment."""
class FileObservation(AbstractObservation): class ConfigSchema(AbstractObservation.ConfigSchema):
"""Observation of a file on a node in the network.""" """Configuration schema for FileObservation."""
def __init__(self, where: Optional[Tuple[str]] = None) -> None: file_name: str
"""Name of the file, used for querying simulation state dictionary."""
include_num_access: Optional[bool] = None
"""Whether to include the number of accesses to the file in the observation."""
def __init__(self, where: WhereType, include_num_access: bool) -> None:
""" """
Initialise file observation. Initialize a file observation instance.
:param where: Store information about where in the simulation state dictionary to find the relevant information. :param where: Where in the simulation state dictionary to find the relevant information for this file.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with A typical location for a file might be
zeroes. ['network', 'nodes', <node_hostname>, 'file_system', 'folder', <folder_name>, 'files', <file_name>].
:type where: WhereType
A typical location for a file looks like this: :param include_num_access: Whether to include the number of accesses to the file in the observation.
['network','nodes',<node_hostname>,'file_system', 'folders',<folder_name>,'files',<file_name>] :type include_num_access: bool
:type where: Optional[List[str]]
""" """
super().__init__() self.where: WhereType = where
self.where: Optional[Tuple[str]] = where self.include_num_access: bool = include_num_access
self.default_observation: spaces.Space = {"health_status": 0}
"Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted."
def observe(self, state: Dict) -> Dict: self.default_observation: ObsType = {"health_status": 0}
"""Generate observation based on the current state of the simulation. if self.include_num_access:
self.default_observation["num_access"] = 0
:param state: Simulation state dictionary def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict :type state: Dict
:return: Observation :return: Observation containing the health status of the file and optionally the number of accesses.
:rtype: Dict :rtype: ObsType
""" """
if self.where is None:
return self.default_observation
file_state = access_from_nested_dict(state, self.where) file_state = access_from_nested_dict(state, self.where)
if file_state is NOT_PRESENT_IN_STATE: if file_state is NOT_PRESENT_IN_STATE:
return self.default_observation return self.default_observation
return {"health_status": file_state["visible_status"]} obs = {"health_status": file_state["visible_status"]}
if self.include_num_access:
obs["num_access"] = file_state["num_access"]
# raise NotImplementedError("TODO: need to fix num_access to use thresholds instead of raw value.")
return obs
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape. """
Gymnasium space object describing the observation space shape.
:return: Gymnasium space :return: Gymnasium space representing the observation space for file status.
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict({"health_status": spaces.Discrete(6)}) space = {"health_status": spaces.Discrete(6)}
if self.include_num_access:
space["num_access"] = spaces.Discrete(4)
return spaces.Dict(space)
@classmethod @classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame", parent_where: List[str] = None) -> "FileObservation": def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FileObservation:
"""Create file observation from a config.
:param config: Dictionary containing the configuration for this file observation.
:type config: Dict
:param game: _description_
:type game: PrimaiteGame
:param parent_where: _description_, defaults to None
:type parent_where: _type_, optional
:return: _description_
:rtype: _type_
""" """
return cls(where=parent_where + ["files", config["file_name"]]) Create a file observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the file observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this file's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed file observation instance.
:rtype: FileObservation
"""
return cls(where=parent_where + ["files", config.file_name], include_num_access=config.include_num_access)
class FolderObservation(AbstractObservation): class FolderObservation(AbstractObservation, identifier="FOLDER"):
"""Folder observation, including files inside of the folder.""" """Folder observation, provides status information about a folder within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for FolderObservation."""
folder_name: str
"""Name of the folder, used for querying simulation state dictionary."""
files: List[FileObservation.ConfigSchema] = []
"""List of file configurations within the folder."""
num_files: Optional[int] = None
"""Number of spaces for file observations in this folder."""
include_num_access: Optional[bool] = None
"""Whether files in this folder should include the number of accesses in their observation."""
def __init__( def __init__(
self, where: Optional[Tuple[str]] = None, files: List[FileObservation] = [], num_files_per_folder: int = 2 self, where: WhereType, files: Iterable[FileObservation], num_files: int, include_num_access: bool
) -> None: ) -> None:
"""Initialise folder Observation, including files inside the folder. """
Initialize a folder observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this folder. :param where: Where in the simulation state dictionary to find the relevant information for this folder.
A typical location for a file looks like this: A typical location for a folder might be ['network', 'nodes', <node_hostname>, 'folders', <folder_name>].
['network','nodes',<node_hostname>,'file_system', 'folders',<folder_name>] :type where: WhereType
:type where: Optional[List[str]] :param files: List of file observation instances within the folder.
:param max_files: As size of the space must remain static, define max files that can be in this folder :type files: Iterable[FileObservation]
, defaults to 5 :param num_files: Number of files expected in the folder.
:type max_files: int, optional :type num_files: int
:param file_positions: Defines the positioning within the observation space of particular files. This ensures :param include_num_access: Whether to include the number of accesses to files in the observation.
that even if new files are created, the existing files will always occupy the same space in the observation :type include_num_access: bool
space. The keys must be between 1 and max_files. Providing file_positions will reserve a spot in the
observation space for a file with that name, even if it's temporarily deleted, if it reappears with the same
name, it will take the position defined in this dict. Defaults to {}
:type file_positions: Dict[int, str], optional
""" """
super().__init__() self.where: WhereType = where
self.where: Optional[Tuple[str]] = where
self.files: List[FileObservation] = files self.files: List[FileObservation] = files
while len(self.files) < num_files_per_folder: while len(self.files) < num_files:
self.files.append(FileObservation()) self.files.append(FileObservation(where=None, include_num_access=include_num_access))
while len(self.files) > num_files_per_folder: while len(self.files) > num_files:
truncated_file = self.files.pop() truncated_file = self.files.pop()
msg = f"Too many files in folder observation. Truncating file {truncated_file}" msg = f"Too many files in folder observation. Truncating file {truncated_file}"
_LOGGER.warning(msg) _LOGGER.warning(msg)
@@ -111,16 +134,15 @@ class FolderObservation(AbstractObservation):
"FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)}, "FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)},
} }
def observe(self, state: Dict) -> Dict: def observe(self, state: Dict) -> ObsType:
"""Generate observation based on the current state of the simulation. """
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict :param state: Simulation state dictionary.
:return: Observation :type state: Dict
:rtype: Dict :return: Observation containing the health status of the folder and status of files within the folder.
:rtype: ObsType
""" """
if self.where is None:
return self.default_observation
folder_state = access_from_nested_dict(state, self.where) folder_state = access_from_nested_dict(state, self.where)
if folder_state is NOT_PRESENT_IN_STATE: if folder_state is NOT_PRESENT_IN_STATE:
return self.default_observation return self.default_observation
@@ -136,9 +158,10 @@ class FolderObservation(AbstractObservation):
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape. """
Gymnasium space object describing the observation space shape.
:return: Gymnasium space :return: Gymnasium space representing the observation space for folder status.
:rtype: spaces.Space :rtype: spaces.Space
""" """
return spaces.Dict( return spaces.Dict(
@@ -149,29 +172,23 @@ class FolderObservation(AbstractObservation):
) )
@classmethod @classmethod
def from_config( def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FolderObservation:
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]], num_files_per_folder: int = 2 """
) -> "FolderObservation": Create a folder observation from a configuration schema.
"""Create folder observation from a config. Also creates child file observations.
:param config: Dictionary containing the configuration for this folder observation. Includes the name of the :param config: Configuration schema containing the necessary information for the folder observation.
folder and the files inside of it. :type config: ConfigSchema
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary to find the information about this folder's :param parent_where: Where in the simulation state dictionary to find the information about this folder's
parent node. A typical location for a node ``where`` can be: parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
['network','nodes',<node_hostname>,'file_system'] :type parent_where: WhereType, optional
:type parent_where: Optional[List[str]] :return: Constructed folder observation instance.
:param num_files_per_folder: How many spaces for files are in this folder observation (to preserve static
observation size) , defaults to 2
:type num_files_per_folder: int, optional
:return: Constructed folder observation
:rtype: FolderObservation :rtype: FolderObservation
""" """
where = parent_where + ["folders", config["folder_name"]] where = parent_where + ["folders", config.folder_name]
file_configs = config["files"] # pass down shared/common config items
files = [FileObservation.from_config(config=f, game=game, parent_where=where) for f in file_configs] for file_config in config.files:
file_config.include_num_access = config.include_num_access
return cls(where=where, files=files, num_files_per_folder=num_files_per_folder) files = [FileObservation.from_config(config=f, parent_where=where) for f in config.files]
return cls(where=where, files=files, num_files=config.num_files, include_num_access=config.include_num_access)

213
src/primaite/game/agent/observations/firewall_observation.py Normal file
View File

@@ -0,0 +1,213 @@
from __future__ import annotations
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.acl_observation import ACLObservation
from primaite.game.agent.observations.nic_observations import PortObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
_LOGGER = getLogger(__name__)
class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
"""Firewall observation, provides status information about a firewall within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for FirewallObservation."""
hostname: str
"""Hostname of the firewall node, used for querying simulation state dictionary."""
ip_list: Optional[List[str]] = None
"""List of IP addresses for encoding ACLs."""
wildcard_list: Optional[List[str]] = None
"""List of IP wildcards for encoding ACLs."""
port_list: Optional[List[int]] = None
"""List of ports for encoding ACLs."""
protocol_list: Optional[List[str]] = None
"""List of protocols for encoding ACLs."""
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""
def __init__(
self,
where: WhereType,
ip_list: List[str],
wildcard_list: List[str],
port_list: List[int],
protocol_list: List[str],
num_rules: int,
) -> None:
"""
Initialize a firewall observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this firewall.
A typical location for a firewall might be ['network', 'nodes', <firewall_hostname>].
:type where: WhereType
:param ip_list: List of IP addresses.
:type ip_list: List[str]
:param wildcard_list: List of wildcard rules.
:type wildcard_list: List[str]
:param port_list: List of port numbers.
:type port_list: List[int]
:param protocol_list: List of protocol types.
:type protocol_list: List[str]
:param num_rules: Number of rules configured in the firewall.
:type num_rules: int
"""
self.where: WhereType = where
self.ports: List[PortObservation] = [
PortObservation(where=self.where + ["port", port_num]) for port_num in (1, 2, 3)
]
# TODO: check what the port nums are for firewall.
self.internal_inbound_acl = ACLObservation(
where=self.where + ["acl", "internal", "inbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.internal_outbound_acl = ACLObservation(
where=self.where + ["acl", "internal", "outbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.dmz_inbound_acl = ACLObservation(
where=self.where + ["acl", "dmz", "inbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.dmz_outbound_acl = ACLObservation(
where=self.where + ["acl", "dmz", "outbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.external_inbound_acl = ACLObservation(
where=self.where + ["acl", "external", "inbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.external_outbound_acl = ACLObservation(
where=self.where + ["acl", "external", "outbound"],
num_rules=num_rules,
ip_list=ip_list,
wildcard_list=wildcard_list,
port_list=port_list,
protocol_list=protocol_list,
)
self.default_observation = {
"PORTS": {i + 1: p.default_observation for i, p in enumerate(self.ports)},
"INTERNAL": {
"INBOUND": self.internal_inbound_acl.default_observation,
"OUTBOUND": self.internal_outbound_acl.default_observation,
},
"DMZ": {
"INBOUND": self.dmz_inbound_acl.default_observation,
"OUTBOUND": self.dmz_outbound_acl.default_observation,
},
"EXTERNAL": {
"INBOUND": self.external_inbound_acl.default_observation,
"OUTBOUND": self.external_outbound_acl.default_observation,
},
}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the status of ports and ACLs for internal, DMZ, and external traffic.
:rtype: ObsType
"""
obs = {
"PORTS": {i + 1: p.observe(state) for i, p in enumerate(self.ports)},
"INTERNAL": {
"INBOUND": self.internal_inbound_acl.observe(state),
"OUTBOUND": self.internal_outbound_acl.observe(state),
},
"DMZ": {
"INBOUND": self.dmz_inbound_acl.observe(state),
"OUTBOUND": self.dmz_outbound_acl.observe(state),
},
"EXTERNAL": {
"INBOUND": self.external_inbound_acl.observe(state),
"OUTBOUND": self.external_outbound_acl.observe(state),
},
}
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for firewall status.
:rtype: spaces.Space
"""
space = spaces.Dict(
{
"PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}),
"INTERNAL": spaces.Dict(
{
"INBOUND": self.internal_inbound_acl.space,
"OUTBOUND": self.internal_outbound_acl.space,
}
),
"DMZ": spaces.Dict(
{
"INBOUND": self.dmz_inbound_acl.space,
"OUTBOUND": self.dmz_outbound_acl.space,
}
),
"EXTERNAL": spaces.Dict(
{
"INBOUND": self.external_inbound_acl.space,
"OUTBOUND": self.external_outbound_acl.space,
}
),
}
)
return space
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FirewallObservation:
"""
Create a firewall observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the firewall observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this firewall's
parent node. A typical location for a node might be ['network', 'nodes', <firewall_hostname>].
:type parent_where: WhereType, optional
:return: Constructed firewall observation instance.
:rtype: FirewallObservation
"""
where = parent_where + ["nodes", config.hostname]
return cls(
where=where,
ip_list=config.ip_list,
wildcard_list=config.wildcard_list,
port_list=config.port_list,
protocol_list=config.protocol_list,
num_rules=config.num_rules,
)

229
src/primaite/game/agent/observations/host_observations.py Normal file
View File

@@ -0,0 +1,229 @@
from __future__ import annotations
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.file_system_observations import FolderObservation
from primaite.game.agent.observations.nic_observations import NICObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.observations.software_observation import ApplicationObservation, ServiceObservation
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
class HostObservation(AbstractObservation, identifier="HOST"):
"""Host observation, provides status information about a host within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for HostObservation."""
hostname: str
"""Hostname of the host, used for querying simulation state dictionary."""
services: List[ServiceObservation.ConfigSchema] = []
"""List of services to observe on the host."""
applications: List[ApplicationObservation.ConfigSchema] = []
"""List of applications to observe on the host."""
folders: List[FolderObservation.ConfigSchema] = []
"""List of folders to observe on the host."""
network_interfaces: List[NICObservation.ConfigSchema] = []
"""List of network interfaces to observe on the host."""
num_services: Optional[int] = None
"""Number of spaces for service observations on this host."""
num_applications: Optional[int] = None
"""Number of spaces for application observations on this host."""
num_folders: Optional[int] = None
"""Number of spaces for folder observations on this host."""
num_files: Optional[int] = None
"""Number of spaces for file observations on this host."""
num_nics: Optional[int] = None
"""Number of spaces for network interface observations on this host."""
include_nmne: Optional[bool] = None
"""Whether network interface observations should include number of malicious network events."""
include_num_access: Optional[bool] = None
"""Whether to include the number of accesses to files observations on this host."""
def __init__(
self,
where: WhereType,
services: List[ServiceObservation],
applications: List[ApplicationObservation],
folders: List[FolderObservation],
network_interfaces: List[NICObservation],
num_services: int,
num_applications: int,
num_folders: int,
num_files: int,
num_nics: int,
include_nmne: bool,
include_num_access: bool,
) -> None:
"""
Initialize a host observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this host.
A typical location for a host might be ['network', 'nodes', <hostname>].
:type where: WhereType
:param services: List of service observations on the host.
:type services: List[ServiceObservation]
:param applications: List of application observations on the host.
:type applications: List[ApplicationObservation]
:param folders: List of folder observations on the host.
:type folders: List[FolderObservation]
:param network_interfaces: List of network interface observations on the host.
:type network_interfaces: List[NICObservation]
:param num_services: Number of services to observe.
:type num_services: int
:param num_applications: Number of applications to observe.
:type num_applications: int
:param num_folders: Number of folders to observe.
:type num_folders: int
:param num_files: Number of files.
:type num_files: int
:param num_nics: Number of network interfaces.
:type num_nics: int
:param include_nmne: Flag to include network metrics and errors.
:type include_nmne: bool
:param include_num_access: Flag to include the number of accesses to files.
:type include_num_access: bool
"""
self.where: WhereType = where
# Ensure lists have lengths equal to specified counts by truncating or padding
self.services: List[ServiceObservation] = services
while len(self.services) < num_services:
self.services.append(ServiceObservation(where=None))
while len(self.services) > num_services:
truncated_service = self.services.pop()
msg = f"Too many services in Node observation space for node. Truncating service {truncated_service.where}"
_LOGGER.warning(msg)
self.applications: List[ApplicationObservation] = applications
while len(self.applications) < num_applications:
self.applications.append(ApplicationObservation(where=None))
while len(self.applications) > num_applications:
truncated_application = self.applications.pop()
msg = f"Too many applications in Node observation space for node. Truncating {truncated_application.where}"
_LOGGER.warning(msg)
self.folders: List[FolderObservation] = folders
while len(self.folders) < num_folders:
self.folders.append(
FolderObservation(where=None, files=[], num_files=num_files, include_num_access=include_num_access)
)
while len(self.folders) > num_folders:
truncated_folder = self.folders.pop()
msg = f"Too many folders in Node observation space for node. Truncating folder {truncated_folder.where}"
_LOGGER.warning(msg)
self.network_interfaces: List[NICObservation] = network_interfaces
while len(self.network_interfaces) < num_nics:
self.network_interfaces.append(NICObservation(where=None, include_nmne=include_nmne))
while len(self.network_interfaces) > num_nics:
truncated_nic = self.network_interfaces.pop()
msg = f"Too many network_interfaces in Node observation space for node. Truncating {truncated_nic.where}"
_LOGGER.warning(msg)
self.default_observation: ObsType = {
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
"APPLICATIONS": {i + 1: a.default_observation for i, a in enumerate(self.applications)},
"FOLDERS": {i + 1: f.default_observation for i, f in enumerate(self.folders)},
"NICS": {i + 1: n.default_observation for i, n in enumerate(self.network_interfaces)},
"operating_status": 0,
"num_file_creations": 0,
"num_file_deletions": 0,
}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the status information about the host.
:rtype: ObsType
"""
node_state = access_from_nested_dict(state, self.where)
if node_state is NOT_PRESENT_IN_STATE:
return self.default_observation
obs = {}
obs["SERVICES"] = {i + 1: service.observe(state) for i, service in enumerate(self.services)}
obs["APPLICATIONS"] = {i + 1: app.observe(state) for i, app in enumerate(self.applications)}
obs["FOLDERS"] = {i + 1: folder.observe(state) for i, folder in enumerate(self.folders)}
obs["operating_status"] = node_state["operating_state"]
obs["NICS"] = {
i + 1: network_interface.observe(state) for i, network_interface in enumerate(self.network_interfaces)
}
obs["num_file_creations"] = node_state["file_system"]["num_file_creations"]
obs["num_file_deletions"] = node_state["file_system"]["num_file_deletions"]
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for host status.
:rtype: spaces.Space
"""
shape = {
"SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
"APPLICATIONS": spaces.Dict({i + 1: app.space for i, app in enumerate(self.applications)}),
"FOLDERS": spaces.Dict({i + 1: folder.space for i, folder in enumerate(self.folders)}),
"operating_status": spaces.Discrete(5),
"NICS": spaces.Dict(
{i + 1: network_interface.space for i, network_interface in enumerate(self.network_interfaces)}
),
"num_file_creations": spaces.Discrete(4),
"num_file_deletions": spaces.Discrete(4),
}
return spaces.Dict(shape)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = None) -> HostObservation:
"""
Create a host observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the host observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this host.
A typical location might be ['network', 'nodes', <hostname>].
:type parent_where: WhereType, optional
:return: Constructed host observation instance.
:rtype: HostObservation
"""
if parent_where is None:
where = ["network", "nodes", config.hostname]
else:
where = parent_where + ["nodes", config.hostname]
# Pass down shared/common config items
for folder_config in config.folders:
folder_config.include_num_access = config.include_num_access
folder_config.num_files = config.num_files
for nic_config in config.network_interfaces:
nic_config.include_nmne = config.include_nmne
services = [ServiceObservation.from_config(config=c, parent_where=where) for c in config.services]
applications = [ApplicationObservation.from_config(config=c, parent_where=where) for c in config.applications]
folders = [FolderObservation.from_config(config=c, parent_where=where) for c in config.folders]
nics = [NICObservation.from_config(config=c, parent_where=where) for c in config.network_interfaces]
return cls(
where=where,
services=services,
applications=applications,
folders=folders,
network_interfaces=nics,
num_services=config.num_services,
num_applications=config.num_applications,
num_folders=config.num_folders,
num_files=config.num_files,
num_nics=config.num_nics,
include_nmne=config.include_nmne,
include_num_access=config.include_num_access,
)

273
src/primaite/game/agent/observations/nic_observations.py
View File

@@ -1,188 +1,157 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING from __future__ import annotations
from typing import Dict, Optional
from gymnasium import spaces from gymnasium import spaces
from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
from primaite.simulator.network.nmne import CAPTURE_NMNE
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class NicObservation(AbstractObservation): class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
"""Observation of a Network Interface Card (NIC) in the network.""" """Status information about a network interface within the simulation environment."""
low_nmne_threshold: int = 0 class ConfigSchema(AbstractObservation.ConfigSchema):
"""The minimum number of malicious network events to be considered low.""" """Configuration schema for NICObservation."""
med_nmne_threshold: int = 5
"""The minimum number of malicious network events to be considered medium."""
high_nmne_threshold: int = 10
"""The minimum number of malicious network events to be considered high."""
global CAPTURE_NMNE nic_num: int
"""Number of the network interface."""
include_nmne: Optional[bool] = None
"""Whether to include number of malicious network events (NMNE) in the observation."""
@property def __init__(self, where: WhereType, include_nmne: bool) -> None:
def default_observation(self) -> Dict:
"""The default NIC observation dict."""
data = {"nic_status": 0}
if CAPTURE_NMNE:
data.update({"NMNE": {"inbound": 0, "outbound": 0}})
return data
def __init__(
self,
where: Optional[Tuple[str]] = None,
low_nmne_threshold: Optional[int] = 0,
med_nmne_threshold: Optional[int] = 5,
high_nmne_threshold: Optional[int] = 10,
) -> None:
"""Initialise NIC observation.
:param where: Where in the simulation state dictionary to find the relevant information for this NIC. A typical
example may look like this:
['network','nodes',<node_hostname>,'NICs',<nic_number>]
If None, this denotes that the NIC does not exist and the observation will be populated with zeroes.
:type where: Optional[Tuple[str]], optional
""" """
super().__init__() Initialize a network interface observation instance.
self.where: Optional[Tuple[str]] = where
global CAPTURE_NMNE :param where: Where in the simulation state dictionary to find the relevant information for this interface.
if CAPTURE_NMNE: A typical location for a network interface might be
self.nmne_inbound_last_step: int = 0 ['network', 'nodes', <node_hostname>, 'NICs', <nic_num>].
"""NMNEs persist for the whole episode, but we want to count per step. Keeping track of last step count lets :type where: WhereType
us find the difference.""" :param include_nmne: Flag to determine whether to include NMNE information in the observation.
self.nmne_outbound_last_step: int = 0 :type include_nmne: bool
"""NMNEs persist for the whole episode, but we want to count per step. Keeping track of last step count lets
us find the difference."""
if low_nmne_threshold or med_nmne_threshold or high_nmne_threshold:
self._validate_nmne_categories(
low_nmne_threshold=low_nmne_threshold,
med_nmne_threshold=med_nmne_threshold,
high_nmne_threshold=high_nmne_threshold,
)
def _validate_nmne_categories(
self, low_nmne_threshold: int = 0, med_nmne_threshold: int = 5, high_nmne_threshold: int = 10
):
""" """
Validates the nmne threshold config. self.where = where
self.include_nmne: bool = include_nmne
If the configuration is valid, the thresholds will be set, otherwise, an exception is raised. self.default_observation: ObsType = {"nic_status": 0}
if self.include_nmne:
self.default_observation.update({"NMNE": {"inbound": 0, "outbound": 0}})
:param: low_nmne_threshold: The minimum number of malicious network events to be considered low def observe(self, state: Dict) -> ObsType:
:param: med_nmne_threshold: The minimum number of malicious network events to be considered medium
:param: high_nmne_threshold: The minimum number of malicious network events to be considered high
""" """
if high_nmne_threshold <= med_nmne_threshold: Generate observation based on the current state of the simulation.
raise Exception(
f"nmne_categories: high nmne count ({high_nmne_threshold}) must be greater "
f"than medium nmne count ({med_nmne_threshold})"
)
if med_nmne_threshold <= low_nmne_threshold: :param state: Simulation state dictionary.
raise Exception(
f"nmne_categories: medium nmne count ({med_nmne_threshold}) must be greater "
f"than low nmne count ({low_nmne_threshold})"
)
self.high_nmne_threshold = high_nmne_threshold
self.med_nmne_threshold = med_nmne_threshold
self.low_nmne_threshold = low_nmne_threshold
def _categorise_mne_count(self, nmne_count: int) -> int:
"""
Categorise the number of Malicious Network Events (NMNEs) into discrete bins.
This helps in classifying the severity or volume of MNEs into manageable levels for the agent.
Bins are defined as follows:
- 0: No MNEs detected (0 events).
- 1: Low number of MNEs (default 1-5 events).
- 2: Moderate number of MNEs (default 6-10 events).
- 3: High number of MNEs (default more than 10 events).
:param nmne_count: Number of MNEs detected.
:return: Bin number corresponding to the number of MNEs. Returns 0, 1, 2, or 3 based on the detected MNE count.
"""
if nmne_count > self.high_nmne_threshold:
return 3
elif nmne_count > self.med_nmne_threshold:
return 2
elif nmne_count > self.low_nmne_threshold:
return 1
return 0
def observe(self, state: Dict) -> Dict:
"""Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary
:type state: Dict :type state: Dict
:return: Observation :return: Observation containing the status of the network interface and optionally NMNE information.
:rtype: Dict :rtype: ObsType
""" """
if self.where is None:
return self.default_observation
nic_state = access_from_nested_dict(state, self.where) nic_state = access_from_nested_dict(state, self.where)
if nic_state is NOT_PRESENT_IN_STATE: if nic_state is NOT_PRESENT_IN_STATE:
return self.default_observation return self.default_observation
else:
obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2} obs = {"nic_status": 1 if nic_state["enabled"] else 2}
if CAPTURE_NMNE: if self.include_nmne:
obs_dict.update({"NMNE": {}}) obs.update({"NMNE": {}})
direction_dict = nic_state["nmne"].get("direction", {}) direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {}) inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
inbound_count = inbound_keywords.get("*", 0) inbound_count = inbound_keywords.get("*", 0)
outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {}) outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {})
outbound_count = outbound_keywords.get("*", 0) outbound_count = outbound_keywords.get("*", 0)
obs_dict["NMNE"]["inbound"] = self._categorise_mne_count(inbound_count - self.nmne_inbound_last_step) obs["NMNE"]["inbound"] = self._categorise_mne_count(inbound_count - self.nmne_inbound_last_step)
obs_dict["NMNE"]["outbound"] = self._categorise_mne_count(outbound_count - self.nmne_outbound_last_step) obs["NMNE"]["outbound"] = self._categorise_mne_count(outbound_count - self.nmne_outbound_last_step)
self.nmne_inbound_last_step = inbound_count self.nmne_inbound_last_step = inbound_count
self.nmne_outbound_last_step = outbound_count self.nmne_outbound_last_step = outbound_count
return obs_dict return obs
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.""" """
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for network interface status and NMNE information.
:rtype: spaces.Space
"""
space = spaces.Dict({"nic_status": spaces.Discrete(3)}) space = spaces.Dict({"nic_status": spaces.Discrete(3)})
if CAPTURE_NMNE: if self.include_nmne:
space["NMNE"] = spaces.Dict({"inbound": spaces.Discrete(4), "outbound": spaces.Discrete(4)}) space["NMNE"] = spaces.Dict({"inbound": spaces.Discrete(4), "outbound": spaces.Discrete(4)})
return space return space
@classmethod @classmethod
def from_config(cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]]) -> "NicObservation": def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> NICObservation:
"""Create NIC observation from a config.
:param config: Dictionary containing the configuration for this NIC observation.
:type config: Dict
:param game: Reference to the PrimaiteGame object that spawned this observation.
:type game: PrimaiteGame
:param parent_where: Where in the simulation state dictionary to find the information about this NIC's parent
node. A typical location for a node ``where`` can be: ['network','nodes',<node_hostname>]
:type parent_where: Optional[List[str]]
:return: Constructed NIC observation
:rtype: NicObservation
""" """
low_nmne_threshold = None Create a network interface observation from a configuration schema.
med_nmne_threshold = None
high_nmne_threshold = None
if game and game.options and game.options.thresholds and game.options.thresholds.get("nmne"): :param config: Configuration schema containing the necessary information for the network interface observation.
threshold = game.options.thresholds["nmne"] :type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this NIC's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed network interface observation instance.
:rtype: NICObservation
"""
return cls(where=parent_where + ["NICs", config.nic_num], include_nmne=config.include_nmne)
low_nmne_threshold = int(threshold.get("low")) if threshold.get("low") is not None else None
med_nmne_threshold = int(threshold.get("medium")) if threshold.get("medium") is not None else None
high_nmne_threshold = int(threshold.get("high")) if threshold.get("high") is not None else None
return cls( class PortObservation(AbstractObservation, identifier="PORT"):
where=parent_where + ["NICs", config["nic_num"]], """Port observation, provides status information about a network port within the simulation environment."""
low_nmne_threshold=low_nmne_threshold,
med_nmne_threshold=med_nmne_threshold, class ConfigSchema(AbstractObservation.ConfigSchema):
high_nmne_threshold=high_nmne_threshold, """Configuration schema for PortObservation."""
)
port_id: int
"""Identifier of the port, used for querying simulation state dictionary."""
def __init__(self, where: WhereType) -> None:
"""
Initialize a port observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this port.
A typical location for a port might be ['network', 'nodes', <node_hostname>, 'NICs', <port_id>].
:type where: WhereType
"""
self.where = where
self.default_observation: ObsType = {"operating_status": 0}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the operating status of the port.
:rtype: ObsType
"""
port_state = access_from_nested_dict(state, self.where)
if port_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {"operating_status": 1 if port_state["enabled"] else 2}
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for port status.
:rtype: spaces.Space
"""
return spaces.Dict({"operating_status": spaces.Discrete(3)})
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> PortObservation:
"""
Create a port observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the port observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this port's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed port observation instance.
:rtype: PortObservation
"""
return cls(where=parent_where + ["NICs", config.port_id])

1197
src/primaite/game/agent/observations/node_observations.py
View File

File diff suppressed because it is too large Load Diff

467
src/primaite/game/agent/observations/observations.py
View File

@@ -1,24 +1,23 @@
"""Manages the observation space for the agent.""" """Manages the observation space for the agent."""
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from ipaddress import IPv4Address from typing import Any, Dict, Iterable, Type
from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Type
from gymnasium import spaces from gymnasium import spaces
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
from primaite import getLogger from primaite import getLogger
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__) _LOGGER = getLogger(__name__)
if TYPE_CHECKING: WhereType = Iterable[str | int] | None
from primaite.game.game import PrimaiteGame
class AbstractObservation(ABC): class AbstractObservation(ABC):
"""Abstract class for an observation space component.""" """Abstract class for an observation space component."""
class ConfigSchema(ABC, BaseModel): class ConfigSchema(ABC, BaseModel):
"""Config schema for observations."""
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
_registry: Dict[str, Type["AbstractObservation"]] = {} _registry: Dict[str, Type["AbstractObservation"]] = {}
@@ -61,269 +60,271 @@ class AbstractObservation(ABC):
@classmethod @classmethod
def from_config(cls, cfg: Dict) -> "AbstractObservation": def from_config(cls, cfg: Dict) -> "AbstractObservation":
"""Create this observation space component form a serialised format.""" """Create this observation space component form a serialised format."""
ObservationType = cls._registry[cfg['type']] ObservationType = cls._registry[cfg["type"]]
return ObservationType.from_config(cfg=cfg) return ObservationType.from_config(cfg=cfg)
# class LinkObservation(AbstractObservation): '''
# """Observation of a link in the network.""" class LinkObservation(AbstractObservation):
"""Observation of a link in the network."""
# default_observation: spaces.Space = {"PROTOCOLS": {"ALL": 0}} default_observation: spaces.Space = {"PROTOCOLS": {"ALL": 0}}
# "Default observation is what should be returned when the link doesn't exist." "Default observation is what should be returned when the link doesn't exist."
# def __init__(self, where: Optional[Tuple[str]] = None) -> None: def __init__(self, where: Optional[Tuple[str]] = None) -> None:
# """Initialise link observation. """Initialise link observation.
# :param where: Store information about where in the simulation state dictionary to find the relevant information. :param where: Store information about where in the simulation state dictionary to find the relevant information.
# Optional. If None, this corresponds that the file does not exist and the observation will be populated with Optional. If None, this corresponds that the file does not exist and the observation will be populated with
# zeroes. zeroes.
# A typical location for a service looks like this: A typical location for a service looks like this:
# `['network','nodes',<node_hostname>,'servics', <service_name>]` `['network','nodes',<node_hostname>,'servics', <service_name>]`
# :type where: Optional[List[str]] :type where: Optional[List[str]]
# """ """
# super().__init__() super().__init__()
# self.where: Optional[Tuple[str]] = where self.where: Optional[Tuple[str]] = where
# def observe(self, state: Dict) -> Dict: def observe(self, state: Dict) -> Dict:
# """Generate observation based on the current state of the simulation. """Generate observation based on the current state of the simulation.
# :param state: Simulation state dictionary :param state: Simulation state dictionary
# :type state: Dict :type state: Dict
# :return: Observation :return: Observation
# :rtype: Dict :rtype: Dict
# """ """
# if self.where is None: if self.where is None:
# return self.default_observation return self.default_observation
# link_state = access_from_nested_dict(state, self.where) link_state = access_from_nested_dict(state, self.where)
# if link_state is NOT_PRESENT_IN_STATE: if link_state is NOT_PRESENT_IN_STATE:
# return self.default_observation return self.default_observation
# bandwidth = link_state["bandwidth"] bandwidth = link_state["bandwidth"]
# load = link_state["current_load"] load = link_state["current_load"]
# if load == 0: if load == 0:
# utilisation_category = 0 utilisation_category = 0
# else: else:
# utilisation_fraction = load / bandwidth utilisation_fraction = load / bandwidth
# # 0 is UNUSED, 1 is 0%-10%. 2 is 10%-20%. 3 is 20%-30%. And so on... 10 is exactly 100% # 0 is UNUSED, 1 is 0%-10%. 2 is 10%-20%. 3 is 20%-30%. And so on... 10 is exactly 100%
# utilisation_category = int(utilisation_fraction * 9) + 1 utilisation_category = int(utilisation_fraction * 9) + 1
# # TODO: once the links support separte load per protocol, this needs amendment to reflect that. # TODO: once the links support separte load per protocol, this needs amendment to reflect that.
# return {"PROTOCOLS": {"ALL": min(utilisation_category, 10)}} return {"PROTOCOLS": {"ALL": min(utilisation_category, 10)}}
# @property @property
# def space(self) -> spaces.Space: def space(self) -> spaces.Space:
# """Gymnasium space object describing the observation space shape. """Gymnasium space object describing the observation space shape.
# :return: Gymnasium space :return: Gymnasium space
# :rtype: spaces.Space :rtype: spaces.Space
# """ """
# return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})}) return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
# @classmethod @classmethod
# def from_config(cls, config: Dict, game: "PrimaiteGame") -> "LinkObservation": def from_config(cls, config: Dict, game: "PrimaiteGame") -> "LinkObservation":
# """Create link observation from a config. """Create link observation from a config.
# :param config: Dictionary containing the configuration for this link observation. :param config: Dictionary containing the configuration for this link observation.
# :type config: Dict :type config: Dict
# :param game: Reference to the PrimaiteGame object that spawned this observation. :param game: Reference to the PrimaiteGame object that spawned this observation.
# :type game: PrimaiteGame :type game: PrimaiteGame
# :return: Constructed link observation :return: Constructed link observation
# :rtype: LinkObservation :rtype: LinkObservation
# """ """
# return cls(where=["network", "links", game.ref_map_links[config["link_ref"]]]) return cls(where=["network", "links", game.ref_map_links[config["link_ref"]]])
# class AclObservation(AbstractObservation): class AclObservation(AbstractObservation):
# """Observation of an Access Control List (ACL) in the network.""" """Observation of an Access Control List (ACL) in the network."""
# # TODO: should where be optional, and we can use where=None to pad the observation space? # TODO: should where be optional, and we can use where=None to pad the observation space?
# # definitely the current approach does not support tracking files that aren't specified by name, for example # definitely the current approach does not support tracking files that aren't specified by name, for example
# # if a file is created at runtime, we have currently got no way of telling the observation space to track it. # if a file is created at runtime, we have currently got no way of telling the observation space to track it.
# # this needs adding, but not for the MVP. # this needs adding, but not for the MVP.
# def __init__( def __init__(
# self, self,
# node_ip_to_id: Dict[str, int], node_ip_to_id: Dict[str, int],
# ports: List[int], ports: List[int],
# protocols: List[str], protocols: List[str],
# where: Optional[Tuple[str]] = None, where: Optional[Tuple[str]] = None,
# num_rules: int = 10, num_rules: int = 10,
# ) -> None: ) -> None:
# """Initialise ACL observation. """Initialise ACL observation.
# :param node_ip_to_id: Mapping between IP address and ID. :param node_ip_to_id: Mapping between IP address and ID.
# :type node_ip_to_id: Dict[str, int] :type node_ip_to_id: Dict[str, int]
# :param ports: List of ports which are part of the game that define the ordering when converting to an ID :param ports: List of ports which are part of the game that define the ordering when converting to an ID
# :type ports: List[int] :type ports: List[int]
# :param protocols: List of protocols which are part of the game, defines ordering when converting to an ID :param protocols: List of protocols which are part of the game, defines ordering when converting to an ID
# :type protocols: list[str] :type protocols: list[str]
# :param where: Where in the simulation state dictionary to find the relevant information for this ACL. A typical :param where: Where in the simulation state dictionary to find the relevant information for this ACL. A typical
# example may look like this: example may look like this:
# ['network','nodes',<router_hostname>,'acl','acl'] ['network','nodes',<router_hostname>,'acl','acl']
# :type where: Optional[Tuple[str]], optional :type where: Optional[Tuple[str]], optional
# :param num_rules: , defaults to 10 :param num_rules: , defaults to 10
# :type num_rules: int, optional :type num_rules: int, optional
# """ """
# super().__init__() super().__init__()
# self.where: Optional[Tuple[str]] = where self.where: Optional[Tuple[str]] = where
# self.num_rules: int = num_rules self.num_rules: int = num_rules
# self.node_to_id: Dict[str, int] = node_ip_to_id self.node_to_id: Dict[str, int] = node_ip_to_id
# "List of node IP addresses, order in this list determines how they are converted to an ID" "List of node IP addresses, order in this list determines how they are converted to an ID"
# self.port_to_id: Dict[int, int] = {port: i + 2 for i, port in enumerate(ports)} self.port_to_id: Dict[int, int] = {port: i + 2 for i, port in enumerate(ports)}
# "List of ports which are part of the game that define the ordering when converting to an ID" "List of ports which are part of the game that define the ordering when converting to an ID"
# self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)} self.protocol_to_id: Dict[str, int] = {protocol: i + 2 for i, protocol in enumerate(protocols)}
# "List of protocols which are part of the game, defines ordering when converting to an ID" "List of protocols which are part of the game, defines ordering when converting to an ID"
# self.default_observation: Dict = { self.default_observation: Dict = {
# i i
# + 1: { + 1: {
# "position": i, "position": i,
# "permission": 0, "permission": 0,
# "source_node_id": 0, "source_node_id": 0,
# "source_port": 0, "source_port": 0,
# "dest_node_id": 0, "dest_node_id": 0,
# "dest_port": 0, "dest_port": 0,
# "protocol": 0, "protocol": 0,
# } }
# for i in range(self.num_rules) for i in range(self.num_rules)
# } }
# def observe(self, state: Dict) -> Dict: def observe(self, state: Dict) -> Dict:
# """Generate observation based on the current state of the simulation. """Generate observation based on the current state of the simulation.
# :param state: Simulation state dictionary :param state: Simulation state dictionary
# :type state: Dict :type state: Dict
# :return: Observation :return: Observation
# :rtype: Dict :rtype: Dict
# """ """
# if self.where is None: if self.where is None:
# return self.default_observation return self.default_observation
# acl_state: Dict = access_from_nested_dict(state, self.where) acl_state: Dict = access_from_nested_dict(state, self.where)
# if acl_state is NOT_PRESENT_IN_STATE: if acl_state is NOT_PRESENT_IN_STATE:
# return self.default_observation return self.default_observation
# # TODO: what if the ACL has more rules than num of max rules for obs space # TODO: what if the ACL has more rules than num of max rules for obs space
# obs = {} obs = {}
# acl_items = dict(acl_state.items()) acl_items = dict(acl_state.items())
# i = 1 # don't show rule 0 for compatibility reasons. i = 1 # don't show rule 0 for compatibility reasons.
# while i < self.num_rules + 1: while i < self.num_rules + 1:
# rule_state = acl_items[i] rule_state = acl_items[i]
# if rule_state is None: if rule_state is None:
# obs[i] = { obs[i] = {
# "position": i - 1, "position": i - 1,
# "permission": 0, "permission": 0,
# "source_node_id": 0, "source_node_id": 0,
# "source_port": 0, "source_port": 0,
# "dest_node_id": 0, "dest_node_id": 0,
# "dest_port": 0, "dest_port": 0,
# "protocol": 0, "protocol": 0,
# } }
# else: else:
# src_ip = rule_state["src_ip_address"] src_ip = rule_state["src_ip_address"]
# src_node_id = 1 if src_ip is None else self.node_to_id[IPv4Address(src_ip)] src_node_id = 1 if src_ip is None else self.node_to_id[IPv4Address(src_ip)]
# dst_ip = rule_state["dst_ip_address"] dst_ip = rule_state["dst_ip_address"]
# dst_node_ip = 1 if dst_ip is None else self.node_to_id[IPv4Address(dst_ip)] dst_node_ip = 1 if dst_ip is None else self.node_to_id[IPv4Address(dst_ip)]
# src_port = rule_state["src_port"] src_port = rule_state["src_port"]
# src_port_id = 1 if src_port is None else self.port_to_id[src_port] src_port_id = 1 if src_port is None else self.port_to_id[src_port]
# dst_port = rule_state["dst_port"] dst_port = rule_state["dst_port"]
# dst_port_id = 1 if dst_port is None else self.port_to_id[dst_port] dst_port_id = 1 if dst_port is None else self.port_to_id[dst_port]
# protocol = rule_state["protocol"] protocol = rule_state["protocol"]
# protocol_id = 1 if protocol is None else self.protocol_to_id[protocol] protocol_id = 1 if protocol is None else self.protocol_to_id[protocol]
# obs[i] = { obs[i] = {
# "position": i - 1, "position": i - 1,
# "permission": rule_state["action"], "permission": rule_state["action"],
# "source_node_id": src_node_id, "source_node_id": src_node_id,
# "source_port": src_port_id, "source_port": src_port_id,
# "dest_node_id": dst_node_ip, "dest_node_id": dst_node_ip,
# "dest_port": dst_port_id, "dest_port": dst_port_id,
# "protocol": protocol_id, "protocol": protocol_id,
# } }
# i += 1 i += 1
# return obs return obs
# @property @property
# def space(self) -> spaces.Space: def space(self) -> spaces.Space:
# """Gymnasium space object describing the observation space shape. """Gymnasium space object describing the observation space shape.
# :return: Gymnasium space :return: Gymnasium space
# :rtype: spaces.Space :rtype: spaces.Space
# """ """
# return spaces.Dict( return spaces.Dict(
# { {
# i i
# + 1: spaces.Dict( + 1: spaces.Dict(
# { {
# "position": spaces.Discrete(self.num_rules), "position": spaces.Discrete(self.num_rules),
# "permission": spaces.Discrete(3), "permission": spaces.Discrete(3),
# # adding two to lengths is to account for reserved values 0 (unused) and 1 (any) # adding two to lengths is to account for reserved values 0 (unused) and 1 (any)
# "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
# "source_port": spaces.Discrete(len(self.port_to_id) + 2), "source_port": spaces.Discrete(len(self.port_to_id) + 2),
# "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2), "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
# "dest_port": spaces.Discrete(len(self.port_to_id) + 2), "dest_port": spaces.Discrete(len(self.port_to_id) + 2),
# "protocol": spaces.Discrete(len(self.protocol_to_id) + 2), "protocol": spaces.Discrete(len(self.protocol_to_id) + 2),
# } }
# ) )
# for i in range(self.num_rules) for i in range(self.num_rules)
# } }
# ) )
# @classmethod @classmethod
# def from_config(cls, config: Dict, game: "PrimaiteGame") -> "AclObservation": def from_config(cls, config: Dict, game: "PrimaiteGame") -> "AclObservation":
# """Generate ACL observation from a config. """Generate ACL observation from a config.
# :param config: Dictionary containing the configuration for this ACL observation. :param config: Dictionary containing the configuration for this ACL observation.
# :type config: Dict :type config: Dict
# :param game: Reference to the PrimaiteGame object that spawned this observation. :param game: Reference to the PrimaiteGame object that spawned this observation.
# :type game: PrimaiteGame :type game: PrimaiteGame
# :return: Observation object :return: Observation object
# :rtype: AclObservation :rtype: AclObservation
# """ """
# max_acl_rules = config["options"]["max_acl_rules"] max_acl_rules = config["options"]["max_acl_rules"]
# node_ip_to_idx = {} node_ip_to_idx = {}
# for ip_idx, ip_map_config in enumerate(config["ip_address_order"]): for ip_idx, ip_map_config in enumerate(config["ip_address_order"]):
# node_ref = ip_map_config["node_hostname"] node_ref = ip_map_config["node_hostname"]
# nic_num = ip_map_config["nic_num"] nic_num = ip_map_config["nic_num"]
# node_obj = game.simulation.network.nodes[game.ref_map_nodes[node_ref]] node_obj = game.simulation.network.nodes[game.ref_map_nodes[node_ref]]
# nic_obj = node_obj.network_interface[nic_num] nic_obj = node_obj.network_interface[nic_num]
# node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2 node_ip_to_idx[nic_obj.ip_address] = ip_idx + 2
# router_hostname = config["router_hostname"] router_hostname = config["router_hostname"]
# return cls( return cls(
# node_ip_to_id=node_ip_to_idx, node_ip_to_id=node_ip_to_idx,
# ports=game.options.ports, ports=game.options.ports,
# protocols=game.options.protocols, protocols=game.options.protocols,
# where=["network", "nodes", router_hostname, "acl", "acl"], where=["network", "nodes", router_hostname, "acl", "acl"],
# num_rules=max_acl_rules, num_rules=max_acl_rules,
# ) )
# class NullObservation(AbstractObservation): class NullObservation(AbstractObservation):
# """Null observation, returns a single 0 value for the observation space.""" """Null observation, returns a single 0 value for the observation space."""
# def __init__(self, where: Optional[List[str]] = None): def __init__(self, where: Optional[List[str]] = None):
# """Initialise null observation.""" """Initialise null observation."""
# self.default_observation: Dict = {} self.default_observation: Dict = {}
# def observe(self, state: Dict) -> Dict: def observe(self, state: Dict) -> Dict:
# """Generate observation based on the current state of the simulation.""" """Generate observation based on the current state of the simulation."""
# return 0 return 0
# @property @property
# def space(self) -> spaces.Space: def space(self) -> spaces.Space:
# """Gymnasium space object describing the observation space shape.""" """Gymnasium space object describing the observation space shape."""
# return spaces.Discrete(1) return spaces.Discrete(1)
# @classmethod @classmethod
# def from_config(cls, config: Dict, game: Optional["PrimaiteGame"] = None) -> "NullObservation": def from_config(cls, config: Dict, game: Optional["PrimaiteGame"] = None) -> "NullObservation":
# """ """
# Create null observation from a config. Create null observation from a config.
# The parameters are ignored, they are here to match the signature of the other observation classes. The parameters are ignored, they are here to match the signature of the other observation classes.
# """ """
# return cls() return cls()
# class ICSObservation(NullObservation): class ICSObservation(NullObservation):
# """ICS observation placeholder, currently not implemented so always returns a single 0.""" """ICS observation placeholder, currently not implemented so always returns a single 0."""
# pass pass
'''

142
src/primaite/game/agent/observations/router_observation.py Normal file
View File

@@ -0,0 +1,142 @@
from __future__ import annotations
from typing import Dict, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.acl_observation import ACLObservation
from primaite.game.agent.observations.nic_observations import PortObservation
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
class RouterObservation(AbstractObservation, identifier="ROUTER"):
"""Router observation, provides status information about a router within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for RouterObservation."""
hostname: str
"""Hostname of the router, used for querying simulation state dictionary."""
ports: Optional[List[PortObservation.ConfigSchema]] = None
"""Configuration of port observations for this router."""
num_ports: Optional[int] = None
"""Number of port observations configured for this router."""
acl: Optional[ACLObservation.ConfigSchema] = None
"""Configuration of ACL observation on this router."""
ip_list: Optional[List[str]] = None
"""List of IP addresses for encoding ACLs."""
wildcard_list: Optional[List[str]] = None
"""List of IP wildcards for encoding ACLs."""
port_list: Optional[List[int]] = None
"""List of ports for encoding ACLs."""
protocol_list: Optional[List[str]] = None
"""List of protocols for encoding ACLs."""
num_rules: Optional[int] = None
"""Number of rules ACL rules to show."""
def __init__(
self,
where: WhereType,
ports: List[PortObservation],
num_ports: int,
acl: ACLObservation,
) -> None:
"""
Initialize a router observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this router.
A typical location for a router might be ['network', 'nodes', <node_hostname>].
:type where: WhereType
:param ports: List of port observations representing the ports of the router.
:type ports: List[PortObservation]
:param num_ports: Number of ports for the router.
:type num_ports: int
:param acl: ACL observation representing the access control list of the router.
:type acl: ACLObservation
"""
self.where: WhereType = where
self.ports: List[PortObservation] = ports
self.acl: ACLObservation = acl
self.num_ports: int = num_ports
while len(self.ports) < num_ports:
self.ports.append(PortObservation(where=None))
while len(self.ports) > num_ports:
self.ports.pop()
msg = "Too many ports in router observation. Truncating."
_LOGGER.warning(msg)
self.default_observation = {
"PORTS": {i + 1: p.default_observation for i, p in enumerate(self.ports)},
"ACL": self.acl.default_observation,
}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the status of ports and ACL configuration of the router.
:rtype: ObsType
"""
router_state = access_from_nested_dict(state, self.where)
if router_state is NOT_PRESENT_IN_STATE:
return self.default_observation
obs = {}
obs["PORTS"] = {i + 1: p.observe(state) for i, p in enumerate(self.ports)}
obs["ACL"] = self.acl.observe(state)
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for router status.
:rtype: spaces.Space
"""
return spaces.Dict(
{"PORTS": spaces.Dict({i + 1: p.space for i, p in enumerate(self.ports)}), "ACL": self.acl.space}
)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> RouterObservation:
"""
Create a router observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the router observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this router's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed router observation instance.
:rtype: RouterObservation
"""
where = parent_where + ["nodes", config.hostname]
if config.acl is None:
config.acl = ACLObservation.ConfigSchema()
if config.acl.num_rules is None:
config.acl.num_rules = config.num_rules
if config.acl.ip_list is None:
config.acl.ip_list = config.ip_list
if config.acl.wildcard_list is None:
config.acl.wildcard_list = config.wildcard_list
if config.acl.port_list is None:
config.acl.port_list = config.port_list
if config.acl.protocol_list is None:
config.acl.protocol_list = config.protocol_list
if config.ports is None:
config.ports = [PortObservation.ConfigSchema(port_id=i + 1) for i in range(config.num_ports)]
ports = [PortObservation.from_config(config=c, parent_where=where) for c in config.ports]
acl = ACLObservation.from_config(config=config.acl, parent_where=where)
return cls(where=where, ports=ports, num_ports=config.num_ports, acl=acl)

192
src/primaite/game/agent/observations/software_observation.py
View File

@@ -1,45 +1,43 @@
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING from __future__ import annotations
from typing import Dict
from gymnasium import spaces from gymnasium import spaces
from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
if TYPE_CHECKING:
from primaite.game.game import PrimaiteGame
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
"""Service observation, shows status of a service in the simulation environment."""
class ServiceObservation(AbstractObservation): class ConfigSchema(AbstractObservation.ConfigSchema):
"""Observation of a service in the network.""" """Configuration schema for ServiceObservation."""
default_observation: spaces.Space = {"operating_status": 0, "health_status": 0} service_name: str
"Default observation is what should be returned when the service doesn't exist." """Name of the service, used for querying simulation state dictionary"""
def __init__(self, where: Optional[Tuple[str]] = None) -> None: def __init__(self, where: WhereType) -> None:
"""Initialise service observation.
:param where: Store information about where in the simulation state dictionary to find the relevant information.
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a service looks like this:
`['network','nodes',<node_hostname>,'services', <service_name>]`
:type where: Optional[List[str]]
""" """
super().__init__() Initialize a service observation instance.
self.where: Optional[Tuple[str]] = where
def observe(self, state: Dict) -> Dict: :param where: Where in the simulation state dictionary to find the relevant information for this service.
"""Generate observation based on the current state of the simulation. A typical location for a service might be ['network', 'nodes', <node_hostname>, 'services', <service_name>].
:type where: WhereType
"""
self.where = where
self.default_observation = {"operating_status": 0, "health_status": 0}
:param state: Simulation state dictionary def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict :type state: Dict
:return: Observation :return: Observation containing the operating status and health status of the service.
:rtype: Dict :rtype: ObsType
""" """
if self.where is None:
return self.default_observation
service_state = access_from_nested_dict(state, self.where) service_state = access_from_nested_dict(state, self.where)
if service_state is NOT_PRESENT_IN_STATE: if service_state is NOT_PRESENT_IN_STATE:
return self.default_observation return self.default_observation
@@ -50,114 +48,96 @@ class ServiceObservation(AbstractObservation):
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.""" """
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for service status.
:rtype: spaces.Space
"""
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)}) return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
@classmethod @classmethod
def from_config( def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ServiceObservation:
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]] = None """
) -> "ServiceObservation": Create a service observation from a configuration schema.
"""Create service observation from a config.
:param config: Dictionary containing the configuration for this service observation. :param config: Configuration schema containing the necessary information for the service observation.
:type config: Dict :type config: ConfigSchema
:param game: Reference to the PrimaiteGame object that spawned this observation. :param parent_where: Where in the simulation state dictionary to find the information about this service's
:type game: PrimaiteGame parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional. :type parent_where: WhereType, optional
:type parent_where: Optional[List[str]], optional :return: Constructed service observation instance.
:return: Constructed service observation
:rtype: ServiceObservation :rtype: ServiceObservation
""" """
return cls(where=parent_where + ["services", config["service_name"]]) return cls(where=parent_where + ["services", config.service_name])
class ApplicationObservation(AbstractObservation): class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
"""Observation of an application in the network.""" """Application observation, shows the status of an application within the simulation environment."""
default_observation: spaces.Space = {"operating_status": 0, "health_status": 0, "num_executions": 0} class ConfigSchema(AbstractObservation.ConfigSchema):
"Default observation is what should be returned when the application doesn't exist." """Configuration schema for ApplicationObservation."""
def __init__(self, where: Optional[Tuple[str]] = None) -> None: application_name: str
"""Initialise application observation. """Name of the application, used for querying simulation state dictionary"""
:param where: Store information about where in the simulation state dictionary to find the relevant information. def __init__(self, where: WhereType) -> None:
Optional. If None, this corresponds that the file does not exist and the observation will be populated with
zeroes.
A typical location for a service looks like this:
`['network','nodes',<node_hostname>,'applications', <application_name>]`
:type where: Optional[List[str]]
""" """
super().__init__() Initialise an application observation instance.
self.where: Optional[Tuple[str]] = where
def observe(self, state: Dict) -> Dict: :param where: Where in the simulation state dictionary to find the relevant information for this application.
"""Generate observation based on the current state of the simulation. A typical location for an application might be
['network', 'nodes', <node_hostname>, 'applications', <application_name>].
:type where: WhereType
"""
self.where = where
self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
:param state: Simulation state dictionary def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict :type state: Dict
:return: Observation :return: Obs containing the operating status, health status, and number of executions of the application.
:rtype: Dict :rtype: ObsType
""" """
if self.where is None: application_state = access_from_nested_dict(state, self.where)
return self.default_observation if application_state is NOT_PRESENT_IN_STATE:
app_state = access_from_nested_dict(state, self.where)
if app_state is NOT_PRESENT_IN_STATE:
return self.default_observation return self.default_observation
return { return {
"operating_status": app_state["operating_state"], "operating_status": application_state["operating_state"],
"health_status": app_state["health_state_visible"], "health_status": application_state["health_state_visible"],
"num_executions": self._categorise_num_executions(app_state["num_executions"]), "num_executions": application_state["num_executions"],
} }
@property @property
def space(self) -> spaces.Space: def space(self) -> spaces.Space:
"""Gymnasium space object describing the observation space shape.""" """
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for application status.
:rtype: spaces.Space
"""
return spaces.Dict( return spaces.Dict(
{ {
"operating_status": spaces.Discrete(7), "operating_status": spaces.Discrete(7),
"health_status": spaces.Discrete(6), "health_status": spaces.Discrete(5),
"num_executions": spaces.Discrete(4), "num_executions": spaces.Discrete(4),
} }
) )
@classmethod @classmethod
def from_config( def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ApplicationObservation:
cls, config: Dict, game: "PrimaiteGame", parent_where: Optional[List[str]] = None """
) -> "ApplicationObservation": Create an application observation from a configuration schema.
"""Create application observation from a config.
:param config: Dictionary containing the configuration for this service observation. :param config: Configuration schema containing the necessary information for the application observation.
:type config: Dict :type config: ConfigSchema
:param game: Reference to the PrimaiteGame object that spawned this observation. :param parent_where: Where in the simulation state dictionary to find the information about this application's
:type game: PrimaiteGame parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:param parent_where: Where in the simulation state dictionary this service's parent node is located. Optional. :type parent_where: WhereType, optional
:type parent_where: Optional[List[str]], optional :return: Constructed application observation instance.
:return: Constructed service observation
:rtype: ApplicationObservation :rtype: ApplicationObservation
""" """
return cls(where=parent_where + ["services", config["application_name"]]) return cls(where=parent_where + ["applications", config.application_name])
@classmethod
def _categorise_num_executions(cls, num_executions: int) -> int:
"""
Categorise the number of executions of an application.
Helps classify the number of application executions into different categories.
Current categories:
- 0: Application is never executed
- 1: Application is executed a low number of times (1-5)
- 2: Application is executed often (6-10)
- 3: Application is executed a high number of times (more than 10)
:param: num_executions: Number of times the application is executed
"""
if num_executions > 10:
return 3
elif num_executions > 5:
return 2
elif num_executions > 0:
return 1
return 0