Finalise actions interface

src/primaite/game/agent/actions.py

@@ -1,21 +1,374 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
-
-from pydantic import BaseModel
+from typing import Any, Dict, List, Optional, Tuple
+import itertools
 
 
-class AbstractAction(BaseModel):
+from primaite.simulator.sim_container import Simulation
+
+from gym import spaces
+
+class ExecutionDefiniton(ABC):
+    """
+    Converter from actions to simulator requests.
+
+    Allows adding extra data/context that defines in more detail what an action means.
+    """
+
+    """
+    Examples:
+    ('node', 'service', 'scan', 2, 0) means scan the first service on node index 2
+        -> ['network', 'nodes', <node-idx-2-uuid>, 'services', <svc-idx-0-uuid>, 'scan'w]
+    """
+    ...
+
+
+class AbstractAction(ABC):
+
     @abstractmethod
-    def __call__(self, action: Any) -> List[str]:
-        """_summary_
-
-        :param action: _description_
-        :type action: Any
-        :return: _description_
-        :rtype: List[str]
+    def __init__(self, manager:"ActionManager", **kwargs) -> None:
         """
+        Init method for action.
+
+        All action init functions should accept **kwargs as a way of ignoring extra arguments.
+
+        Since many parameters are defined for the action space as a whole (such as max files per folder, max services
+        per node), we need to pass those options to every action that gets created. To pervent verbosity, these
+        parameters are just broadcasted to all actions and the actions can pay attention to the ones that apply.
+        """
+        self.name:str = ""
+        """Human-readable action identifier used for printing, logging, and reporting."""
+        self.shape = (0,)
+        """Tuple describing number of options for each parameter of this action. Can be passed to
+        gym.spaces.MultiDiscrete to form a valid space."""
+        self.manager:ActionManager = manager
+
+
+    @abstractmethod
+    def form_request(self) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
+        return []
+
+
+class DoNothingAction(AbstractAction):
+    def __init__(self, manager:"ActionManager", **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.name = "DONOTHING"
+        self.shape = (1,)
+
+    def form_request(self) -> List[str]:
+        return ["do_nothing"]
+
+class NodeServiceAbstractAction(AbstractAction):
+    """
+    Base class for service actions.
+
+    Any action which applies to a service and uses node_id and service_id as its only two parameters can inherit from
+    this base class.
+    """
+    @abstractmethod
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Tuple[int] = (num_nodes, num_services)
+        self.verb:str
+
+    def form_request(self, node_id:int, service_id:int) -> List[str]:
+        node_uuid = self.manager.get_node_uuid_by_idx(node_id)
+        service_uuid = self.manager.get_service_uuid_by_idx(node_id, service_id)
+        if node_uuid is None or service_uuid is None:
+            return ["do_nothing"]
+        return ['network', 'node', node_uuid, 'services', service_uuid, self.verb]
+
+class NodeServiceScanAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "scan"
+
+class NodeServiceStopAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "stop"
+
+class NodeServiceStartAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "start"
+
+class NodeServicePauseAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "pause"
+
+class NodeServiceResumeAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "resume"
+
+class NodeServiceRestartAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "restart"
+
+class NodeServiceDisableAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "disable"
+
+class NodeServiceEnableAction(NodeServiceAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = "enable"
+
+
+
+class NodeFolderAbstractAction(AbstractAction):
+    @abstractmethod
+    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape = (num_nodes, num_folders)
+        self.verb: str
+
+    def form_request(self, node_id:int, folder_id:int) -> List[str]:
+        node_uuid = self.manager.get_node_uuid_by_idx(node_id)
+        folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id)
+        if node_uuid is None or folder_uuid is None:
+            return ["do_nothing"]
+        return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, self.verb]
+
+class NodeFolderScanAction(NodeFolderAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, **kwargs)
+        self.verb:str = "scan"
+
+class NodeFolderCheckhashAction(NodeFolderAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, **kwargs)
+        self.verb:str = "checkhash"
+
+class NodeFolderRepairAction(NodeFolderAbstractAction):
+    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, **kwargs)
+        self.verb:str = "repair"
+
+class NodeFolderRestoreAction(NodeFolderAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, **kwargs)
+        self.verb:str = "restore"
+
+
+class NodeFileAbstractAction(AbstractAction):
+    @abstractmethod
+    def __init__(self, manager:"ActionManager", num_nodes:int, num_folders:int, num_files:int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape:Tuple[int] = (num_nodes, num_folders, num_files)
+        self.verb:str
+
+    def form_request(self, node_id:int, folder_id:int, file_id:int) -> List[str]:
+        node_uuid = self.manager.get_node_uuid_by_idx(node_id)
+        folder_uuid = self.manager.get_folder_uuid_by_idx(node_idx=node_id, folder_idx=folder_id)
+        file_uuid = self.manager.get_file_uuid_by_idx(node_idx=node_id, folder_idx=folder_id, file_idx=file_id)
+        if node_uuid is None or folder_uuid is None or file_uuid is None:
+            return ["do_nothing"]
+        return ['network', 'node', node_uuid, 'file_system', 'folder', folder_uuid, 'files', file_uuid, self.verb]
+
+class NodeFileScanAction(NodeFileAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, num_files, **kwargs)
+        self.verb = "scan"
+
+class NodeFileCheckhashAction(NodeFileAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, num_files, **kwargs)
+        self.verb = "checkhash"
+
+class NodeFileDeleteAction(NodeFileAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, num_files, **kwargs)
+        self.verb = "delete"
+
+class NodeFileRepairAction(NodeFileAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, num_files, **kwargs)
+        self.verb = "repair"
+
+class NodeFileRestoreAction(NodeFileAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
+        super().__init__(manager, num_nodes, num_folders, num_files, **kwargs)
+        self.verb = "restore"
+
+class NodeAbstractAction(AbstractAction):
+    @abstractmethod
+    def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Tuple[int] = (num_nodes,)
+        self.verb: str
+
+    def form_request(self, node_id:int) -> List[str]:
+        node_uuid = self.manager.get_node_uuid_by_idx(node_id)
+        return ["network", "node", node_uuid, self.verb]
+
+class NodeOSScanAction(NodeAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = 'scan'
+
+class NodeShutdownAction(NodeAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = 'shutdown'
+
+class NodeStartupAction(NodeAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = 'start'
+
+class NodeResetAction(NodeAbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.verb = 'reset'
+
+class NetworkACLAddRuleAction(AbstractAction):
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
+        super().__init__(manager=manager)
+        num_permissions = 2
+        self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_nics, num_nics, num_ports, num_ports, num_protocols)
+
+
+
+
+
+
+class ActionManager:
+    # let the action manager handle the conversion of action spaces into a single discrete integer space.
+    #
+
+
+    # when action space is created, it will take subspaces and generate an action map by enumerating all possibilities,
+    # BUT, the action map can be provided in the config, in which case it will use that.
+
+    # action map is basically just a mapping between integer and CAOS action (incl. parameter values)
+    # for example the action map can be:
+    # 0: DONOTHING
+    # 1: NODE, FILE, SCAN, NODEID=2, FOLDERID=1, FILEID=0
+    # 2: ......
+    __act_class_identifiers:Dict[str,type] = {
+        "DONOTHING": DoNothingAction,
+        "NODE_SERVICE_SCAN": NodeServiceScanAction,
+        "NODE_SERVICE_STOP": NodeServiceStopAction,
+        # "NODE_SERVICE_START": NodeServiceStartAction,
+        # "NODE_SERVICE_PAUSE": NodeServicePauseAction,
+        # "NODE_SERVICE_RESUME": NodeServiceResumeAction,
+        # "NODE_SERVICE_RESTART": NodeServiceRestartAction,
+        # "NODE_SERVICE_DISABLE": NodeServiceDisableAction,
+        # "NODE_SERVICE_ENABLE": NodeServiceEnableAction,
+        # "NODE_FILE_SCAN": NodeFileScanAction,
+        # "NODE_FILE_CHECKHASH": NodeFileCheckhashAction,
+        # "NODE_FILE_DELETE": NodeFileDeleteAction,
+        # "NODE_FILE_REPAIR": NodeFileRepairAction,
+        # "NODE_FILE_RESTORE": NodeFileRestoreAction,
+        "NODE_FOLDER_SCAN": NodeFolderScanAction,
+        # "NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction,
+        # "NODE_FOLDER_REPAIR": NodeFolderRepairAction,
+        # "NODE_FOLDER_RESTORE": NodeFolderRestoreAction,
+        # "NODE_OS_SCAN": NodeOSScanAction,
+        # "NODE_SHUTDOWN": NodeShutdownAction,
+        # "NODE_STARTUP": NodeStartupAction,
+        # "NODE_RESET": NodeResetAction,
+        # "NETWORK_ACL_ADDRULE": NetworkACLAddRuleAction,
+        # "NETWORK_ACL_REMOVERULE": NetworkACLRemoveRuleAction,
+        # "NETWORK_NIC_ENABLE": NetworkNICEnable,
+        # "NETWORK_NIC_DISABLE": NetworkNICDisable,
+    }
+
+
+    def __init__(self,
+                 sim:Simulation,
+                 actions:List[str],
+                 node_uuids:List[str],
+                 max_folders_per_node:int = 2,
+                 max_files_per_folder:int = 2,
+                 max_services_per_node:int = 2,
+                 max_nics_per_node:int=8,
+                 max_acl_rules:int=10,
+                 act_map:Optional[Dict[int, Dict]]=None) -> None:
+        self.sim: Simulation = sim
+        self.node_uuids:List[str] = node_uuids
+
+        action_args = {
+            "num_nodes": len(node_uuids),
+            "num_folders":max_folders_per_node,
+            "num_files": max_files_per_folder,
+            "num_services": max_services_per_node,
+            "num_nics": max_nics_per_node,
+            "num_acl_rules": max_acl_rules}
+        self.actions: Dict[str, AbstractAction] = {}
+        for act_type in actions:
+            self.actions[act_type] = self.__act_class_identifiers[act_type](self, **action_args)
+
+        self.action_map:Dict[int, Tuple[str, Dict]] = {}
+        """
+        Action mapping that converts an integer to a specific action and parameter choice.
+
+        For example :
+        {0: ("NODE_SERVICE_SCAN", {node_id:0, service_id:2})}
+        """
+        if act_map is None:
+            self.action_map = self._enumerate_actions()
+        else:
+            self.action_map = {i:(a['action'], a['options']) for i,a in act_map.items()}
+        # make sure all numbers between 0 and N are represented as dict keys in action map
+        assert all([i in self.action_map.keys() for i in range(len(self.action_map))])
+
+    def _enumerate_actions(self,) -> Dict[int, Tuple[AbstractAction, Dict]]:
         ...
 
+    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"""
+        """The caos format is basically a action identifier, followed by parameters stored in a dictionary"""
+        act_identifier, act_options = self.action_map[action]
+        return act_identifier, act_options
 
-class ActionSpace:
+    def form_request(self, action_identifier:str, action_options:Dict):
+        """Take action in CAOS format and use the execution definition to change it into PrimAITE request format"""
+        act_obj = self.actions[action_identifier]
+        return act_obj.form_request(**action_options)
+
+    @property
+    def space(self) -> spaces.Space:
+        return spaces.Discrete(len(self.action_map))
+
+    def get_node_uuid_by_idx(self, node_idx):
+        return self.node_uuids[node_idx]
+
+    def get_folder_uuid_by_idx(self, node_idx, folder_idx) -> Optional[str]:
+        node_uuid = self.get_node_uuid_by_idx(node_idx)
+        node = self.sim.network.nodes[node_uuid]
+        folder_uuids = list(node.file_system.folders.keys())
+        return folder_uuids[folder_idx] if len(folder_uuids)>folder_idx else None
+
+    def get_file_uuid_by_idx(self, node_idx, folder_idx, file_idx) -> Optional[str]:
+        node_uuid = self.get_node_uuid_by_idx(node_idx)
+        node = self.sim.network.nodes[node_uuid]
+        folder_uuids = list(node.file_system.folders.keys())
+        if len(folder_uuids)<=folder_idx:
+            return None
+        folder = node.file_system.folders[folder_uuids[folder_idx]]
+        file_uuids = list(folder.files.keys())
+        return file_uuids[file_idx] if len(file_uuids)>file_idx else None
+
+    def get_service_uuid_by_idx(self, node_idx, service_idx) -> Optional[str]:
+        node_uuid = self.get_node_uuid_by_idx(node_idx)
+        node = self.sim.network.nodes[node_uuid]
+        service_uuids = list(node.services.keys())
+        return service_uuids[service_idx] if len(service_uuids)>service_idx else None
+
+
+
+
+
+
+class UC2RedActions(AbstractAction):
+    ...
+
+class UC2GreenActionSpace(ActionManager):
     ...

src/primaite/game/agent/interface.py

@@ -2,32 +2,70 @@
 # That's because I want to point out that this is disctinct from 'agent' in the reinforcement learning sense of the word
 # If you disagree, make a comment in the PR review and we can discuss
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union, TypeAlias
+import numpy as np
 
-from primaite.game.agent.actions import ActionSpace
+from primaite.game.agent.actions import ActionManager
 from primaite.game.agent.observations import ObservationSpace
 from primaite.game.agent.rewards import RewardFunction
 
+ObsType:TypeAlias = Union[Dict, np.ndarray]
 
 class AbstractAgent(ABC):
     """Base class for scripted and RL agents."""
 
     def __init__(
         self,
-        action_space: Optional[ActionSpace],
+        action_space: Optional[ActionManager],
         observation_space: Optional[ObservationSpace],
         reward_function: Optional[RewardFunction],
     ) -> None:
-        self.action_space: Optional[ActionSpace] = action_space
+        self.action_space: Optional[ActionManager] = action_space
         self.observation_space: Optional[ObservationSpace] = observation_space
         self.reward_function: Optional[RewardFunction] = reward_function
 
+        # exection definiton converts CAOS action to Primaite simulator request, sometimes having to enrich the info
+        # by for example specifying target ip addresses, or converting a node ID into a uuid
+        self.execution_definition = None
+
+    def get_obs_from_state(self, state:Dict) -> ObsType:
+        """
+        state : dict state directly from simulation.describe_state
+        output : dict state according to CAOS.
+        """
+        return self.observation_space.observe(state)
+
+    def get_reward_from_state(self, state:Dict) -> float:
+        return self.reward_function.calculate(state)
+
+    @abstractmethod
+    def get_action(self, obs:ObsType, reward:float=None):
+        # in RL agent, this method will send CAOS observation to GATE RL agent, then receive a int 1-40,
+        # then use a bespoke conversion to take 1-40 int back into CAOS action
+        return ('NODE', 'SERVICE', 'SCAN', '<fake-node-sid>', '<fake-service-sid>')
+
+    @abstractmethod
+    def format_request(self, action) -> List[str]:
+        # this will take something like APPLICATION.EXECUTE and add things like target_ip_address in simulator.
+        # therefore the execution definition needs to be a mapping from CAOS into SIMULATOR
+        """Format action into format expected by the simulator, and apply execution definition if applicable."""
+        return ['network', 'nodes', '<fake-node-uuid>', 'file_system', 'folder', 'root', 'scan']
+
+
+
+
 
 class AbstractScriptedAgent(AbstractAgent):
     """Base class for actors which generate their own behaviour."""
 
     ...
 
+class RandomAgent(AbstractScriptedAgent):
+    """Agent that ignores its observation and acts completely at random."""
+
+    def get_action(self, obs:ObsType, reward:float=None):
+        return self.action_space.space.sample()
+
 
 class AbstractGATEAgent(AbstractAgent):
     """Base class for actors controlled via external messages, such as RL policies."""

src/primaite/game/agent/observations.py

@@ -55,7 +55,7 @@ class AbstractObservation(ABC):
 
 
 class FileObservation(AbstractObservation):
-    def __init__(self, where: List[str] = []) -> None:
+    def __init__(self, where: Optional[List[str]] = None) -> None:
         """
         _summary_
 
@@ -68,12 +68,12 @@ class FileObservation(AbstractObservation):
         :type where: Optional[List[str]]
         """
         super().__init__()
-        self.where: List[str] = where
+        self.where: Optional[List[str]] = where
         self.default_observation: spaces.Space = {"health_status": 0}
         "Default observation is what should be returned when the file doesn't exist, e.g. after it has been deleted."
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
         file_state = access_from_nested_dict(state, self.where)
         if file_state is NOT_PRESENT_IN_STATE:
@@ -89,21 +89,21 @@ class ServiceObservation(AbstractObservation):
     default_observation: spaces.Space = {"operating_status": 0, "health_status": 0}
     "Default observation is what should be returned when the service doesn't exist."
 
-    def __init__(self, where: List[str] = []) -> None:
+    def __init__(self, where: Optional[List[str]] = None) -> None:
         """
         :param where: Store information about where in the simulation state dictionary to find the relevant information.
             Optional. If None, this corresponds that the file does not exist and the observation will be populated with
             zeroes.
 
             A typical location for a service looks like this:
-            `['network','nodes',<node_uuid>,'servics', <service_uuid>]`
+            `['network','nodes',<node_uuid>,'services', <service_uuid>]`
         :type where: Optional[List[str]]
         """
         super().__init__()
-        self.where: List[str] = where
+        self.where: Optional[List[str]] = where
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
 
         service_state = access_from_nested_dict(state, self.where)
@@ -120,7 +120,7 @@ class LinkObservation(AbstractObservation):
     default_observation: spaces.Space = {"protocols": {"all": {"load": 0}}}
     "Default observation is what should be returned when the link doesn't exist."
 
-    def __init__(self, where: List[str] = []) -> None:
+    def __init__(self, where: Optional[List[str]] = None) -> None:
         """
         :param where: Store information about where in the simulation state dictionary to find the relevant information.
             Optional. If None, this corresponds that the file does not exist and the observation will be populated with
@@ -131,10 +131,10 @@ class LinkObservation(AbstractObservation):
         :type where: Optional[List[str]]
         """
         super().__init__()
-        self.where: List[str] = where
+        self.where: Optional[List[str]] = where
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
 
         link_state = access_from_nested_dict(state, self.where)
@@ -156,7 +156,7 @@ class LinkObservation(AbstractObservation):
 
 
 class FolderObservation(AbstractObservation):
-    def __init__(self, where: List[str] = [], files: List[FileObservation] = []) -> None:
+    def __init__(self, where: Optional[List[str]] = None, files: List[FileObservation] = []) -> None:
         """Initialise folder Observation, including files inside of the folder.
 
         :param where: Where in the simulation state dictionary to find the relevant information for this folder.
@@ -175,7 +175,7 @@ class FolderObservation(AbstractObservation):
         """
         super().__init__()
 
-        self.where: List[str] = where
+        self.where: Optional[List[str]] = where
 
         self.files: List[FileObservation] = files
 
@@ -185,7 +185,7 @@ class FolderObservation(AbstractObservation):
         }
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
         folder_state = access_from_nested_dict(state, self.where)
         if folder_state is NOT_PRESENT_IN_STATE:
@@ -213,12 +213,12 @@ class FolderObservation(AbstractObservation):
 class NicObservation(AbstractObservation):
     default_observation: spaces.Space = {"nic_status": 0}
 
-    def __init__(self, where: List[str] = []) -> None:
-        super.__init__()
-        self.where: List[str] = where
+    def __init__(self, where: Optional[List[str]] = None) -> None:
+        super().__init__()
+        self.where: Optional[List[str]] = where
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
         nic_state = access_from_nested_dict(state, self.where)
         if nic_state is NOT_PRESENT_IN_STATE:
@@ -234,10 +234,11 @@ class NicObservation(AbstractObservation):
 class NodeObservation(AbstractObservation):
     def __init__(
         self,
-        where: List[str] = [],
+        where: Optional[List[str]] = None,
         services: List[ServiceObservation] = [],
         folders: List[FolderObservation] = [],
         nics: List[NicObservation] = [],
+        logon_status:bool=False
     ) -> None:
         """
         Configurable observation for a node in the simulation.
@@ -259,12 +260,13 @@ class NodeObservation(AbstractObservation):
         :param max_nics: Max number of NICS in this node's obs space, defaults to 5
         :type max_nics: int, optional
         """
-        super.__init__()
-        self.where: List[str] = where
+        super().__init__()
+        self.where: Optional[List[str]] = where
 
         self.services: List[ServiceObservation] = services
         self.folders: List[FolderObservation] = folders
         self.nics: List[NicObservation] = nics
+        self.logon_status:bool=logon_status
 
         self.default_observation: Dict = {
             "SERVICES": {i + 1: s.default_observation for i, s in enumerate(self.services)},
@@ -272,9 +274,11 @@ class NodeObservation(AbstractObservation):
             "NICS": {i + 1: n.default_observation for i, n in enumerate(self.nics)},
             "operating_status": 0,
         }
+        if self.logon_status:
+            self.default_observation['logon_status']=0
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
 
         node_state = access_from_nested_dict(state, self.where)
@@ -288,18 +292,24 @@ class NodeObservation(AbstractObservation):
         obs["operating_status"] = node_state["operating_state"]
         obs["NICS"] = {i + 1: nic.observe(state) for i, nic in enumerate(self.nics)}
 
+        if self.logon_status:
+            obs['logon_status'] = 0
+
         return obs
 
     @property
     def space(self) -> spaces.Space:
-        return spaces.Dict(
-            {
-                "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(0),
-                "NICS": spaces.Dict({i + 1: nic.space for i, nic in enumerate(self.nics)}),
-            }
-        )
+        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: nic.space for i, nic in enumerate(self.nics)}),
+        }
+        if self.logon_status:
+            space_shape['logon_status'] = spaces.Discrete(3)
+
+        return spaces.Dict(space_shape)
+
 
 
 class AclObservation(AbstractObservation):
@@ -308,41 +318,33 @@ 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, nodes: List[str], ports: List[int], protocols: list[str], where: List[str] = [], num_rules: int = 10
+        self, node_ip_to_id: Dict[str,int], ports: List[int], protocols: list[str], where: Optional[List[str]] = None, num_rules: int = 10
     ) -> None:
         super().__init__()
-        self.where: List[str] = where
+        self.where: Optional[List[str]] = where
         self.num_rules: int = num_rules
-        self.node_to_id: Dict[str, int] = {node: i + 1 for i, node in enumerate(nodes)}
+        self.node_to_id: Dict[str, int] = node_ip_to_id
         "List of node IP addresses, order in this list determines how they are converted to an ID"
-        self.port_to_id: Dict[int, int] = {port: i + 1 for i, port in enumerate(ports)}
+        self.port_to_id: Dict[int, int] = {port: i + 2 for i, port in enumerate(ports)}
         "List of ports which are part of the game that define the ordering when converting to an ID"
-        self.protocol_to_id: Dict[str, int] = {protocol: i + 1 for i, protocol in enumerate(protocols)}
+        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: spaces.Space = spaces.Dict(
-            {
-                "RULES": spaces.Dict(
-                    {
-                        i
-                        + 1: spaces.Dict(
-                            {
-                                "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)
-                    }
-                )
+        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,
+                }
+            for i in range(self.num_rules)
             }
-        )
+        }
 
     def observe(self, state: Dict) -> Dict:
-        if not self.where:
+        if self.where is None:
             return self.default_observation
         acl_state: Dict = access_from_nested_dict(state, self.where)
         if acl_state is NOT_PRESENT_IN_STATE:
@@ -379,16 +381,16 @@ class AclObservation(AbstractObservation):
             {
                 "RULE": spaces.Dict(
                     {
-                        i
-                        + 1: spaces.Dict(
+                        i + 1: spaces.Dict(
                             {
                                 "position": spaces.Discrete(self.num_rules),
                                 "permission": spaces.Discrete(3),
-                                "source_node_id": spaces.Discrete(len(self.nodes) + 1),
-                                "source_port": spaces.Discrete(len(self.ports) + 1),
-                                "dest_node_id": spaces.Discrete(len(self.nodes) + 1),
-                                "dest_port": spaces.Discrete(len(self.ports) + 1),
-                                "protocol": spaces.Discrete(len(self.protocols) + 1),
+                                # adding two to lengths is to account for reserved values 0 (unused) and 1 (any)
+                                "source_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
+                                "source_port": spaces.Discrete(len(self.port_to_id) + 2),
+                                "dest_node_id": spaces.Discrete(len(set(self.node_to_id.values())) + 2),
+                                "dest_port": spaces.Discrete(len(self.port_to_id) + 2),
+                                "protocol": spaces.Discrete(len(self.protocol_to_id) + 2),
                             }
                         )
                         for i in range(self.num_rules)
@@ -398,14 +400,96 @@ class AclObservation(AbstractObservation):
         )
 
 
-class ICSObservation(AbstractObservation):
-    def observe(self, state: Dict) -> Any:
-        return 0
+
+
+class NullObservation(AbstractObservation):
+    def __init__(self, where:Optional[List[str]]=None):
+        self.default_observation: Dict = {}
+
+    def observe(self, state: Dict) -> Dict:
+        return {}
 
     @property
     def space(self) -> spaces.Space:
-        return spaces.Discrete(1)
+        return spaces.Dict({})
 
+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,
+    ) -> None:
+        super().__init__()
+        self.where: Optional[List[str]] = where
+
+        self.nodes: List[NodeObservation] = nodes
+        self.links: List[LinkObservation] = links
+        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)},
+            "ACL": self.acl.default_observation,
+            "ICS": self.ics.default_observation,
+        }
+
+    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)}
+
+        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,
+        })
+
+    @classmethod
+    def from_config(cls, config:Dict, sim:Simulation):
+        nodes = ...
+        links = ...
+        acl = ...
+        ics = ...
+        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:
+        super().__init__()
+        self.where:Optional[List[str]] = where
+        self.nodes: List[NodeObservation] = nodes
+
+        self.default_observation=...#TODO
+
+    def observe(self, state: Dict) -> Any:
+        return super().observe(state)
+
+    @property
+    def space(self) -> spaces.Space:
+        ... #TODO
+
+    @classmethod
+    def from_config(cls, config: Dict, sim:Simulation):
+        ... #TODO
 
 class ObservationSpace:
     """
@@ -422,29 +506,12 @@ class ObservationSpace:
     # what this class does:
     # keep a list of observations
     # create observations for an actor from the config
-    def __init__(
-        self,
-        simulation: Simulation,
-        nodes: List[NodeObservation] = [],
-        links: List[LinkObservation] = [],
-        acl: Optional[AclObservation] = None,
-        ics: Optional[ICSObservation] = None,
-    ) -> None:
-        self.simulation: Simulation = simulation
-        self.parts: Dict[str, AbstractObservation] = {}
+    def __init__(self, observation:AbstractObservation) -> None:
+        self.obs: AbstractObservation = observation
 
-        self.nodes: List[NodeObservation] = nodes
-        self.links: List[LinkObservation] = links
-        self.acl: Optional[AclObservation] = acl
-        self.ics: Optional[ICSObservation] = ics
-
-    def observe(self) -> None:
-        ...
+    def observe(self, state) -> Dict:
+        return self.obs.observe(state)
 
     @property
     def space(self) -> None:
-        ...
-
-    @classmethod
-    def from_config(self) -> None:
-        ...
+        return self.obs.space

src/primaite/game/agent/rewards.py

@@ -1,20 +1,17 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
-from pydantic import BaseModel
-
-
-class AbstractReward(BaseModel):
-    def __call__(self, states: List[Dict]) -> float:
-        """_summary_
-
-        :param state: _description_
-        :type state: Dict
-        :return: _description_
-        :rtype: float
-        """
+class AbstractReward():
+    def __init__(self):
         ...
 
+    def calculate(self, state:Dict) -> float:
+        return 0.3
 
-class RewardFunction(BaseModel):
-    ...
+
+class RewardFunction():
+    def __init__(self, reward_function:AbstractReward):
+        self.reward: AbstractReward = reward_function
+
+    def calculate(self, state:Dict) -> float:
+        return self.reward.calculate(state)

src/primaite/game/session.py

@@ -4,3 +4,49 @@
 # 3. create actors and configure their actions/observations/rewards/ anything else
 # 4. Create connection with ARCD GATE
 # 5. idk
+
+from primaite.simulator.sim_container import Simulation
+from primaite.game.agent.interface import AbstractAgent
+
+from typing import List
+
+class PrimaiteSession:
+    def __init__(self):
+        self.simulation: Simulation = Simulation()
+        self.agents:List[AbstractAgent] = []
+        self.step_counter:int = 0
+        self.episode_counter:int = 0
+
+
+    def step(self):
+        # 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
+            sim_state = self.simulation.describe_state()
+
+            # 6. each agent takes most recent state and converts it to CAOS observation
+            agent_obs = agent.get_obs_from_state(sim_state)
+
+            # 7. meanwhile each agent also takes state and calculates reward
+            agent_reward = agent.get_reward_from_state(sim_state)
+
+            # 8. each agent takes observation and applies decision rule to observation to create CAOS
+            #    action(such as random, rulebased, or send to GATE) (therefore, converting CAOS action
+            #    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
+            agent_action = 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)
+            agent_request = agent.format_request(agent_action)
+
+            # 10. primaite session receives the action from the agents and asks the simulation to apply each
+            self.simulation.apply_action(agent_request)
+
+        self.simulation.apply_timestep(self.step_counter)
+        self.step_counter += 1
+