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

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

@@ -4,8 +4,9 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
from gym import spaces
from primaite.simulator.sim_container import Simulation
from primaite import getLogger
from primaite.simulator.sim_container import Simulation
_LOGGER = getLogger(__name__)
if TYPE_CHECKING:
@@ -394,6 +395,7 @@ class NetworkNICDisableAction(NetworkNICAbstractAction):
super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs)
self.verb = "disable"
# class NetworkNICDisableAction(AbstractAction):
# def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
# super().__init__(manager=manager)
@@ -507,8 +509,8 @@ class ActionManager:
# option_2: value2
# where `type` decides which AbstractAction subclass should be used
# and `options` is an optional dict of options to pass to the init method of the action class
act_type = act_spec.get('type')
act_options = act_spec.get('options', {})
act_type = act_spec.get("type")
act_options = act_spec.get("options", {})
self.actions[act_type] = self.__act_class_identifiers[act_type](self, **global_action_args, **act_options)
self.action_map: Dict[int, Tuple[str, Dict]] = {}
@@ -555,14 +557,13 @@ class ActionManager:
param_combinations = list(itertools.product(*possibilities))
all_action_possibilities.extend(
[
(
act_name, {param_names[i]:param_combinations[j][i] for i in range(len(param_names))}
) for j in range(len(param_combinations))]
(act_name, {param_names[i]: param_combinations[j][i] for i in range(len(param_names))})
for j in range(len(param_combinations))
]
)
return {i: p for i, p in enumerate(all_action_possibilities)}
def get_action(self, action: int) -> Tuple[str, Dict]:
"""Produce action in CAOS format"""
"""the agent chooses an action (as an integer), this is converted into an action in CAOS format"""
@@ -627,7 +628,7 @@ class ActionManager:
session=session,
actions=cfg["action_list"],
# node_uuids=cfg["options"]["node_uuids"],
**cfg['options'],
**cfg["options"],
protocols=session.options.protocols,
ports=session.options.ports,
ip_address_list=None,

99
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
@@ -133,10 +134,7 @@ class ServiceObservation(AbstractObservation):
@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]
)
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):
@@ -238,7 +236,7 @@ class FolderObservation(AbstractObservation):
@classmethod
def from_config(cls, config: Dict, session: "PrimaiteSession", parent_where: Optional[List[str]]):
where = parent_where + ["folders", config['folder_name']]
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]
@@ -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.
@@ -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', {})
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,7 +391,9 @@ 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:{
"RULES": {
i
+ 1: {
"position": i,
"permission": 0,
"source_node_id": 0,
@@ -406,7 +413,6 @@ 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
obs = {}
obs["RULES"] = {}
@@ -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,19 +467,19 @@ 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):
@@ -490,7 +497,9 @@ class NullObservation(AbstractObservation):
def from_config(cls, config: Dict, session: Optional["PrimaiteSession"] = None) -> "NullObservation":
return cls()
class ICSObservation(NullObservation): pass
class ICSObservation(NullObservation):
pass
class UC2BlueObservation(AbstractObservation):
@@ -522,21 +531,23 @@ class UC2BlueObservation(AbstractObservation):
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({
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"):
@@ -551,7 +562,7 @@ 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
@@ -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({
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:
"""
@@ -615,11 +630,11 @@ class ObservationSpace:
@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":
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")

18
src/primaite/game/agent/rewards.py
View File

@@ -1,7 +1,8 @@
from abc import ABC, abstractmethod
from typing import Any, Dict, List
class AbstractReward():
class AbstractReward:
def __init__(self):
...
@@ -9,15 +10,15 @@ class AbstractReward():
def calculate(self, state: Dict) -> float:
return 0.3
class DummyReward(AbstractReward):
class DummyReward(AbstractReward):
def calculate(self, state: Dict) -> float:
return -0.1
class RewardFunction():
__rew_class_identifiers:Dict[str,type[AbstractReward]] = {
"DUMMY" : DummyReward
}
class RewardFunction:
__rew_class_identifiers: Dict[str, type[AbstractReward]] = {"DUMMY": DummyReward}
def __init__(self, reward_function: AbstractReward):
self.reward: AbstractReward = reward_function
@@ -26,9 +27,8 @@ class RewardFunction():
@classmethod
def from_config(cls, cfg: Dict) -> "RewardFunction":
for rew_component_cfg in cfg['reward_components']:
rew_type = rew_component_cfg['type']
for rew_component_cfg in cfg["reward_components"]:
rew_type = rew_component_cfg["type"]
rew_component = cls.__rew_class_identifiers[rew_type]()
new = cls(reward_function=rew_component)
return new

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

@@ -10,6 +10,7 @@ from typing import Dict, List
from pydantic import BaseModel
from primaite import getLogger
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.interface import AbstractAgent, RandomAgent
from primaite.game.agent.observations import (
@@ -37,14 +38,12 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.sim_container import Simulation
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.database_service import DatabaseService
from primaite.simulator.system.services.dns_client import DNSClient
from primaite.simulator.system.services.dns_server import DNSServer
from primaite.simulator.system.services.database.database_service import DatabaseService
from primaite.simulator.system.services.dns.dns_client import DNSClient
from primaite.simulator.system.services.dns.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__)
@@ -106,8 +105,8 @@ class PrimaiteSession:
def from_config(cls, cfg: dict) -> "PrimaiteSession":
sess = cls()
sess.options = PrimaiteSessionOptions(
ports = cfg['game_config']['ports'],
protocols = cfg['game_config']['protocols'],
ports=cfg["game_config"]["ports"],
protocols=cfg["game_config"]["protocols"],
)
sim = sess.simulation
net = sim.network
@@ -331,23 +330,23 @@ class PrimaiteSession:
"""
# CREATE ACTION SPACE
action_space_cfg['options']['node_uuids'] = []
action_space_cfg["options"]["node_uuids"] = []
# if a list of nodes is defined, convert them from node references to node UUIDs
for action_node_option in action_space_cfg.get('options',{}).pop('nodes', {}):
if 'node_ref' in action_node_option:
node_uuid = sess.ref_map_nodes[action_node_option['node_ref']]
action_space_cfg['options']['node_uuids'].append(node_uuid)
for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}):
if "node_ref" in action_node_option:
node_uuid = sess.ref_map_nodes[action_node_option["node_ref"]]
action_space_cfg["options"]["node_uuids"].append(node_uuid)
# Each action space can potentially have a different list of nodes that it can apply to. Therefore,
# we will pass node_uuids as a part of the action space config.
# However, it's not possible to specify the node uuids directly in the config, as they are generated
# dynamically, so we have to translate node references to uuids before passing this config on.
if 'action_list' in action_space_cfg:
for action_config in action_space_cfg['action_list']:
if 'options' in action_config:
if 'target_router_ref' in action_config['options']:
_target = action_config['options']['target_router_ref']
action_config['options']['target_router_uuid'] = sess.ref_map_nodes[_target]
if "action_list" in action_space_cfg:
for action_config in action_space_cfg["action_list"]:
if "options" in action_config:
if "target_router_ref" in action_config["options"]:
_target = action_config["options"]["target_router_ref"]
action_config["options"]["target_router_uuid"] = sess.ref_map_nodes[_target]
action_space = ActionManager.from_config(sess, action_space_cfg)
@@ -357,16 +356,30 @@ class PrimaiteSession:
# CREATE AGENT
if agent_type == "GreenWebBrowsingAgent":
# 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)
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)
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)
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")
return sess