From d4eee36b7bc921af0791e9b1a4aa062a0b97111d Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 23 Oct 2023 17:23:14 +0100 Subject: [PATCH] Fix software registration for game layer and simulator interface --- example_config.yaml | 2 +- src/primaite/game/agent/observations.py | 2 +- .../simulator/network/hardware/base.py | 36 ++++++++++++++++++- .../system/applications/application.py | 2 +- .../simulator/system/core/software_manager.py | 15 ++++++++ .../system/services/web_server/web_server.py | 20 ++++++++++- src/primaite/simulator/system/software.py | 6 ++-- 7 files changed, 75 insertions(+), 8 deletions(-) diff --git a/example_config.yaml b/example_config.yaml index f3d8dc10..00beaa1e 100644 --- a/example_config.yaml +++ b/example_config.yaml @@ -457,7 +457,7 @@ game_config: weight: 0.5 options: node_ref: web_server - service_ref: web_server_database_client + service_ref: web_server_web_service agent_settings: diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index af398fc9..35fe8ac5 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -139,7 +139,7 @@ class ServiceObservation(AbstractObservation): service_state = access_from_nested_dict(state, self.where) if service_state is NOT_PRESENT_IN_STATE: return self.default_observation - return {"operating_status": service_state["operating_state"], "health_status": service_state["health_status"]} + return {"operating_status": service_state["operating_state"], "health_status": service_state["health_state"]} @property def space(self) -> spaces.Space: diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 78ae228e..607e348b 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -938,6 +938,7 @@ class Node(SimComponent): kwargs["file_system"] = FileSystem(sys_log=kwargs["sys_log"], sim_root=kwargs["root"] / "fs") if not kwargs.get("software_manager"): kwargs["software_manager"] = SoftwareManager( + parent_node=self, sys_log=kwargs.get("sys_log"), session_manager=kwargs.get("session_manager"), file_system=kwargs.get("file_system"), @@ -1199,7 +1200,8 @@ class Node(SimComponent): self._service_request_manager.add_request(service.uuid, RequestType(func=service._request_manager)) def uninstall_service(self, service: Service) -> None: - """Uninstall and completely remove service from this node. + """ + Uninstall and completely remove service from this node. :param service: Service object that is currently associated with this node. :type service: Service @@ -1214,6 +1216,38 @@ class Node(SimComponent): _LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}") self._service_request_manager.remove_request(service.uuid) + def install_application(self, application: Application) -> None: + """ + Install an application on this node. + + :param application: Application instance that has not been installed on any node yet. + :type application: Application + """ + if application in self: + _LOGGER.warning(f"Can't add application {application.uuid} to node {self.uuid}. It's already installed.") + return + self.applications[application.uuid] = application + application.parent = self + self.sys_log.info(f"Installed application {application.name}") + _LOGGER.info(f"Added application {application.uuid} to node {self.uuid}") + self._application_request_manager.add_request(application.uuid, RequestType(func=application._request_manager)) + + def uninstall_application(self, application: Application) -> None: + """ + Uninstall and completely remove application from this node. + + :param application: Application object that is currently associated with this node. + :type application: Application + """ + if application not in self: + _LOGGER.warning(f"Can't remove application {application.uuid} from node {self.uuid}. It's not installed.") + return + self.applications.pop(application.uuid) + application.parent = None + self.sys_log.info(f"Uninstalled application {application.name}") + _LOGGER.info(f"Removed application {application.uuid} from node {self.uuid}") + self._application_request_manager.remove_request(application.uuid) + def __contains__(self, item: Any) -> bool: if isinstance(item, Service): return item.uuid in self.services diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 69b64aac..e3da6f01 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -45,7 +45,7 @@ class Application(IOSoftware): state = super().describe_state() state.update( { - "opearting_state": self.operating_state.name, + "opearting_state": self.operating_state.value, "execution_control_status": self.execution_control_status, "num_executions": self.num_executions, "groups": list(self.groups), diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index 973b17b4..8b8fe599 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -14,6 +14,7 @@ from primaite.simulator.system.software import IOSoftware if TYPE_CHECKING: from primaite.simulator.system.core.session_manager import SessionManager from primaite.simulator.system.core.sys_log import SysLog + from primaite.simulator.network.hardware.base import Node from typing import Type, TypeVar @@ -25,6 +26,7 @@ class SoftwareManager: def __init__( self, + parent_node: "Node", session_manager: "SessionManager", sys_log: SysLog, file_system: FileSystem, @@ -35,6 +37,7 @@ class SoftwareManager: :param session_manager: The session manager handling network communications. """ + self.node = parent_node self.session_manager = session_manager self.software: Dict[str, Union[Service, Application]] = {} self._software_class_to_name_map: Dict[Type[IOSoftwareClass], str] = {} @@ -62,6 +65,8 @@ class SoftwareManager: :param software_class: The software class. """ + # TODO: Software manager and node itself both have an install method. Need to refactor to have more logical + # separation of concerns. if software_class in self._software_class_to_name_map: self.sys_log.info(f"Cannot install {software_class} as it is already installed") return @@ -77,6 +82,12 @@ class SoftwareManager: if isinstance(software, Application): software.operating_state = ApplicationOperatingState.CLOSED + # add the software to the node's registry after it has been fully initialized + if isinstance(software, Service): + self.node.install_service(software) + elif isinstance(software, Application): + self.node.install_application(software) + def uninstall(self, software_name: str): """ Uninstall an Application or Service. @@ -85,6 +96,10 @@ class SoftwareManager: """ if software_name in self.software: software = self.software.pop(software_name) # noqa + if isinstance(software, Application): + self.node.uninstall_application(software) + elif isinstance(software, Service): + self.node.uninstall_service(software) del software self.sys_log.info(f"Deleted {software_name}") return diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py index f63d5169..5957e4cb 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/src/primaite/simulator/system/services/web_server/web_server.py @@ -1,5 +1,5 @@ from ipaddress import IPv4Address -from typing import Any, Optional +from typing import Any, Dict, Optional from urllib.parse import urlparse from primaite.simulator.network.protocols.http import ( @@ -17,6 +17,23 @@ from primaite.simulator.system.services.service import Service class WebServer(Service): """Class used to represent a Web Server Service in simulation.""" + last_response_status_code: Optional[HttpStatusCode] = None + + def describe_state(self) -> Dict: + """ + Produce a dictionary describing the current state of this object. + + Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation. + + :return: Current state of this object and child objects. + :rtype: Dict + """ + state = super().describe_state() + state["last_response_status_code"] = ( + self.last_response_status_code.value if self.last_response_status_code else None + ) + return state + def __init__(self, **kwargs): kwargs["name"] = "WebServer" kwargs["protocol"] = IPProtocol.TCP @@ -66,6 +83,7 @@ class WebServer(Service): self.send(payload=response, session_id=session_id) # return true if response is OK + self.last_response_status_code = response.status_code return response.status_code == HttpStatusCode.OK def _handle_get_request(self, payload: HttpRequestPacket) -> HttpResponsePacket: diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 25f764e4..8cd13d1a 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -119,9 +119,9 @@ class Software(SimComponent): state = super().describe_state() state.update( { - "health_state": self.health_state_actual.name, - "health_state_red_view": self.health_state_visible.name, - "criticality": self.criticality.name, + "health_state": self.health_state_actual.value, + "health_state_red_view": self.health_state_visible.value, + "criticality": self.criticality.value, "patching_count": self.patching_count, "scanning_count": self.scanning_count, "revealed_to_red": self.revealed_to_red,