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

Fix formatting with precommit

Browse Source
This commit is contained in:
Marek Wolan
2023-10-09 18:35:30 +01:00
parent c9bc8fbf3d
commit 91f06c15f6
6 changed files with 189 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,10 +1,11 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING, Sequence, Tuple
from typing import Any, Dict, Hashable, List, Optional, Sequence, Tuple, TYPE_CHECKING
from gym import spaces
from pydantic import BaseModel
from primaite.simulator.sim_container import Simulation
if TYPE_CHECKING:
from primaite.game.session import PrimaiteSession
@@ -29,7 +30,7 @@ def access_from_nested_dict(dictionary: Dict, keys: Sequence[Hashable]) -> Any:
:return: The value in the dictionary
:rtype: Any
"""
key_list = [*keys] # copy keys to a new list to prevent editing original list
key_list = [*keys] # copy keys to a new list to prevent editing original list
if len(key_list) == 0:
return dictionary
k = key_list.pop(0)
@@ -58,7 +59,7 @@ class AbstractObservation(ABC):
@classmethod
@abstractmethod
def from_config(cls, config:Dict, session:"PrimaiteSession"):
def from_config(cls, config: Dict, session: "PrimaiteSession"):
"""Create this observation space component form a serialised format.
The `session` parameter is for a the PrimaiteSession object that spawns this component. During deserialisation,
@@ -98,7 +99,7 @@ class FileObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None):
return cls(where=parent_where+["files", config["file_name"]])
return cls(where=parent_where + ["files", config["file_name"]])
class ServiceObservation(AbstractObservation):
@@ -132,11 +133,8 @@ class ServiceObservation(AbstractObservation):
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(6)})
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]=None):
return cls(
where=parent_where+["services",session.ref_map_services[config['service_ref']].uuid]
)
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None):
return cls(where=parent_where + ["services", session.ref_map_services[config["service_ref"]].uuid])
class LinkObservation(AbstractObservation):
@@ -179,7 +177,7 @@ class LinkObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession"):
return cls(where=['network','links', session.ref_map_links[config['link_ref']]])
return cls(where=["network", "links", session.ref_map_links[config["link_ref"]]])
class FolderObservation(AbstractObservation):
@@ -237,13 +235,13 @@ class FolderObservation(AbstractObservation):
)
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]):
where = parent_where + ["folders", config['folder_name']]
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]):
where = parent_where + ["folders", config["folder_name"]]
file_configs = config["files"]
files = [FileObservation.from_config(config=f, session=session, parent_where=where) for f in file_configs]
return cls(where=where,files=files)
return cls(where=where, files=files)
class NicObservation(AbstractObservation):
@@ -267,7 +265,7 @@ class NicObservation(AbstractObservation):
return spaces.Dict({"nic_status": spaces.Discrete(3)})
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]):
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]):
return cls(where=parent_where + ["NICs", config["nic_uuid"]])
@@ -278,7 +276,7 @@ class NodeObservation(AbstractObservation):
services: List[ServiceObservation] = [],
folders: List[FolderObservation] = [],
nics: List[NicObservation] = [],
logon_status:bool=False
logon_status: bool = False,
) -> None:
"""
Configurable observation for a node in the simulation.
@@ -306,7 +304,7 @@ class NodeObservation(AbstractObservation):
self.services: List[ServiceObservation] = services
self.folders: List[FolderObservation] = folders
self.nics: List[NicObservation] = nics
self.logon_status:bool=logon_status
self.logon_status: bool = logon_status
self.default_observation: Dict = {
"SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
@@ -315,7 +313,7 @@ class NodeObservation(AbstractObservation):
"operating_status": 0,
}
if self.logon_status:
self.default_observation['logon_status']=0
self.default_observation["logon_status"] = 0
def observe(self, state: Dict) -> Dict:
if self.where is None:
@@ -332,7 +330,7 @@ class NodeObservation(AbstractObservation):
obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)}
if self.logon_status:
obs['logon_status'] = 0
obs["logon_status"] = 0
return obs
@@ -345,26 +343,28 @@ class NodeObservation(AbstractObservation):
"NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}),
}
if self.logon_status:
space_shape['logon_status'] = spaces.Discrete(3)
space_shape["logon_status"] = spaces.Discrete(3)
return spaces.Dict(space_shape)
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None) -> "NodeObservation":
node_uuid = session.ref_map_nodes[config['node_ref']]
def from_config(
cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]] = None
) -> "NodeObservation":
node_uuid = session.ref_map_nodes[config["node_ref"]]
if parent_where is None:
where = ["network", "nodes", node_uuid]
else:
where = parent_where + ["nodes", node_uuid]
svc_configs = config.get('services', {})
svc_configs = config.get("services", {})
services = [ServiceObservation.from_config(config=c, session=session, parent_where=where) for c in svc_configs]
folder_configs = config.get('folders', {})
folders = [FolderObservation.from_config(config=c,session=session, parent_where=where) for c in folder_configs]
folder_configs = config.get("folders", {})
folders = [FolderObservation.from_config(config=c, session=session, parent_where=where) for c in folder_configs]
nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys()
nic_configs = [{'nic_uuid':n for n in nic_uuids }] if nic_uuids else []
nic_configs = [{"nic_uuid": n for n in nic_uuids}] if nic_uuids else []
nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs]
logon_status = config.get('logon_status',False)
logon_status = config.get("logon_status", False)
return cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status)
@@ -374,7 +374,12 @@ class AclObservation(AbstractObservation):
# 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.
def __init__(
self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[Tuple[str]] = None, num_rules: int = 10
self,
node_ip_to_id: Dict[str, int],
ports: List[int],
protocols: list[str],
where: Optional[Tuple[str]] = None,
num_rules: int = 10,
) -> None:
super().__init__()
self.where: Optional[Tuple[str]] = where
@@ -386,16 +391,18 @@ class AclObservation(AbstractObservation):
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"
self.default_observation: Dict = {
"RULES": {i+ 1:{
"position": i,
"permission": 0,
"source_node_id": 0,
"source_port": 0,
"dest_node_id": 0,
"dest_port": 0,
"protocol": 0,
"RULES": {
i
+ 1: {
"position": i,
"permission": 0,
"source_node_id": 0,
"source_port": 0,
"dest_node_id": 0,
"dest_port": 0,
"protocol": 0,
}
for i in range(self.num_rules)
for i in range(self.num_rules)
}
}
@@ -406,8 +413,7 @@ class AclObservation(AbstractObservation):
if acl_state is NOT_PRESENT_IN_STATE:
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["RULES"] = {}
for i, rule_state in acl_state.items():
@@ -439,7 +445,8 @@ class AclObservation(AbstractObservation):
{
"RULE": spaces.Dict(
{
i + 1: spaces.Dict(
i
+ 1: spaces.Dict(
{
"position": spaces.Discrete(self.num_rules),
"permission": spaces.Discrete(3),
@@ -460,23 +467,23 @@ class AclObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "AclObservation":
node_ip_to_idx = {}
for node_idx, node_cfg in enumerate(config['node_order']):
for node_idx, node_cfg in enumerate(config["node_order"]):
n_ref = node_cfg["node_ref"]
n_obj = session.simulation.network.nodes[session.ref_map_nodes[n_ref]]
for nic_uuid, nic_obj in n_obj.nics.items():
node_ip_to_idx[nic_obj.ip_address] = node_idx + 2
router_uuid = session.ref_map_nodes[config['router_node_ref']]
router_uuid = session.ref_map_nodes[config["router_node_ref"]]
return cls(
node_ip_to_id=node_ip_to_idx,
ports=session.options.ports,
protocols=session.options.protocols,
where=["network", "nodes", router_uuid, "acl", "acl"])
where=["network", "nodes", router_uuid, "acl", "acl"],
)
class NullObservation(AbstractObservation):
def __init__(self, where:Optional[List[str]]=None):
def __init__(self, where: Optional[List[str]] = None):
self.default_observation: Dict = {}
def observe(self, state: Dict) -> Dict:
@@ -487,20 +494,22 @@ class NullObservation(AbstractObservation):
return spaces.Dict({})
@classmethod
def from_config(cls, config:Dict, session:Optional["PrimaiteSession"]=None) -> "NullObservation":
def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation":
return cls()
class ICSObservation(NullObservation): pass
class ICSObservation(NullObservation):
pass
class UC2BlueObservation(AbstractObservation):
def __init__(
self,
nodes: List[NodeObservation],
links: List[LinkObservation],
acl: AclObservation,
ics: ICSObservation,
where:Optional[List[str]] = None,
self,
nodes: List[NodeObservation],
links: List[LinkObservation],
acl: AclObservation,
ics: ICSObservation,
where: Optional[List[str]] = None,
) -> None:
super().__init__()
self.where: Optional[Tuple[str]] = where
@@ -510,36 +519,38 @@ class UC2BlueObservation(AbstractObservation):
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)},
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:
def observe(self, state: Dict) -> 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)
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:
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,
})
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, session:"PrimaiteSession"):
def from_config(cls, config: Dict, session: "PrimaiteSession"):
node_configs = config["nodes"]
nodes = [NodeObservation.from_config(config=n, session=session) for n in node_configs]
@@ -551,18 +562,18 @@ class UC2BlueObservation(AbstractObservation):
ics_config = config["ics"]
ics = ICSObservation.from_config(config=ics_config, session=session)
new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network'])
new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"])
return new
class UC2RedObservation(AbstractObservation):
def __init__(self, nodes:List[NodeObservation], where:Optional[List[str]] = None) -> None:
def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None:
super().__init__()
self.where:Optional[List[str]] = where
self.where: Optional[List[str]] = where
self.nodes: List[NodeObservation] = nodes
self.default_observation : Dict = {
"NODES": {i+1: n.default_observation for i,n in enumerate(self.nodes)},
self.default_observation: Dict = {
"NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
}
def observe(self, state: Dict) -> Dict:
@@ -570,14 +581,16 @@ class UC2RedObservation(AbstractObservation):
return self.default_observation
obs = {}
obs['NODES'] = {i+1: node.observe(state) for i, node in enumerate(self.nodes)}
obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
return obs
@property
def space(self) -> spaces.Space:
return spaces.Dict({
"NODES": spaces.Dict({i+1: node.space for i, node in enumerate(self.nodes)}),
})
return spaces.Dict(
{
"NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
}
)
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession"):
@@ -586,7 +599,9 @@ class UC2RedObservation(AbstractObservation):
return cls(nodes=nodes, where=["network"])
class UC2GreenObservation(NullObservation): pass
class UC2GreenObservation(NullObservation):
pass
class ObservationSpace:
"""
@@ -603,7 +618,7 @@ class ObservationSpace:
# what this class does:
# keep a list of observations
# create observations for an actor from the config
def __init__(self, observation:AbstractObservation) -> None:
def __init__(self, observation: AbstractObservation) -> None:
self.obs: AbstractObservation = observation
def observe(self, state) -> Dict:
@@ -614,12 +629,12 @@ class ObservationSpace:
return self.obs.space
@classmethod
def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace":
if config['type'] == "UC2BlueObservation":
return cls(UC2BlueObservation.from_config(config.get('options',{}), session=session))
elif config['type'] == "UC2RedObservation":
return cls(UC2RedObservation.from_config(config.get('options',{}), session=session))
elif config['type'] == "UC2GreenObservation":
return cls(UC2GreenObservation.from_config(config.get("options",{}), session=session))
def from_config(cls, config: Dict, session: "PrimaiteSession") -> "ObservationSpace":
if config["type"] == "UC2BlueObservation":
return cls(UC2BlueObservation.from_config(config.get("options", {}), session=session))
elif config["type"] == "UC2RedObservation":
return cls(UC2RedObservation.from_config(config.get("options", {}), session=session))
elif config["type"] == "UC2GreenObservation":
return cls(UC2GreenObservation.from_config(config.get("options", {}), session=session))
else:
raise ValueError("Observation space type invalid")