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

Change observations to make loading from config better

Browse Source
This commit is contained in:
Marek Wolan
2023-10-08 17:02:54 +01:00
parent 3dea9743c3
commit ccb36f8400
5 changed files with 246 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -18,10 +18,12 @@ class AbstractAgent(ABC):
def __init__(
self,
agent_name: Optional[str],
action_space: Optional[ActionManager],
observation_space: Optional[ObservationSpace],
reward_function: Optional[RewardFunction],
) -> None:
self.agent_name:str = agent_name or "unnamed_agent"
self.action_space: Optional[ActionManager] = action_space
self.observation_space: Optional[ObservationSpace] = observation_space
self.reward_function: Optional[RewardFunction] = reward_function

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

@@ -1,10 +1,13 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, Hashable, List, Optional
from typing import Any, Dict, Hashable, List, Optional, TYPE_CHECKING
from gym import spaces
from pydantic import BaseModel
from primaite.game.session import PrimaiteSession
from primaite.simulator.sim_container import Simulation
if TYPE_CHECKING:
from primaite.game.session import PrimaiteSession
NOT_PRESENT_IN_STATE = object()
"""
@@ -53,6 +56,15 @@ class AbstractObservation(ABC):
"""Subclasses must define the shape that they expect"""
...
@abstractmethod
@classmethod
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,
a subclass of this class may need to translate from a 'reference' to a UUID.
"""
class FileObservation(AbstractObservation):
def __init__(self, where: Optional[List[str]] = None) -> None:
@@ -84,6 +96,10 @@ class FileObservation(AbstractObservation):
def space(self) -> spaces.Space:
return spaces.Dict({"health_status": spaces.Discrete(6)})
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where=None):
return cls(where=parent_where+["files", config["file_name"]])
class ServiceObservation(AbstractObservation):
default_observation: spaces.Space = {"operating_status": 0, "health_status": 0}
@@ -115,6 +131,11 @@ class ServiceObservation(AbstractObservation):
def space(self) -> spaces.Space:
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']]])
class LinkObservation(AbstractObservation):
default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}}
@@ -154,6 +175,10 @@ class LinkObservation(AbstractObservation):
def space(self) -> spaces.Space:
return spaces.Dict({"protocols": spaces.Dict({"all": spaces.Dict({"load": spaces.Discrete(11)})})})
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession"):
return cls(where=['network','links', session.ref_map_links[config['link_ref']]])
class FolderObservation(AbstractObservation):
def __init__(self, where: Optional[List[str]] = None, files: List[FileObservation] = []) -> None:
@@ -209,6 +234,15 @@ class FolderObservation(AbstractObservation):
}
)
@classmethod
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)
class NicObservation(AbstractObservation):
default_observation: spaces.Space = {"nic_status": 0}
@@ -230,6 +264,10 @@ class NicObservation(AbstractObservation):
def space(self) -> spaces.Space:
return spaces.Dict({"nic_status": spaces.Discrete(3)})
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]):
return cls(where=parent_where + ["NICs", config["nic_uuid"]])
class NodeObservation(AbstractObservation):
def __init__(
@@ -310,6 +348,25 @@ class NodeObservation(AbstractObservation):
return spaces.Dict(space_shape)
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where:Optional[List[str]]= None):
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', {})
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]
nic_uuids = session.simulation.network.nodes[node_uuid].nics.keys()
nic_configs = [{'nic_uuid':n for n in nic_uuids }]
nics = [NicObservation.from_config(config=c, session=session, parent_where=where) for c in nic_configs]
logon_status = config.get('logon_status',False)
cls(where=where, services=services, folders=folders, nics=nics, logon_status=logon_status)
return super().from_config(config, session)
class AclObservation(AbstractObservation):
@@ -399,6 +456,21 @@ 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']):
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']]
return cls(
node_ip_to_id=node_ip_to_idx,
ports=session.options.ports,
protocols=session.options.protocols,
where=["network", "nodes", router_uuid])
@@ -413,6 +485,10 @@ class NullObservation(AbstractObservation):
def space(self) -> spaces.Space:
return spaces.Dict({})
@classmethod
def from_config(cls, cfg:Dict) -> "NullObservation":
return cls()
class ICSObservation(NullObservation): pass
@@ -463,11 +539,18 @@ class UC2BlueObservation(AbstractObservation):
})
@classmethod
def from_config(cls, config:Dict, sim:Simulation):
nodes = ...
links = ...
acl = ...
ics = ...
def from_config(cls, config:Dict, sess:"PrimaiteSession"):
node_configs = config["nodes"]
nodes = [NodeObservation.from_config(n) for n in node_configs]
link_configs = config["links"]
links = [LinkObservation.from_config(l) for l in link_configs]
acl_config = config["acl"]
acl = AclObservation.from_config(acl_config)
ics_config = config["ics"]
ics = ICSObservation.from_config(ics_config)
new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=['network'])
return new
@@ -489,8 +572,11 @@ class UC2RedObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, sim:Simulation):
... #TODO
class UC2GreenObservation(NullObservation): pass
class ObservationSpace:
"""
Manage the observations of an Actor.
@@ -515,3 +601,14 @@ class ObservationSpace:
@property
def space(self) -> None:
return self.obs.space
@classmethod
def from_config(cls, config:Dict, session:"PrimaiteSession") -> "ObservationSpace":
if config['type'] == "UC2BlueObservation":
return cls(UC2BlueObservation(config['options']))
elif config['type'] == "UC2RedObservation":
return cls(UC2RedObservation(config['options']))
elif config['type'] == "UC2GreenObservation":
return cls(UC2GreenObservation(config["options"]))
else:
raise ValueError("Observation space type invalid")

32
src/primaite/game/session.py
View File

@@ -23,6 +23,7 @@ from primaite.game.agent.observations import (
NullObservation,
ServiceObservation,
UC2BlueObservation,
UC2GreenObservation,
UC2RedObservation,
)
from primaite.game.agent.rewards import RewardFunction
@@ -41,6 +42,10 @@ from primaite.simulator.system.services.dns_server import DNSServer
from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
from primaite.simulator.system.services.service import Service
from primaite import getLogger
_LOGGER = getLogger(__name__)
class PrimaiteSessionOptions(BaseModel):
ports: List[str]
@@ -55,13 +60,19 @@ class PrimaiteSession:
self.episode_counter: int = 0
self.options: PrimaiteSessionOptions
self.ref_map_nodes: Dict[str, Node] = {}
self.ref_map_services: Dict[str, Service] = {}
self.ref_map_links: Dict[str, Link] = {}
def step(self):
_LOGGER.debug(f"Stepping primaite session. Step counter: {self.step_counter}")
# currently designed with assumption that all agents act once per step in order
for agent in self.agents:
# 3. primaite session asks simulation to provide initial state
# 4. primate session gives state to all agents
# 5. primaite session asks agents to produce an action based on most recent state
_LOGGER.debug(f"Sending simulation state to agent {agent.agent_name}")
sim_state = self.simulation.describe_state()
# 6. each agent takes most recent state and converts it to CAOS observation
@@ -75,14 +86,18 @@ class PrimaiteSession:
# to discrete(40) is only necessary for purposes of RL learning, therefore that bit of
# code should live inside of the GATE agent subclass)
# gets action in CAOS format
_LOGGER.debug(f"Getting agent action")
agent_action, action_options = agent.get_action(agent_obs, agent_reward)
# 9. CAOS action is converted into request (extra information might be needed to enrich
# the request, this is what the execution definition is there for)
_LOGGER.debug(f"Formatting agent action {agent_action}") # maybe too many debug log statements
agent_request = agent.format_request(agent_action, action_options)
# 10. primaite session receives the action from the agents and asks the simulation to apply each
_LOGGER.debug(f"Sending request to simulation: {agent_request}")
self.simulation.apply_action(agent_request)
_LOGGER.debug(f"Initiating simulation step {self.step_counter}")
self.simulation.apply_timestep(self.step_counter)
self.step_counter += 1
@@ -96,9 +111,9 @@ class PrimaiteSession:
sim = sess.simulation
net = sim.network
ref_map_nodes: Dict[str, Node] = {}
ref_map_services: Dict[str, Service] = {}
ref_map_links: Dict[str, Link] = {}
sess.ref_map_nodes: Dict[str, Node] = {}
sess.ref_map_services: Dict[str, Service] = {}
sess.ref_map_links: Dict[str, Link] = {}
nodes_cfg = cfg["simulation"]["network"]["nodes"]
links_cfg = cfg["simulation"]["network"]["links"]
@@ -304,6 +319,8 @@ class PrimaiteSession:
)
elif observation_space_cfg["type"] == "UC2RedObservation":
obs_space = UC2RedObservation.from_config(observation_space_cfg["options"], sim=sim)
elif observation_space_cfg["type"] == "UC2GreenObservation":
obs_space = UC2GreenObservation.from_config(observation_space_cfg.get('options',{}))
else:
print("observation space config not specified correctly.")
obs_space = NullObservation()
@@ -334,12 +351,15 @@ class PrimaiteSession:
# CREATE AGENT
if agent_type == "GreenWebBrowsingAgent":
new_agent = RandomAgent(action_space=action_space, observation_space=obs_space, reward_function=rew_function)
# TODO: implement non-random agents and fix this parsing
new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function)
sess.agents.append(new_agent)
elif agent_type == "GATERLAgent":
...
new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function)
sess.agents.append(new_agent)
elif agent_type == "RedDatabaseCorruptingAgent":
...
new_agent = RandomAgent(agent_name=agent_cfg['ref'], action_space=action_space, observation_space=obs_space, reward_function=rew_function)
sess.agents.append(new_agent)
else:
print("agent type not found")