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

Add service patch and fix other bugs

Browse Source
This commit is contained in:
Marek Wolan
2023-12-21 15:07:41 +00:00
parent e33f74e3f2
commit 96f8435c5e
6 changed files with 135 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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

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