oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merged PR 241: Add service patch and fix apply timestep issues.

Browse Source 
## Summary
Replaces Action ID 13 in UC2 with SERVICE.PATCH.
Fixes issue where apply_timestep was not happening at all
Fixes issue where file actions were not correctly being forwarded

## Test process
All pytests passing including a new one for service patch.

## Checklist
- [x] PR is linked to a **work item**
- [x] **acceptance criteria** of linked ticket are met
- [x] performed **self-review** of the code
- [x] written **tests** for any new functionality added with this PR
- [ ] updated the **documentation** if this PR changes or adds functionality
- [ ] written/updated **design docs** if this PR implements new functionality
- [ ] updated the **change log**
- [x] ran **pre-commit** checks for code style
- [x] attended to any **TO-DOs** left in the code

Related work items: #2143
This commit is contained in:
Marek Wolan
2024-01-05 09:58:16 +00:00
parent 7d81b1cb07 3adf1f9f6a
commit 81734e7ad3
9 changed files with 264 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/VERSION
View File

@@ -1 +1 @@
3.0.0b2dev
3.0.0b4dev

115
src/primaite/config/_package_data/example_config.yaml
View File

@@ -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

230
src/primaite/config/_package_data/example_config_2_rl_agents.yaml
View File

@@ -165,6 +165,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
@@ -195,111 +196,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:
@@ -403,93 +403,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
@@ -597,6 +598,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
@@ -627,111 +629,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:
@@ -835,93 +836,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

20
src/primaite/game/agent/actions.py
View File

@@ -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())

14
src/primaite/game/game.py
View File

@@ -366,6 +366,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", {}):
@@ -380,10 +381,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

7
src/primaite/simulator/file_system/folder.py
View File

@@ -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.")

5
src/primaite/simulator/sim_container.py
View File

@@ -55,3 +55,8 @@ 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)

33
src/primaite/simulator/system/software.py
View File

@@ -91,6 +91,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."""
@@ -113,6 +117,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
@@ -183,10 +193,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):
"""

13
tests/unit_tests/_primaite/_simulator/_system/_services/test_service_actions.py
View File

@@ -78,3 +78,16 @@ def test_service_enable(service):
service.apply_request(["enable"])
assert service.operating_state == ServiceOperatingState.STOPPED
def test_service_patch(service):
"""Test that a service can be patched and that it takes two timesteps to complete."""
service.start()
assert service.health_state_actual == SoftwareHealthState.GOOD
service.apply_request(["patch"])
assert service.health_state_actual == SoftwareHealthState.PATCHING
service.apply_timestep(1)
assert service.health_state_actual == SoftwareHealthState.PATCHING
service.apply_timestep(2)
assert service.health_state_actual == SoftwareHealthState.GOOD