Add service patch and fix other bugs

src/primaite/config/_package_data/example_config.yaml

@@ -169,6 +169,7 @@ agents:
         - type: NODE_SERVICE_RESTART
         - type: NODE_SERVICE_DISABLE
         - type: NODE_SERVICE_ENABLE
+        - type: NODE_SERVICE_PATCH
         - type: NODE_FILE_SCAN
         - type: NODE_FILE_CHECKHASH
         - type: NODE_FILE_DELETE
@@ -199,111 +200,110 @@ agents:
           1:
             action: NODE_SERVICE_SCAN
             options:
-              node_id: 2
+              node_id: 1
-              service_id: 1
+              service_id: 0
           # stop webapp service
           2:
             action: NODE_SERVICE_STOP
             options:
-              node_id: 2
+              node_id: 1
-              service_id: 1
+              service_id: 0
           # start webapp service
           3:
             action: "NODE_SERVICE_START"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
           4:
             action: "NODE_SERVICE_PAUSE"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
           5:
             action: "NODE_SERVICE_RESUME"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
           6:
             action: "NODE_SERVICE_RESTART"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
           7:
             action: "NODE_SERVICE_DISABLE"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
           8:
             action: "NODE_SERVICE_ENABLE"
             options:
-                node_id: 2
+                node_id: 1
-                service_id: 1
+                service_id: 0
-          9:
+          9: # check database.db file
             action: "NODE_FILE_SCAN"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
-                file_id: 1
+                file_id: 0
           10:
             action: "NODE_FILE_CHECKHASH"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
-                file_id: 1
+                file_id: 0
           11:
             action: "NODE_FILE_DELETE"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
-                file_id: 1
+                file_id: 0
           12:
             action: "NODE_FILE_REPAIR"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
-                file_id: 1
+                file_id: 0
           13:
-            action: "NODE_FILE_RESTORE"
+            action: "NODE_SERVICE_PATCH"
             options:
-                node_id: 3
+                node_id: 2
-                folder_id: 1
+                service_id: 0
-                file_id: 1
           14:
             action: "NODE_FOLDER_SCAN"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
           15:
             action: "NODE_FOLDER_CHECKHASH"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
           16:
             action: "NODE_FOLDER_REPAIR"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
           17:
             action: "NODE_FOLDER_RESTORE"
             options:
-                node_id: 3
+                node_id: 2
                 folder_id: 1
           18:
             action: "NODE_OS_SCAN"
             options:
-                node_id: 3
+                node_id: 2
-          19:
+          19: # shutdown client 1
             action: "NODE_SHUTDOWN"
             options:
-                node_id: 6
+                node_id: 5
           20:
             action: "NODE_STARTUP"
             options:
-                node_id: 6
+                node_id: 5
           21:
             action: "NODE_RESET"
             options:
-                node_id: 6
+                node_id: 5
           22:
             action: "NETWORK_ACL_ADDRULE"
             options:
@@ -407,93 +407,94 @@ agents:
           38:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 1
+                node_id: 0
                 nic_id: 1
           39:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 1
+                node_id: 0
                 nic_id: 1
           40:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 2
+                node_id: 1
                 nic_id: 1
           41:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 2
+                node_id: 1
                 nic_id: 1
           42:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 3
+                node_id: 2
                 nic_id: 1
           43:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 3
+                node_id: 2
                 nic_id: 1
           44:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 4
+                node_id: 3
                 nic_id: 1
           45:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 4
+                node_id: 3
                 nic_id: 1
           46:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 5
+                node_id: 4
                 nic_id: 1
           47:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 5
+                node_id: 4
                 nic_id: 1
           48:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 5
+                node_id: 4
                 nic_id: 2
           49:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 5
+                node_id: 4
                 nic_id: 2
           50:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 6
+                node_id: 5
                 nic_id: 1
           51:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 6
+                node_id: 5
                 nic_id: 1
           52:
             action: "NETWORK_NIC_DISABLE"
             options:
-                node_id: 7
+                node_id: 6
                 nic_id: 1
           53:
             action: "NETWORK_NIC_ENABLE"
             options:
-                node_id: 7
+                node_id: 6
                 nic_id: 1
 
 
       options:
         nodes:
-        - node_ref: router_1
-        - node_ref: switch_1
-        - node_ref: switch_2
         - node_ref: domain_controller
         - node_ref: web_server
+          services:
+          - service_ref: web_server_web_service
         - node_ref: database_server
+          services:
+          - service_ref: database_service
         - node_ref: backup_server
         - node_ref: security_suite
         - node_ref: client_1

src/primaite/game/agent/actions.py

@@ -89,7 +89,7 @@ class NodeServiceAbstractAction(AbstractAction):
         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, "service", service_uuid, self.verb]
 
 
 class NodeServiceScanAction(NodeServiceAbstractAction):
@@ -156,6 +156,14 @@ class NodeServiceEnableAction(NodeServiceAbstractAction):
         self.verb: str = "enable"
 
 
+class NodeServicePatchAction(NodeServiceAbstractAction):
+    """Action which patches 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: str = "patch"
+
+
 class NodeApplicationAbstractAction(AbstractAction):
     """
     Base class for application actions.
@@ -262,7 +270,7 @@ class NodeFileAbstractAction(AbstractAction):
         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, "file", file_uuid, self.verb]
 
 
 class NodeFileScanAction(NodeFileAbstractAction):
@@ -566,6 +574,7 @@ class ActionManager:
         "NODE_SERVICE_RESTART": NodeServiceRestartAction,
         "NODE_SERVICE_DISABLE": NodeServiceDisableAction,
         "NODE_SERVICE_ENABLE": NodeServiceEnableAction,
+        "NODE_SERVICE_PATCH": NodeServicePatchAction,
         "NODE_APPLICATION_EXECUTE": NodeApplicationExecuteAction,
         "NODE_FILE_SCAN": NodeFileScanAction,
         "NODE_FILE_CHECKHASH": NodeFileCheckhashAction,
@@ -594,6 +603,7 @@ class ActionManager:
         actions: List[str],  # stores list of actions available to agent
         node_uuids: List[str],  # allows mapping index to node
         application_uuids: List[List[str]],  # allows mapping index to application
+        service_uuids: List[List[str]],  # allows mapping index to service
         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
@@ -635,6 +645,7 @@ class ActionManager:
         self.game: "PrimaiteGame" = game
         self.node_uuids: List[str] = node_uuids
         self.application_uuids: List[List[str]] = application_uuids
+        self.service_uuids: List[List[str]] = service_uuids
         self.protocols: List[str] = protocols
         self.ports: List[str] = ports
 
@@ -804,6 +815,11 @@ class ActionManager:
         :return: The UUID of the service. Or None if the node has fewer services than the given index.
         :rtype: Optional[str]
         """
+        # if a mapping was specified, use that mapping, otherwise just use the list of all installed services
+        if self.service_uuids:
+            if self.service_uuids[node_idx]:
+                return self.service_uuids[node_idx][service_idx]
+
         node_uuid = self.get_node_uuid_by_idx(node_idx)
         node = self.game.simulation.network.nodes[node_uuid]
         service_uuids = list(node.services.keys())

src/primaite/game/game.py

@@ -361,6 +361,7 @@ class PrimaiteGame:
             # CREATE ACTION SPACE
             action_space_cfg["options"]["node_uuids"] = []
             action_space_cfg["options"]["application_uuids"] = []
+            action_space_cfg["options"]["service_uuids"] = []
 
             # if a list of nodes is defined, convert them from node references to node UUIDs
             for action_node_option in action_space_cfg.get("options", {}).pop("nodes", {}):
@@ -375,10 +376,21 @@ class PrimaiteGame:
                         #  node_uuid, whereas here the application gets added by uuid.
                         application_uuid = game.ref_map_applications[application_option["application_ref"]]
                         node_application_uuids.append(application_uuid)
-
                     action_space_cfg["options"]["application_uuids"].append(node_application_uuids)
+
                 else:
                     action_space_cfg["options"]["application_uuids"].append([])
+
+                if "services" in action_node_option:
+                    node_service_uuids = []
+                    for service_option in action_node_option["services"]:
+                        service_uuid = game.ref_map_services[service_option["service_ref"]]
+                        node_service_uuids.append(service_uuid)
+                    action_space_cfg["options"]["service_uuids"].append(node_service_uuids)
+
+                else:
+                    action_space_cfg["options"]["service_uuids"].append([])
+
             # Each action space can potentially have a different list of nodes that it can apply to. Therefore,
             # we will pass node_uuids as a part of the action space config.
             # However, it's not possible to specify the node uuids directly in the config, as they are generated

src/primaite/simulator/file_system/folder.py

@@ -102,6 +102,11 @@ class Folder(FileSystemItemABC):
             name="delete",
             request_type=RequestType(func=lambda request, context: self.remove_file_by_id(file_uuid=request[0])),
         )
+        self._file_request_manager = RequestManager()
+        rm.add_request(
+            name="file",
+            request_type=RequestType(func=lambda request, context: self._file_request_manager),
+        )
         return rm
 
     def describe_state(self) -> Dict:
@@ -254,6 +259,7 @@ class Folder(FileSystemItemABC):
         # add to list
         self.files[file.uuid] = file
         self._files_by_name[file.name] = file
+        self._file_request_manager.add_request(file.uuid, RequestType(func=file._request_manager))
         file.folder = self
 
     def remove_file(self, file: Optional[File]):
@@ -273,6 +279,7 @@ class Folder(FileSystemItemABC):
             self.deleted_files[file.uuid] = file
             file.delete()
             self.sys_log.info(f"Removed file {file.name} (id: {file.uuid})")
+            self._file_request_manager.remove_request(file.uuid)
         else:
             _LOGGER.debug(f"File with UUID {file.uuid} was not found.")
 

src/primaite/simulator/sim_container.py

@@ -55,3 +55,9 @@ class Simulation(SimComponent):
             }
         )
         return state
+
+    def apply_timestep(self, timestep: int) -> None:
+        """Apply a timestep to the simulation."""
+        super().apply_timestep(timestep)
+        self.network.apply_timestep(timestep)
+        # self.domain.apply_timestep(timestep)

src/primaite/simulator/system/software.py

@@ -89,6 +89,10 @@ class Software(SimComponent):
     "The FileSystem of the Node the Software is installed on."
     folder: Optional[Folder] = None
     "The folder on the file system the Software uses."
+    patching_duration: int = 2
+    "The number of ticks it takes to patch the software."
+    _patching_countdown: Optional[int] = None
+    "Current number of ticks left to patch the software."
 
     def set_original_state(self):
         """Sets the original state."""
@@ -111,6 +115,12 @@ class Software(SimComponent):
                 func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED),
             ),
         )
+        rm.add_request(
+            "patch",
+            RequestType(
+                func=lambda request, context: self.patch(),
+            ),
+        )
         rm.add_request("scan", RequestType(func=lambda request, context: self.scan()))
         return rm
 
@@ -181,10 +191,33 @@ class Software(SimComponent):
         """Update the observed health status to match the actual health status."""
         self.health_state_visible = self.health_state_actual
 
+    def patch(self) -> None:
+        """Perform a patch on the software."""
+        self._patching_countdown = self.patching_duration
+        self.set_health_state(SoftwareHealthState.PATCHING)
+
+    def _update_patch_status(self) -> None:
+        """Update the patch status of the software."""
+        self._patching_countdown -= 1
+        if self._patching_countdown <= 0:
+            self.set_health_state(SoftwareHealthState.GOOD)
+            self._patching_countdown = None
+            self.patching_count += 1
+
     def reveal_to_red(self) -> None:
         """Reveals the software to the red agent."""
         self.revealed_to_red = True
 
+    def apply_timestep(self, timestep: int) -> None:
+        """
+        Apply a single timestep to the software.
+
+        :param timestep: The current timestep of the simulation.
+        """
+        super().apply_timestep(timestep)
+        if self.health_state_actual == SoftwareHealthState.PATCHING:
+            self._update_patch_status()
+
 
 class IOSoftware(Software):
     """