diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index 24f9945d..a764703b 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/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 - service_id: 1 + node_id: 1 + service_id: 0 # stop webapp service 2: action: NODE_SERVICE_STOP options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 # start webapp service 3: action: "NODE_SERVICE_START" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 4: action: "NODE_SERVICE_PAUSE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 5: action: "NODE_SERVICE_RESUME" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 6: action: "NODE_SERVICE_RESTART" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 7: action: "NODE_SERVICE_DISABLE" options: - node_id: 2 - service_id: 1 + node_id: 1 + service_id: 0 8: action: "NODE_SERVICE_ENABLE" options: - node_id: 2 - service_id: 1 - 9: + node_id: 1 + service_id: 0 + 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 - folder_id: 1 - file_id: 1 + node_id: 2 + service_id: 0 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 - 19: + node_id: 2 + 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 diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index 8eed3ba4..4c47bfaa 100644 --- a/src/primaite/game/agent/actions.py +++ b/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()) diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index 8c32f41d..d2db4bea 100644 --- a/src/primaite/game/game.py +++ b/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 diff --git a/src/primaite/simulator/file_system/folder.py b/src/primaite/simulator/file_system/folder.py index fd18e154..d4e72f63 100644 --- a/src/primaite/simulator/file_system/folder.py +++ b/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.") diff --git a/src/primaite/simulator/sim_container.py b/src/primaite/simulator/sim_container.py index c529ed04..db8a718c 100644 --- a/src/primaite/simulator/sim_container.py +++ b/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) diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 87802a7b..048e6fec 100644 --- a/src/primaite/simulator/system/software.py +++ b/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): """