New observations
This commit is contained in:
Marek Wolan
@@ -1,21 +1,473 @@
|
||||
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
|
||||
from __future__ import annotations
|
||||
from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, TYPE_CHECKING, Union
|
||||
|
||||
from gymnasium import spaces
|
||||
from gymnasium.core import ObsType
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
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
|
||||
from primaite.game.agent.observations.software_observation import ServiceObservation
|
||||
# from primaite.game.agent.observations.file_system_observations import FolderObservation
|
||||
# from primaite.game.agent.observations.nic_observations import NicObservation
|
||||
# from primaite.game.agent.observations.software_observation import ServiceObservation
|
||||
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
|
||||
|
||||
_LOGGER = getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from primaite.game.game import PrimaiteGame
|
||||
WhereType = Iterable[str | int] | None
|
||||
|
||||
|
||||
class NodeObservation(AbstractObservation):
|
||||
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
service_name: str
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
self.where = where
|
||||
self.default_observation = {"operating_status": 0, "health_status": 0}
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
service_state = access_from_nested_dict(state, self.where)
|
||||
if service_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
return {
|
||||
"operating_status": service_state["operating_state"],
|
||||
"health_status": service_state["health_state_visible"],
|
||||
}
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
"""Gymnasium space object describing the observation space shape."""
|
||||
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
return cls(where=parent_where+["services", config.service_name])
|
||||
|
||||
|
||||
class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
application_name: str
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
self.where = where
|
||||
self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
# raise NotImplementedError("TODO NUM EXECUTIONS NEEDS TO BE CONVERTED TO A CATEGORICAL")
|
||||
application_state = access_from_nested_dict(state, self.where)
|
||||
if application_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
return {
|
||||
"operating_status": application_state["operating_state"],
|
||||
"health_status": application_state["health_state_visible"],
|
||||
"num_executions": application_state["num_executions"],
|
||||
}
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
"""Gymnasium space object describing the observation space shape."""
|
||||
return spaces.Dict({
|
||||
"operating_status": spaces.Discrete(7),
|
||||
"health_status": spaces.Discrete(5),
|
||||
"num_executions": spaces.Discrete(4)
|
||||
})
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ApplicationObservation:
|
||||
return cls(where=parent_where+["applications", config.application_name])
|
||||
|
||||
|
||||
class FileObservation(AbstractObservation, identifier="FILE"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
file_name: str
|
||||
include_num_access : bool = False
|
||||
|
||||
def __init__(self, where: WhereType, include_num_access: bool)->None:
|
||||
self.where: WhereType = where
|
||||
self.include_num_access :bool = include_num_access
|
||||
|
||||
self.default_observation: ObsType = {"health_status": 0}
|
||||
if self.include_num_access:
|
||||
self.default_observation["num_access"] = 0
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
file_state = access_from_nested_dict(state, self.where)
|
||||
if file_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
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
|
||||
def space(self) -> spaces.Space:
|
||||
space = {"health_status": spaces.Discrete(6)}
|
||||
if self.include_num_access:
|
||||
space["num_access"] = spaces.Discrete(4)
|
||||
return spaces.Dict(space)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> FileObservation:
|
||||
return cls(where=parent_where+["files", config.file_name], include_num_access=config.include_num_access)
|
||||
|
||||
|
||||
class FolderObservation(AbstractObservation, identifier="FOLDER"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
folder_name: str
|
||||
files: List[FileObservation.ConfigSchema] = []
|
||||
num_files : int = 0
|
||||
include_num_access : bool = False
|
||||
|
||||
def __init__(self, where: WhereType, files: Iterable[FileObservation], num_files: int, include_num_access: bool)->None:
|
||||
self.where: WhereType = where
|
||||
|
||||
self.files: List[FileObservation] = files
|
||||
while len(self.files) < num_files:
|
||||
self.files.append(FileObservation(where=None,include_num_access=include_num_access))
|
||||
while len(self.files) > num_files:
|
||||
truncated_file = self.files.pop()
|
||||
msg = f"Too many files in folder observation. Truncating file {truncated_file}"
|
||||
_LOGGER.warning(msg)
|
||||
|
||||
self.default_observation = {
|
||||
"health_status": 0,
|
||||
"FILES": {i + 1: f.default_observation for i, f in enumerate(self.files)},
|
||||
}
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
folder_state = access_from_nested_dict(state, self.where)
|
||||
if folder_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
|
||||
health_status = folder_state["health_status"]
|
||||
|
||||
obs = {}
|
||||
|
||||
obs["health_status"] = health_status
|
||||
obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)}
|
||||
|
||||
return obs
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
"""Gymnasium space object describing the observation space shape.
|
||||
|
||||
:return: Gymnasium space
|
||||
:rtype: spaces.Space
|
||||
"""
|
||||
return spaces.Dict(
|
||||
{
|
||||
"health_status": spaces.Discrete(6),
|
||||
"FILES": spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)}),
|
||||
}
|
||||
)
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> FileObservation:
|
||||
where = parent_where + ["folders", config.folder_name]
|
||||
|
||||
#pass down shared/common config items
|
||||
for file_config in config.files:
|
||||
file_config.include_num_access = config.include_num_access
|
||||
|
||||
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)
|
||||
|
||||
|
||||
class NICObservation(AbstractObservation, identifier="NETWORK_INTERFACE"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
nic_num: int
|
||||
include_nmne: bool = False
|
||||
|
||||
|
||||
def __init__(self, where: WhereType, include_nmne: bool)->None:
|
||||
self.where = where
|
||||
self.include_nmne : bool = include_nmne
|
||||
|
||||
self.default_observation: ObsType = {"nic_status": 0}
|
||||
if self.include_nmne:
|
||||
self.default_observation.update({"NMNE":{"inbound":0, "outbound":0}})
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
# raise NotImplementedError("TODO: CATEGORISATION")
|
||||
nic_state = access_from_nested_dict(state, self.where)
|
||||
|
||||
if nic_state is NOT_PRESENT_IN_STATE:
|
||||
return self.default_observation
|
||||
|
||||
obs = {"nic_status": 1 if nic_state["enabled"] else 2}
|
||||
if self.include_nmne:
|
||||
obs.update({"NMNE": {}})
|
||||
direction_dict = nic_state["nmne"].get("direction", {})
|
||||
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
|
||||
inbound_count = inbound_keywords.get("*", 0)
|
||||
outbound_keywords = direction_dict.get("outbound", {}).get("keywords", {})
|
||||
outbound_count = outbound_keywords.get("*", 0)
|
||||
obs["NMNE"]["inbound"] = self._categorise_mne_count(inbound_count - self.nmne_inbound_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_outbound_last_step = outbound_count
|
||||
return obs
|
||||
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
space = spaces.Dict({"nic_status": spaces.Discrete(3)})
|
||||
|
||||
if self.include_nmne:
|
||||
space["NMNE"] = spaces.Dict({"inbound": spaces.Discrete(4), "outbound": spaces.Discrete(4)})
|
||||
|
||||
return space
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
return cls(where = parent_where+["NICs", config.nic_num], include_nmne=config.include_nmne)
|
||||
|
||||
|
||||
class HostObservation(AbstractObservation, identifier="HOST"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
hostname: str
|
||||
services: List[ServiceObservation.ConfigSchema] = []
|
||||
applications: List[ApplicationObservation.ConfigSchema] = []
|
||||
folders: List[FolderObservation.ConfigSchema] = []
|
||||
network_interfaces: List[NICObservation.ConfigSchema] = []
|
||||
num_services: int
|
||||
num_applications: int
|
||||
num_folders: int
|
||||
num_files: int
|
||||
num_nics: int
|
||||
include_nmne: bool
|
||||
include_num_access: bool
|
||||
|
||||
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:
|
||||
|
||||
self.where : WhereType = where
|
||||
|
||||
# ensure service list has length equal to num_services 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)
|
||||
|
||||
# ensure application list has length equal to num_applications by truncating or padding
|
||||
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 application {truncated_application.where}"
|
||||
_LOGGER.warning(msg)
|
||||
|
||||
# ensure folder list has length equal to num_folders by truncating or padding
|
||||
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)
|
||||
|
||||
# ensure network_interface list has length equal to num_network_interfaces by truncating or padding
|
||||
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_folder.where}"
|
||||
_LOGGER.warning(msg)
|
||||
|
||||
self.default_observation: ObsType = {
|
||||
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
|
||||
"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) -> Any:
|
||||
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["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:
|
||||
shape = {
|
||||
"SERVICES": spaces.Dict({i + 1: service.space for i, service in enumerate(self.services)}),
|
||||
"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 ) -> ServiceObservation:
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
class PortObservation(AbstractObservation, identifier="PORT"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
pass
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
pass
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
pass
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
pass
|
||||
|
||||
class ACLObservation(AbstractObservation, identifier="ACL"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
pass
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
pass
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
pass
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
pass
|
||||
|
||||
class RouterObservation(AbstractObservation, identifier="ROUTER"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
hostname: str
|
||||
ports: List[PortObservation.ConfigSchema]
|
||||
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
pass
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
pass
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
pass
|
||||
|
||||
class FirewallObservation(AbstractObservation, identifier="FIREWALL"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
hostname: str
|
||||
ports: List[PortObservation.ConfigSchema] = []
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
pass
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
pass
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
pass
|
||||
|
||||
class NodesObservation(AbstractObservation, identifier="NODES"):
|
||||
class ConfigSchema(AbstractObservation.ConfigSchema):
|
||||
"""Config"""
|
||||
hosts: List[HostObservation.ConfigSchema] = []
|
||||
routers: List[RouterObservation.ConfigSchema] = []
|
||||
firewalls: List[FirewallObservation.ConfigSchema] = []
|
||||
num_services: int = 1
|
||||
|
||||
|
||||
def __init__(self, where: WhereType)->None:
|
||||
pass
|
||||
|
||||
def observe(self, state: Dict) -> Any:
|
||||
pass
|
||||
|
||||
@property
|
||||
def space(self) -> spaces.Space:
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: ConfigSchema, parent_where: WhereType = [] ) -> ServiceObservation:
|
||||
pass
|
||||
|
||||
############################ OLD
|
||||
|
||||
class NodeObservation(AbstractObservation, identifier= "OLD"):
|
||||
"""Observation of a node in the network. Includes services, folders and NICs."""
|
||||
|
||||
def __init__(
|
||||
|
||||
Reference in New Issue
Block a user