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

Merged PR 240: Add service patch action and fix some clangers.

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
Unit tests passing
Running primaite session works without errors
Ran primaite session on UC2 in debug mode and inspected that correct changes are being reflected after actions are applied.

## Checklist
- [x] PR is linked to a **work item**
- [x] **acceptance criteria** of linked ticket are met
- [x] performed **self-review** of the code
- [ ] 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**
- [ ] ran **pre-commit** checks for code style
- [ ] attended to any **TO-DOs** left in the code

Related work items: #2143
This commit is contained in:
Marek Wolan
2024-01-05 09:56:33 +00:00
parent 75f732dacf 2135bdcd10
commit bc08db6b7e
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.0b2
3.0.0b3dev

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

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

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

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

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