Add documentation

src/primaite/game/agent/actions.py

@@ -1,6 +1,16 @@
+"""
+This module contains the ActionManager class which belongs to the Agent class.
+
+An agent's action space is made up of a collection of actions. Each action is an instance of a subclass of
+AbstractAction. The ActionManager is responsible for:
+    1. Creating the action space from a list of action types.
+    2. Converting an integer action choice into a specific action and parameter choice.
+    3. Converting an action and parameter choice into a request which can be ingested by the PrimAITE simulation. This
+        ensures that requests conform to the simulator's request format.
+"""
 import itertools
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
 
 from gymnasium import spaces
 
@@ -13,22 +23,9 @@ if TYPE_CHECKING:
     from primaite.game.session import PrimaiteSession
 
 
-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):
+    """Base class for actions."""
+
     @abstractmethod
     def __init__(self, manager: "ActionManager", **kwargs) -> None:
         """
@@ -46,6 +43,8 @@ class AbstractAction(ABC):
         """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
+        """Reference to the ActionManager which created this action. This is used to access the session and simulation
+        objects."""
 
     @abstractmethod
     def form_request(self) -> List[str]:
@@ -54,6 +53,8 @@ class AbstractAction(ABC):
 
 
 class DoNothingAction(AbstractAction):
+    """Action which does nothing. This is here to allow agents to be idle if they choose to."""
+
     def __init__(self, manager: "ActionManager", **kwargs) -> None:
         super().__init__(manager=manager)
         self.name = "DONOTHING"
@@ -65,6 +66,7 @@ class DoNothingAction(AbstractAction):
         # with one option. This just aids the Action Manager to enumerate all possibilities.
 
     def form_request(self, **kwargs) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         return ["do_nothing"]
 
 
@@ -77,12 +79,13 @@ class NodeServiceAbstractAction(AbstractAction):
     """
 
     @abstractmethod
-    def __init__(self, manager: "ActionManager", num_nodes, num_services, **kwargs) -> None:
+    def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **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]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         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:
@@ -91,54 +94,77 @@ class NodeServiceAbstractAction(AbstractAction):
 
 
 class NodeServiceScanAction(NodeServiceAbstractAction):
+    """Action which scans a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "scan"
 
 
 class NodeServiceStopAction(NodeServiceAbstractAction):
+    """Action which stops a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "stop"
 
 
 class NodeServiceStartAction(NodeServiceAbstractAction):
+    """Action which starts a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "start"
 
 
 class NodeServicePauseAction(NodeServiceAbstractAction):
+    """Action which pauses a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "pause"
 
 
 class NodeServiceResumeAction(NodeServiceAbstractAction):
+    """Action which resumes a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "resume"
 
 
 class NodeServiceRestartAction(NodeServiceAbstractAction):
+    """Action which restarts a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "restart"
 
 
 class NodeServiceDisableAction(NodeServiceAbstractAction):
+    """Action which disables a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "disable"
 
 
 class NodeServiceEnableAction(NodeServiceAbstractAction):
+    """Action which enables a service."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_services: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, num_services=num_services)
         self.verb = "enable"
 
 
 class NodeFolderAbstractAction(AbstractAction):
+    """
+    Base class for folder actions.
+
+    Any action which applies to a folder and uses node_id and folder_id as its only two parameters can inherit from
+    this base class.
+    """
+
     @abstractmethod
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
         super().__init__(manager=manager)
@@ -146,6 +172,7 @@ class NodeFolderAbstractAction(AbstractAction):
         self.verb: str
 
     def form_request(self, node_id: int, folder_id: int) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         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:
@@ -154,30 +181,44 @@ class NodeFolderAbstractAction(AbstractAction):
 
 
 class NodeFolderScanAction(NodeFolderAbstractAction):
+    """Action which scans a folder."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
         self.verb: str = "scan"
 
 
 class NodeFolderCheckhashAction(NodeFolderAbstractAction):
+    """Action which checks the hash of a folder."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
         self.verb: str = "checkhash"
 
 
 class NodeFolderRepairAction(NodeFolderAbstractAction):
+    """Action which repairs a folder."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
         self.verb: str = "repair"
 
 
 class NodeFolderRestoreAction(NodeFolderAbstractAction):
+    """Action which restores a folder."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
         self.verb: str = "restore"
 
 
 class NodeFileAbstractAction(AbstractAction):
+    """Abstract base class for file actions.
+
+    Any action which applies to a file and uses node_id, folder_id, and file_id as its only three parameters can inherit
+    from this base class.
+    """
+
     @abstractmethod
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager=manager)
@@ -185,6 +226,7 @@ class NodeFileAbstractAction(AbstractAction):
         self.verb: str
 
     def form_request(self, node_id: int, folder_id: int, file_id: int) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         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)
@@ -194,42 +236,60 @@ class NodeFileAbstractAction(AbstractAction):
 
 
 class NodeFileScanAction(NodeFileAbstractAction):
+    """Action which scans a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "scan"
 
 
 class NodeFileCheckhashAction(NodeFileAbstractAction):
+    """Action which checks the hash of a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "checkhash"
 
 
 class NodeFileDeleteAction(NodeFileAbstractAction):
+    """Action which deletes a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "delete"
 
 
 class NodeFileRepairAction(NodeFileAbstractAction):
+    """Action which repairs a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "repair"
 
 
 class NodeFileRestoreAction(NodeFileAbstractAction):
+    """Action which restores a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "restore"
 
 
 class NodeFileCorruptAction(NodeFileAbstractAction):
+    """Action which corrupts a file."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, num_files: int, **kwargs) -> None:
         super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, num_files=num_files, **kwargs)
         self.verb = "corrupt"
 
 
 class NodeAbstractAction(AbstractAction):
+    """
+    Abstract base class for node actions.
+
+    Any action which applies to a node and uses node_id as its only parameter can inherit from this base class.
+    """
+
     @abstractmethod
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager)
@@ -237,35 +297,46 @@ class NodeAbstractAction(AbstractAction):
         self.verb: str
 
     def form_request(self, node_id: int) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         node_uuid = self.manager.get_node_uuid_by_idx(node_id)
         return ["network", "node", node_uuid, self.verb]
 
 
 class NodeOSScanAction(NodeAbstractAction):
+    """Action which scans a node's OS."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes)
         self.verb = "scan"
 
 
 class NodeShutdownAction(NodeAbstractAction):
+    """Action which shuts down a node."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes)
         self.verb = "shutdown"
 
 
 class NodeStartupAction(NodeAbstractAction):
+    """Action which starts up a node."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes)
         self.verb = "startup"
 
 
 class NodeResetAction(NodeAbstractAction):
+    """Action which resets a node."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes)
         self.verb = "reset"
 
 
 class NetworkACLAddRuleAction(AbstractAction):
+    """Action which adds a rule to a router's ACL."""
+
     def __init__(
         self,
         manager: "ActionManager",
@@ -276,6 +347,21 @@ class NetworkACLAddRuleAction(AbstractAction):
         num_protocols: int,
         **kwargs,
     ) -> None:
+        """Init method for NetworkACLAddRuleAction.
+
+        :param manager: Reference to the ActionManager which created this action.
+        :type manager: ActionManager
+        :param target_router_uuid: UUID of the router to which the ACL rule should be added.
+        :type target_router_uuid: str
+        :param max_acl_rules: Maximum number of ACL rules that can be added to the router.
+        :type max_acl_rules: int
+        :param num_ips: Number of IP addresses in the simulation.
+        :type num_ips: int
+        :param num_ports: Number of ports in the simulation.
+        :type num_ports: int
+        :param num_protocols: Number of protocols in the simulation.
+        :type num_protocols: int
+        """
         super().__init__(manager=manager)
         num_permissions = 3
         self.shape: Dict[str, int] = {
@@ -290,8 +376,16 @@ class NetworkACLAddRuleAction(AbstractAction):
         self.target_router_uuid: str = target_router_uuid
 
     def form_request(
-        self, position, permission, source_ip_id, dest_ip_id, source_port_id, dest_port_id, protocol_id
+        self,
+        position: int,
+        permission: int,
+        source_ip_id: int,
+        dest_ip_id: int,
+        source_port_id: int,
+        dest_port_id: int,
+        protocol_id: int,
     ) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         if permission == 0:
             permission_str = "UNUSED"
             return ["do_nothing"]  # NOT SUPPORTED, JUST DO NOTHING IF WE COME ACROSS THIS
@@ -354,22 +448,51 @@ class NetworkACLAddRuleAction(AbstractAction):
 
 
 class NetworkACLRemoveRuleAction(AbstractAction):
+    """Action which removes a rule from a router's ACL."""
+
     def __init__(self, manager: "ActionManager", target_router_uuid: str, max_acl_rules: int, **kwargs) -> None:
+        """Init method for NetworkACLRemoveRuleAction.
+
+        :param manager: Reference to the ActionManager which created this action.
+        :type manager: ActionManager
+        :param target_router_uuid: UUID of the router from which the ACL rule should be removed.
+        :type target_router_uuid: str
+        :param max_acl_rules: Maximum number of ACL rules that can be added to the router.
+        :type max_acl_rules: int
+        """
         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 the action formatted as a request which can be ingested by the PrimAITE simulation."""
         return ["network", "node", self.target_router_uuid, "acl", "remove_rule", position]
 
 
 class NetworkNICAbstractAction(AbstractAction):
+    """
+    Abstract base class for NIC actions.
+
+    Any action which applies to a NIC and uses node_id and nic_id as its only two parameters can inherit from this base
+    class.
+    """
+
     def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
+        """Init method for NetworkNICAbstractAction.
+
+        :param manager: Reference to the ActionManager which created this action.
+        :type manager: ActionManager
+        :param num_nodes: Number of nodes in the simulation.
+        :type num_nodes: int
+        :param max_nics_per_node: Maximum number of NICs per node.
+        :type max_nics_per_node: int
+        """
         super().__init__(manager=manager)
         self.shape: Dict[str, int] = {"node_id": num_nodes, "nic_id": max_nics_per_node}
         self.verb: str
 
     def form_request(self, node_id: int, nic_id: int) -> List[str]:
+        """Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
         node_uuid = self.manager.get_node_uuid_by_idx(node_idx=node_id)
         nic_uuid = self.manager.get_nic_uuid_by_idx(node_idx=node_id, nic_idx=nic_id)
         if node_uuid is None or nic_uuid is None:
@@ -385,45 +508,24 @@ class NetworkNICAbstractAction(AbstractAction):
 
 
 class NetworkNICEnableAction(NetworkNICAbstractAction):
+    """Action which enables a NIC."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
         super().__init__(manager=manager, num_nodes=num_nodes, max_nics_per_node=max_nics_per_node, **kwargs)
         self.verb = "enable"
 
 
 class NetworkNICDisableAction(NetworkNICAbstractAction):
+    """Action which disables a NIC."""
+
     def __init__(self, manager: "ActionManager", num_nodes: int, max_nics_per_node: int, **kwargs) -> None:
         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)
-#         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.
-    #
+    """Class which manages the action space for an agent."""
 
-    # 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,
@@ -453,6 +555,7 @@ class ActionManager:
         "NETWORK_NIC_ENABLE": NetworkNICEnableAction,
         "NETWORK_NIC_DISABLE": NetworkNICDisableAction,
     }
+    """Dictionary which maps action type strings to the corresponding action class."""
 
     def __init__(
         self,
@@ -469,6 +572,33 @@ class ActionManager:
         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:
+        """Init method for ActionManager.
+
+        :param session: Reference to the session to which the agent belongs.
+        :type session: PrimaiteSession
+        :param actions: List of action types which should be made available to the agent.
+        :type actions: List[str]
+        :param node_uuids: List of node UUIDs that this agent can act on.
+        :type node_uuids: List[str]
+        :param max_folders_per_node: Maximum number of folders per node. Used for calculating action shape.
+        :type max_folders_per_node: int
+        :param max_files_per_folder: Maximum number of files per folder. Used for calculating action shape.
+        :type max_files_per_folder: int
+        :param max_services_per_node: Maximum number of services per node. Used for calculating action shape.
+        :type max_services_per_node: int
+        :param max_nics_per_node: Maximum number of NICs per node. Used for calculating action shape.
+        :type max_nics_per_node: int
+        :param max_acl_rules: Maximum number of ACL rules per router. Used for calculating action shape.
+        :type max_acl_rules: int
+        :param protocols: List of protocols that are available in the simulation. Used for calculating action shape.
+        :type protocols: List[str]
+        :param ports: List of ports that are available in the simulation. Used for calculating action shape.
+        :type ports: List[str]
+        :param ip_address_list: List of IP addresses that known to this agent. Used for calculating action shape.
+        :type ip_address_list: Optional[List[str]]
+        :param act_map: Action map which maps integers to actions. Used for restricting the set of possible actions.
+        :type act_map: Optional[Dict[int, Dict]]
+        """
         self.session: "PrimaiteSession" = session
         self.sim: Simulation = self.session.simulation
         self.node_uuids: List[str] = node_uuids
@@ -578,18 +708,48 @@ class ActionManager:
 
     @property
     def space(self) -> spaces.Space:
+        """Return the gymnasium action space for this agent."""
         return spaces.Discrete(len(self.action_map))
 
-    def get_node_uuid_by_idx(self, node_idx):
+    def get_node_uuid_by_idx(self, node_idx: int) -> str:
+        """Get the node UUID corresponding to the given index.
+
+        :param node_idx: The index of the node to retrieve.
+        :type node_idx: int
+        :return: The node UUID.
+        :rtype: str
+        """
         return self.node_uuids[node_idx]
 
-    def get_folder_uuid_by_idx(self, node_idx, folder_idx) -> Optional[str]:
+    def get_folder_uuid_by_idx(self, node_idx: int, folder_idx: int) -> Optional[str]:
+        """Get the folder UUID corresponding to the given node and folder indices.
+
+        :param node_idx: The index of the node.
+        :type node_idx: int
+        :param folder_idx: The index of the folder on the node.
+        :type folder_idx: int
+        :return: The UUID of the folder. Or None if the node has fewer folders than the given index.
+        :rtype: 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]:
+    def get_file_uuid_by_idx(self, node_idx: int, folder_idx: int, file_idx: int) -> Optional[str]:
+        """Get the file UUID corresponding to the given node, folder, and file indices.
+
+        :param node_idx: The index of the node.
+        :type node_idx: int
+        :param folder_idx: The index of the folder on the node.
+        :type folder_idx: int
+        :param file_idx: The index of the file in the folder.
+        :type file_idx: int
+        :return: The UUID of the file. Or None if the node has fewer folders than the given index, or the folder has
+            fewer files than the given index.
+        :rtype: 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())
@@ -599,22 +759,64 @@ class ActionManager:
         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]:
+    def get_service_uuid_by_idx(self, node_idx: int, service_idx: int) -> Optional[str]:
+        """Get the service UUID corresponding to the given node and service indices.
+
+        :param node_idx: The index of the node.
+        :type node_idx: int
+        :param service_idx: The index of the service on the node.
+        :type service_idx: int
+        :return: The UUID of the service. Or None if the node has fewer services than the given index.
+        :rtype: 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
 
     def get_internet_protocol_by_idx(self, protocol_idx: int) -> str:
+        """Get the internet protocol corresponding to the given index.
+
+        :param protocol_idx: The index of the protocol to retrieve.
+        :type protocol_idx: int
+        :return: The protocol.
+        :rtype: str
+        """
         return self.protocols[protocol_idx]
 
     def get_ip_address_by_idx(self, ip_idx: int) -> str:
+        """
+        Get the IP address corresponding to the given index.
+
+        :param ip_idx: The index of the IP address to retrieve.
+        :type ip_idx: int
+        :return: The IP address.
+        :rtype: str
+        """
         return self.ip_address_list[ip_idx]
 
     def get_port_by_idx(self, port_idx: int) -> str:
+        """
+        Get the port corresponding to the given index.
+
+        :param port_idx: The index of the port to retrieve.
+        :type port_idx: int
+        :return: The port.
+        :rtype: str
+        """
         return self.ports[port_idx]
 
     def get_nic_uuid_by_idx(self, node_idx: int, nic_idx: int) -> str:
+        """
+        Get the NIC UUID corresponding to the given node and NIC indices.
+
+        :param node_idx: The index of the node.
+        :type node_idx: int
+        :param nic_idx: The index of the NIC on the node.
+        :type nic_idx: int
+        :return: The NIC UUID.
+        :rtype: 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())
@@ -624,6 +826,31 @@ class ActionManager:
 
     @classmethod
     def from_config(cls, session: "PrimaiteSession", cfg: Dict) -> "ActionManager":
+        """
+        Construct an ActionManager from a config definition.
+
+        The action space config supports the following three sections:
+            1. ``action_list``
+                ``action_list`` contians a list action components which need to be included in the action space.
+                Each action component has a ``type`` which maps to a subclass of AbstractAction, and additional options
+                which will be passed to the action class's __init__ method during initialisation.
+            2. ``action_map``
+                Since the agent uses a discrete action space which acts as a flattened version of the component-based
+                action space, action_map provides a mapping between an integer (chosen by the agent) and a meaningful
+                action and values of parameters. For example action 0 can correspond to do nothing, action 1 can
+                correspond to "NODE_SERVICE_SCAN" with ``node_id=1`` and ``service_id=1``, action 2 can be "
+            3. ``options``
+                ``options`` contains a dictionary of options which are passed to the ActionManager's __init__ method.
+                These options are used to calculate the shape of the action space, and to provide additional information
+                to the ActionManager which is required to convert the agent's action choice into a CAOS request.
+
+        :param session: The Primaite Session to which the agent belongs.
+        :type session: PrimaiteSession
+        :param cfg: The action space config.
+        :type cfg: Dict
+        :return: The constructed ActionManager.
+        :rtype: ActionManager
+        """
         obj = cls(
             session=session,
             actions=cfg["action_list"],

src/primaite/game/agent/interface.py

@@ -1,6 +1,4 @@
-# TODO: remove this comment... This is just here to point out that I've named this 'actor' rather than 'agent'
-# 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
+"""Interface for agents."""
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple, TypeAlias, Union
 

src/primaite/game/agent/rewards.py

@@ -1,8 +1,9 @@
-from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
-
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Tuple, TYPE_CHECKING
+
 from primaite import getLogger
+from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
+
 _LOGGER = getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -10,14 +11,13 @@ if TYPE_CHECKING:
 
 
 class AbstractReward:
-
     @abstractmethod
     def calculate(self, state: Dict) -> float:
         return 0.0
 
     @classmethod
     @abstractmethod
-    def from_config(cls, config:dict, session:"PrimaiteSession") -> "AbstractReward":
+    def from_config(cls, config: dict, session: "PrimaiteSession") -> "AbstractReward":
         return cls()
 
 
@@ -26,16 +26,26 @@ class DummyReward(AbstractReward):
         return 0.0
 
     @classmethod
-    def from_config(cls, config: dict, session:"PrimaiteSession") -> "DummyReward":
+    def from_config(cls, config: dict, session: "PrimaiteSession") -> "DummyReward":
         return cls()
 
+
 class DatabaseFileIntegrity(AbstractReward):
-    def __init__(self, node_uuid:str, folder_name:str, file_name:str) -> None:
-        self.location_in_state = ["network", "nodes", node_uuid, "file_system", "folders",folder_name, "files", file_name]
+    def __init__(self, node_uuid: str, folder_name: str, file_name: str) -> None:
+        self.location_in_state = [
+            "network",
+            "nodes",
+            node_uuid,
+            "file_system",
+            "folders",
+            folder_name,
+            "files",
+            file_name,
+        ]
 
     def calculate(self, state: Dict) -> float:
         database_file_state = access_from_nested_dict(state, self.location_in_state)
-        health_status = database_file_state['health_status']
+        health_status = database_file_state["health_status"]
         if health_status == "corrupted":
             return -1
         elif health_status == "good":
@@ -48,29 +58,39 @@ class DatabaseFileIntegrity(AbstractReward):
         node_ref = config.get("node_ref")
         folder_name = config.get("folder_name")
         file_name = config.get("file_name")
-        if not (node_ref):
-            _LOGGER.error(f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified")
-            return DummyReward() #TODO: better error handling
+        if not node_ref:
+            _LOGGER.error(
+                f"{cls.__name__} could not be initialised from config because node_ref parameter was not specified"
+            )
+            return DummyReward()  # TODO: better error handling
         if not folder_name:
-            _LOGGER.error(f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified")
-            return DummyReward() # TODO: better error handling
+            _LOGGER.error(
+                f"{cls.__name__} could not be initialised from config because folder_name parameter was not specified"
+            )
+            return DummyReward()  # TODO: better error handling
         if not file_name:
-            _LOGGER.error(f"{cls.__name__} could not be initialised from config because file_name parameter was not specified")
-            return DummyReward() # TODO: better error handling
+            _LOGGER.error(
+                f"{cls.__name__} could not be initialised from config because file_name parameter was not specified"
+            )
+            return DummyReward()  # TODO: better error handling
         node_uuid = session.ref_map_nodes[node_ref]
         if not node_uuid:
-            _LOGGER.error(f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation")
-            return DummyReward() # TODO: better error handling
+            _LOGGER.error(
+                f"{cls.__name__} could not be initialised from config because the referenced node could not be found in the simulation"
+            )
+            return DummyReward()  # TODO: better error handling
+
+        return cls(node_uuid=node_uuid, folder_name=folder_name, file_name=file_name)
 
-        return cls(node_uuid = node_uuid, folder_name=folder_name, file_name=file_name)
 
 class WebServer404Penalty(AbstractReward):
-    def __init__(self, node_uuid:str, service_uuid:str) -> None:
-        self.location_in_state = ['network','nodes', node_uuid, 'services', service_uuid]
+    def __init__(self, node_uuid: str, service_uuid: str) -> None:
+        self.location_in_state = ["network", "nodes", node_uuid, "services", service_uuid]
 
     def calculate(self, state: Dict) -> float:
         web_service_state = access_from_nested_dict(state, self.location_in_state)
-        most_recent_return_code = web_service_state['most_recent_return_code']
+        most_recent_return_code = web_service_state["most_recent_return_code"]
+        # TODO: reward needs to use the current web state. Observation should return web state at the time of last scan.
         if most_recent_return_code == 200:
             return 1
         elif most_recent_return_code == 404:
@@ -85,13 +105,13 @@ class WebServer404Penalty(AbstractReward):
         if not (node_ref and service_ref):
             msg = f"{cls.__name__} could not be initialised from config because node_ref and service_ref were not found in reward config."
             _LOGGER.warn(msg)
-            return DummyReward() #TODO: should we error out with incorrect inputs? Probably!
+            return DummyReward()  # TODO: should we error out with incorrect inputs? Probably!
         node_uuid = session.ref_map_nodes[node_ref]
         service_uuid = session.ref_map_services[service_ref].uuid
         if not (node_uuid and service_uuid):
             msg = f"{cls.__name__} could not be initialised because node {node_ref} and service {service_ref} were not found in the simulator."
             _LOGGER.warn(msg)
-            return DummyReward() # TODO: consider erroring here as well
+            return DummyReward()  # TODO: consider erroring here as well
 
         return cls(node_uuid=node_uuid, service_uuid=service_uuid)
 
@@ -101,13 +121,13 @@ class RewardFunction:
         "DUMMY": DummyReward,
         "DATABASE_FILE_INTEGRITY": DatabaseFileIntegrity,
         "WEB_SERVER_404_PENALTY": WebServer404Penalty,
-        }
+    }
 
     def __init__(self):
         self.reward_components: List[Tuple[AbstractReward, float]] = []
         "attribute reward_components keeps track of reward components and the weights assigned to each."
 
-    def regsiter_component(self, component:AbstractReward, weight:float=1.0) -> None:
+    def regsiter_component(self, component: AbstractReward, weight: float = 1.0) -> None:
         self.reward_components.append((component, weight))
 
     def calculate(self, state: Dict) -> float:
@@ -124,8 +144,8 @@ class RewardFunction:
 
         for rew_component_cfg in config["reward_components"]:
             rew_type = rew_component_cfg["type"]
-            weight = rew_component_cfg.get("weight",1.0)
+            weight = rew_component_cfg.get("weight", 1.0)
             rew_class = cls.__rew_class_identifiers[rew_type]
-            rew_instance = rew_class.from_config(config=rew_component_cfg.get('options',{}), session=session)
+            rew_instance = rew_class.from_config(config=rew_component_cfg.get("options", {}), session=session)
             new.regsiter_component(component=rew_instance, weight=weight)
         return new