Fix formatting with precommit

src/primaite/game/agent/actions.py

@@ -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_type = act_spec.get("type")
-            act_options = act_spec.get('options', {})
+            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))})
-                    act_name, {param_names[i]:param_combinations[j][i] for i in range(len(param_names))}
+                    for j in range(len(param_combinations))
-                ) 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,

src/primaite/game/agent/observations.py

@@ -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(
+        return cls(where=parent_where + ["services", session.ref_map_services[config["service_ref"]].uuid])
-            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":
+    def from_config(
-        node_uuid = session.ref_map_nodes[config['node_ref']]
+        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["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["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)}
-        obs['ACL'] = self.acl.observe(state)
+        obs["ACL"] = self.acl.observe(state)
-        obs['ICS'] = self.ics.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":
+        if config["type"] == "UC2BlueObservation":
-            return cls(UC2BlueObservation.from_config(config.get('options',{}), session=session))
+            return cls(UC2BlueObservation.from_config(config.get("options", {}), session=session))
-        elif config['type'] == "UC2RedObservation":
+        elif config["type"] == "UC2RedObservation":
-            return cls(UC2RedObservation.from_config(config.get('options',{}), session=session))
+            return cls(UC2RedObservation.from_config(config.get("options", {}), session=session))
-        elif config['type'] == "UC2GreenObservation":
+        elif config["type"] == "UC2GreenObservation":
             return cls(UC2GreenObservation.from_config(config.get("options", {}), session=session))
         else:
             raise ValueError("Observation space type invalid")

src/primaite/game/agent/rewards.py

@@ -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]] = {
+class RewardFunction:
-        "DUMMY" : DummyReward
+    __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']:
+        for rew_component_cfg in cfg["reward_components"]:
-            rew_type = rew_component_cfg['type']
+            rew_type = rew_component_cfg["type"]
             rew_component = cls.__rew_class_identifiers[rew_type]()
             new = cls(reward_function=rew_component)
         return new
-

src/primaite/game/session.py

@@ -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.database.database_service import DatabaseService
-from primaite.simulator.system.services.dns_client import DNSClient
+from primaite.simulator.system.services.dns.dns_client import DNSClient
-from primaite.simulator.system.services.dns_server import DNSServer
+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'],
+            ports=cfg["game_config"]["ports"],
-            protocols = cfg['game_config']['protocols'],
+            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', {}):
+            for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}):
-                if 'node_ref' in action_node_option:
+                if "node_ref" in action_node_option:
-                    node_uuid = sess.ref_map_nodes[action_node_option['node_ref']]
+                    node_uuid = sess.ref_map_nodes[action_node_option["node_ref"]]
-                    action_space_cfg['options']['node_uuids'].append(node_uuid)
+                    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:
+            if "action_list" in action_space_cfg:
-                for action_config in action_space_cfg['action_list']:
+                for action_config in action_space_cfg["action_list"]:
-                    if 'options' in action_config:
+                    if "options" in action_config:
-                        if 'target_router_ref' in action_config['options']:
+                        if "target_router_ref" in action_config["options"]:
-                            _target = action_config['options']['target_router_ref']
+                            _target = action_config["options"]["target_router_ref"]
-                            action_config['options']['target_router_uuid'] = sess.ref_map_nodes[_target]
+                            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