Add network action

src/primaite/game/agent/actions.py

@@ -1,12 +1,13 @@
+import itertools
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple
-import itertools
-
-
-from primaite.simulator.sim_container import Simulation
 
 from gym import spaces
 
+from primaite.game.session import PrimaiteSession
+from primaite.simulator.sim_container import Simulation
+
+
 class ExecutionDefiniton(ABC):
     """
     Converter from actions to simulator requests.
@@ -23,9 +24,8 @@ class ExecutionDefiniton(ABC):
 
 
 class AbstractAction(ABC):
-
     @abstractmethod
-    def __init__(self, manager:"ActionManager", **kwargs) -> None:
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
         """
         Init method for action.
 
@@ -35,13 +35,12 @@ class AbstractAction(ABC):
         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 = ""
+        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
-
+        self.shape: Dict[str, int] = {}
+        """Dictionary describing the number of options for each parameter of this action. The keys of this dict must
+        align with the keyword args of the form_request method."""
+        self.manager: ActionManager = manager
 
     @abstractmethod
     def form_request(self) -> List[str]:
@@ -50,14 +49,20 @@ class AbstractAction(ABC):
 
 
 class DoNothingAction(AbstractAction):
-    def __init__(self, manager:"ActionManager", **kwargs) -> None:
+    def __init__(self, manager: "ActionManager", **kwargs) -> None:
         super().__init__(manager=manager)
         self.name = "DONOTHING"
-        self.shape = (1,)
+        self.shape: Dict[str, int] = {
+            "dummy": 1,
+        }
+        # This action does not accept any parameters, therefore it technically has a gymnasium shape of Discrete(1),
+        # i.e. a choice between one option. To make enumerating this action easier, we are adding a 'dummy' paramter
+        # with one option. This just aids the Action Manager to enumerate all possibilities.
 
     def form_request(self) -> List[str]:
         return ["do_nothing"]
 
+
 class NodeServiceAbstractAction(AbstractAction):
     """
     Base class for service actions.
@@ -65,211 +70,284 @@ class NodeServiceAbstractAction(AbstractAction):
     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]:
+    @abstractmethod
+    def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Dict[str, int] = {"node_id": num_nodes, "service_id": 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]
+        return ["network", "node", node_uuid, "services", service_uuid, self.verb]
+
 
 class NodeServiceScanAction(NodeServiceAbstractAction):
-    def __init__(self, manager:"ActionManager", num_nodes, num_services, **kwargs) -> None:
+    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:
+    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:
+    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:
+    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:
+    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:
+    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:
+    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:
+    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:
+    def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None:
         super().__init__(manager=manager)
-        self.shape = (num_nodes, num_folders)
+        self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders}
         self.verb: str
 
-    def form_request(self, node_id:int, folder_id:int) -> List[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]
+        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:
+    def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None:
         super().__init__(manager, num_nodes, num_folders, **kwargs)
-        self.verb:str = "scan"
+        self.verb: str = "scan"
+
 
 class NodeFolderCheckhashAction(NodeFolderAbstractAction):
-    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+    def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None:
         super().__init__(manager, num_nodes, num_folders, **kwargs)
-        self.verb:str = "checkhash"
+        self.verb: str = "checkhash"
+
 
 class NodeFolderRepairAction(NodeFolderAbstractAction):
-    def __init__(self, manager:"ActionManager", num_nodes, num_folders, **kwargs) -> None:
+    def __init__(self, manager: "ActionManager", num_nodes, num_folders, **kwargs) -> None:
         super().__init__(manager, num_nodes, num_folders, **kwargs)
-        self.verb:str = "repair"
+        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"
+        self.verb: str = "restore"
 
 
 class NodeFileAbstractAction(AbstractAction):
     @abstractmethod
-    def __init__(self, manager:"ActionManager", num_nodes:int, num_folders:int, num_files:int, **kwargs) -> None:
+    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
+        self.shape: Dict[str, int] = {"node_id": num_nodes, "folder_id": num_folders, "file_id": num_files}
+        self.verb: str
 
-    def form_request(self, node_id:int, folder_id:int, file_id:int) -> List[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]
+        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.shape: Dict[str, int] = {"node_id": num_nodes}
         self.verb: str
 
-    def form_request(self, node_id:int) -> List[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'
+        self.verb = "scan"
+
 
 class NodeShutdownAction(NodeAbstractAction):
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager)
-        self.verb = 'shutdown'
+        self.verb = "shutdown"
+
 
 class NodeStartupAction(NodeAbstractAction):
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager)
-        self.verb = 'start'
+        self.verb = "start"
+
 
 class NodeResetAction(NodeAbstractAction):
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager)
-        self.verb = 'reset'
+        self.verb = "reset"
+
 
 class NetworkACLAddRuleAction(AbstractAction):
-    def __init__(self,
-                 manager: "ActionManager",
-                 target_router_uuid:str,
-                 max_acl_rules:int,
-                 num_ips:int,
-                 num_ports:int,
-                 num_protocols:int,
-                 **kwargs) -> None:
+    def __init__(
+        self,
+        manager: "ActionManager",
+        target_router_uuid: str,
+        max_acl_rules: int,
+        num_ips: int,
+        num_ports: int,
+        num_protocols: int,
+        **kwargs,
+    ) -> None:
         super().__init__(manager=manager)
         num_permissions = 2
-        self.shape: Tuple[int] = (max_acl_rules, num_permissions, num_ips, num_ips, num_ports, num_ports, num_protocols)
-        self.target_router_uuid:str = target_router_uuid
+        self.shape: Dict[str, int] = {
+            "position": max_acl_rules,
+            "permission": num_permissions,
+            "source_ip_idx": num_ips,
+            "dest_ip_idx": num_ips,
+            "source_port_idx": num_ports,
+            "dest_port_idx": num_ports,
+            "protocol_idx": num_protocols,
+        }
+        self.target_router_uuid: str = target_router_uuid
 
-    def form_request(self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx) -> List[str]:
+    def form_request(
+        self, position, permission, source_ip_idx, dest_ip_idx, source_port_idx, dest_port_idx, protocol_idx
+    ) -> List[str]:
         protocol = self.manager.get_internet_protocol_by_idx(protocol_idx)
         src_ip = self.manager.get_ip_address_by_idx(source_ip_idx)
         src_port = self.manager.get_port_by_idx(source_port_idx)
         dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx)
         dst_port = self.manager.get_port_by_idx(dest_port_idx)
         return [
-            'network',
-            'node',
+            "network",
+            "node",
             self.target_router_uuid,
-            'acl',
-            'add_rule',
+            "acl",
+            "add_rule",
             permission,
             protocol,
             src_ip,
             src_port,
             dst_ip,
             dst_port,
-            position
+            position,
         ]
 
 
+class NetworkACLRemoveRuleAction(AbstractAction):
+    def __init__(self, manager: "ActionManager", target_router_uuid: str, max_acl_rules: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Dict[str, int] = {"position": max_acl_rules}
+        self.target_router_uuid: str = target_router_uuid
 
+    def form_request(self, position: int) -> List[str]:
+        return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position]
+
+
+class NetworkNICEnableAction(AbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
+
+    def form_request(self, node_id: int, nic_id: int) -> List[str]:
+        return [
+            "network",
+            "node",
+            self.manager.get_node_uuid_by_idx(node_idx=node_id),
+            "nic",
+            self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id),
+            "enable",
+        ]
+
+
+class NetworkNICDisableAction(AbstractAction):
+    def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
+        super().__init__(manager=manager)
+        self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
+
+    def form_request(self, node_id: int, nic_id: int) -> List[str]:
+        return [
+            "network",
+            "node",
+            self.manager.get_node_uuid_by_idx(node_idx=node_id),
+            "nic",
+            self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id),
+            "disable",
+        ]
 
 
 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.
 
@@ -278,69 +356,83 @@ class ActionManager:
     # 0: DONOTHING
     # 1: NODE, FILE, SCAN, NODEID=2, FOLDERID=1, FILEID=0
     # 2: ......
-    __act_class_identifiers:Dict[str,type] = {
+    __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_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,
+        "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": NetworkNICEnableAction,
+        "NETWORK_NIC_DISABLE": NetworkNICDisableAction,
     }
 
+    def __init__(
+        self,
+        session: PrimaiteSession,  # reference to session for looking up stuff
+        actions: List[str],  # stores list of actions available to agent
+        node_uuids: List[str],  # allows mapping index to node
+        max_folders_per_node: int = 2,  # allows calculating shape
+        max_files_per_folder: int = 2,  # allows calculating shape
+        max_services_per_node: int = 2,  # allows calculating shape
+        max_nics_per_node: int = 8,  # allows calculating shape
+        max_acl_rules: int = 10,  # allows calculating shape
+        protocols: List[str] = ["TCP", "UDP", "ICMP"],  # allow mapping index to protocol
+        ports: List[str] = ["HTTP", "DNS", "ARP", "FTP"],  # allow mapping index to port
+        ip_address_list: Optional[List[str]] = None,  # to allow us to map an index to an ip address.
+        act_map: Optional[Dict[int, Dict]] = None,  # allows restricting set of possible actions
+    ) -> None:
+        self.session: PrimaiteSession = session
+        self.sim: Simulation = self.session.simulation
+        self.node_uuids: List[str] = node_uuids
+        self.protocols: List[str] = protocols
+        self.ports: List[str] = ports
 
-    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,
-                 protocols:List[str]=['TCP','UDP','ICMP'],
-                 ports:List[str]=['HTTP','DNS','ARP','FTP'],
-                 ip_address_list:Optional[List[str]]=None,
-                 act_map:Optional[Dict[int, Dict]]=None) -> None:
-        self.sim: Simulation = sim
-        self.node_uuids:List[str] = node_uuids
-        self.protocols:List[str] = protocols
-        self.ports:List[str] = ports
+        self.ip_address_list: List[str]
+        if ip_address_list is not None:
+            self.ip_address_list = ip_address_list
+        else:
+            self.ip_address_list = []
+            for node_uuid in self.node_uuids:
+                node_obj = self.sim.network.nodes[node_uuid]
+                nics = node_obj.nics
+                for nic_uuid, nic_obj in nics.items():
+                    self.ip_address_list.append(nic_obj.ip_address)
 
         action_args = {
             "num_nodes": len(node_uuids),
-            "num_folders":max_folders_per_node,
+            "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,
             "num_protocols": len(self.protocols),
             "num_ports": len(self.protocols),
-            "num_ips":}
+            "num_ips": len(self.ip_address_list),
+        }
         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]] = {}
+        self.action_map: Dict[int, Tuple[str, Dict]] = {}
         """
         Action mapping that converts an integer to a specific action and parameter choice.
 
@@ -350,21 +442,30 @@ class ActionManager:
         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()}
+            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 _enumerate_actions(
+        self,
+    ) -> Dict[int, Tuple[AbstractAction, Dict]]:
+        all_action_possibilities = []
+        for action in self.actions.values():
+            param_names = (list(action.shape.keys()),)
+            num_possibilities = list(action.shape.values())
+            possibilities = [range(n) for n in num_possibilities]
 
-    def get_action(self, action: int) -> Tuple[str,Dict]:
+            itertools.product(action.shape.values())
+            all_action_possibilities.append((action, {}))
+
+    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
 
-    def form_request(self, action_identifier:str, action_options:Dict):
+    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)
@@ -380,37 +481,57 @@ class ActionManager:
         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
+        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:
+        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
+        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
+        return service_uuids[service_idx] if len(service_uuids) > service_idx else None
 
-    def get_internet_protocol_by_idx(self, protocol_idx:int) -> str:
+    def get_internet_protocol_by_idx(self, protocol_idx: int) -> str:
+        return self.protocols[protocol_idx]
 
+    def get_ip_address_by_idx(self, ip_idx: int) -> str:
+        return self.ip_address_list[ip_idx]
 
-        # protocol = self.manager.get_internet_protocol_by_idx(protocol_idx)
-        # src_ip = self.manager.get_ip_address_by_idx(source_ip_idx)
-        # src_port = self.manager.get_port_by_idx(source_port_idx)
-        # dst_ip = self.manager.get_ip_address_by_idx(dest_ip_idx)
-        # dst_port = self.manager.get_port_by_idx(dest_port_idx)
+    def get_port_by_idx(self, port_idx: int) -> str:
+        return self.ports[port_idx]
 
+    def get_nic_uuid_by_idx(self, node_idx: int, nic_idx: int) -> str:
+        node_uuid = self.get_node_uuid_by_idx(node_idx)
+        node_obj = self.sim.network.nodes[node_uuid]
+        nics = list(node_obj.nics.keys())
+        if len(nics) <= nic_idx:
+            return None
+        return nics[nic_idx]
 
+    @classmethod
+    def from_config(cls, session: PrimaiteSession, cfg: Dict) -> "ActionManager":
+        obj = cls(
+            session=session,
+            actions=cfg["action_list"],
+            node_uuids=cfg["options"]["nodes"],
+            max_folders_per_node=cfg["options"]["max_folders_per_node"],
+            max_files_per_folder=cfg["options"]["max_files_per_folder"],
+            max_services_per_node=cfg["options"]["max_services_per_node"],
+            max_nics_per_node=cfg["options"]["max_nics_per_node"],
+            max_acl_rules=cfg["options"]["max_acl_rules"],
+            max_X=cfg["options"]["max_X"],
+            protocols=session.options.ports,
+            ports=session.options.protocols,
+            ip_address_list=None,
+            act_map=cfg["action_map"],
+        )
 
-class UC2RedActions(AbstractAction):
-    ...
-
-class UC2GreenActionSpace(ActionManager):
-    ...
+        return obj

src/primaite/game/agent/interface.py

@@ -2,14 +2,16 @@
 # 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, Union, TypeAlias
+from typing import Any, Dict, List, Optional, TypeAlias, Union
+
 import numpy as np
 
 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]
+ObsType: TypeAlias = Union[Dict, np.ndarray]
+
 
 class AbstractAgent(ABC):
     """Base class for scripted and RL agents."""
@@ -28,31 +30,28 @@ class AbstractAgent(ABC):
         # 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:
+    def convert_state_to_obs(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:
+    def calculate_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,
+    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 0-39,
         # then use a bespoke conversion to take 1-40 int back into CAOS action
-        return ('NODE', 'SERVICE', 'SCAN', '<fake-node-sid>', '<fake-service-sid>')
+        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']
-
-
-
+        return ["network", "nodes", "<fake-node-uuid>", "file_system", "folder", "root", "scan"]
 
 
 class AbstractScriptedAgent(AbstractAgent):
@@ -60,10 +59,11 @@ class AbstractScriptedAgent(AbstractAgent):
 
     ...
 
+
 class RandomAgent(AbstractScriptedAgent):
     """Agent that ignores its observation and acts completely at random."""
 
-    def get_action(self, obs:ObsType, reward:float=None):
+    def get_action(self, obs: ObsType, reward: float = None):
         return self.action_space.space.sample()