diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 15c44821..fd3f369d 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1455,74 +1455,6 @@ class Node(SimComponent): else: return - def install_service(self, service: Service) -> None: - """ - Install a service on this node. - - :param service: Service instance that has not been installed on any node yet. - :type service: Service - """ - if service in self: - _LOGGER.warning(f"Can't add service {service.name} to node {self.hostname}. It's already installed.") - return - self.services[service.uuid] = service - service.parent = self - service.install() # Perform any additional setup, such as creating files for this service on the node. - self.sys_log.info(f"Installed service {service.name}") - _LOGGER.debug(f"Added service {service.name} to node {self.hostname}") - self._service_request_manager.add_request(service.name, RequestType(func=service._request_manager)) - - def uninstall_service(self, service: Service) -> None: - """ - Uninstall and completely remove service from this node. - - :param service: Service object that is currently associated with this node. - :type service: Service - """ - if service not in self: - _LOGGER.warning(f"Can't remove service {service.name} from node {self.hostname}. It's not installed.") - return - service.uninstall() # Perform additional teardown, such as removing files or restarting the machine. - self.services.pop(service.uuid) - service.parent = None - self.sys_log.info(f"Uninstalled service {service.name}") - self._service_request_manager.remove_request(service.name) - - 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.name} to node {self.hostname}. It's already installed." - ) - return - self.applications[application.uuid] = application - application.parent = self - self.sys_log.info(f"Installed application {application.name}") - _LOGGER.debug(f"Added application {application.name} to node {self.hostname}") - self._application_request_manager.add_request(application.name, 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.name} from node {self.hostname}. It's not installed." - ) - return - self.applications.pop(application.uuid) - application.parent = None - self.sys_log.info(f"Uninstalled application {application.name}") - self._application_request_manager.remove_request(application.name) - def _shut_down_actions(self): """Actions to perform when the node is shut down.""" # Turn off all the services in the node diff --git a/src/primaite/simulator/system/core/software_manager.py b/src/primaite/simulator/system/core/software_manager.py index e2266c2d..9c4d7cf6 100644 --- a/src/primaite/simulator/system/core/software_manager.py +++ b/src/primaite/simulator/system/core/software_manager.py @@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING, Union from prettytable import MARKDOWN, PrettyTable +from primaite.simulator.core import RequestType from primaite.simulator.file_system.file_system import FileSystem from primaite.simulator.network.transmission.data_link_layer import Frame from primaite.simulator.network.transmission.network_layer import IPProtocol @@ -20,9 +21,7 @@ if TYPE_CHECKING: from primaite.simulator.system.services.arp.arp import ARP from primaite.simulator.system.services.icmp.icmp import ICMP -from typing import Type, TypeVar - -IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware) +from typing import Type class SoftwareManager: @@ -51,7 +50,7 @@ class SoftwareManager: 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] = {} + self._software_class_to_name_map: Dict[Type[IOSoftware], str] = {} self.port_protocol_mapping: Dict[Tuple[Port, IPProtocol], Union[Service, Application]] = {} self.sys_log: SysLog = sys_log self.file_system: FileSystem = file_system @@ -104,33 +103,34 @@ class SoftwareManager: return True return False - def install(self, software_class: Type[IOSoftwareClass]): + def install(self, software_class: Type[IOSoftware]): """ Install an Application or Service. :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.warning(f"Cannot install {software_class} as it is already installed") return software = software_class( software_manager=self, sys_log=self.sys_log, file_system=self.file_system, dns_server=self.dns_server ) + software.parent = self.node if isinstance(software, Application): - software.install() + self.node.applications[software.uuid] = software + self.node._application_request_manager.add_request( + software.name, RequestType(func=software._request_manager) + ) + elif isinstance(software, Service): + self.node.services[software.uuid] = software + self.node._service_request_manager.add_request(software.name, RequestType(func=software._request_manager)) + software.install() software.software_manager = self self.software[software.name] = software self.port_protocol_mapping[(software.port, software.protocol)] = software 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) + self.node.sys_log.info(f"Installed {software.name}") def uninstall(self, software_name: str): """ @@ -138,25 +138,31 @@ class SoftwareManager: :param software_name: The software name. """ - if software_name in self.software: - self.software[software_name].uninstall() - 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) - for key, value in self.port_protocol_mapping.items(): - if value.name == software_name: - self.port_protocol_mapping.pop(key) - break - for key, value in self._software_class_to_name_map.items(): - if value == software_name: - self._software_class_to_name_map.pop(key) - break - del software - self.sys_log.info(f"Uninstalled {software_name}") + if software_name not in self.software: + self.sys_log.error(f"Cannot uninstall {software_name} as it is not installed") return - self.sys_log.error(f"Cannot uninstall {software_name} as it is not installed") + + self.software[software_name].uninstall() + software = self.software.pop(software_name) # noqa + if isinstance(software, Application): + self.node.applications.pop(software.uuid) + self.node._application_request_manager.remove_request(software.name) + elif isinstance(software, Service): + self.node.services.pop(software.uuid) + software.uninstall() + self.node._service_request_manager.remove_request(software.name) + software.parent = None + for key, value in self.port_protocol_mapping.items(): + if value.name == software_name: + self.port_protocol_mapping.pop(key) + break + for key, value in self._software_class_to_name_map.items(): + if value == software_name: + self._software_class_to_name_map.pop(key) + break + del software + self.sys_log.info(f"Uninstalled {software_name}") + return def send_internal_payload(self, target_software: str, payload: Any): """ diff --git a/tests/conftest.py b/tests/conftest.py index 54519e2b..ca704461 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -37,14 +37,14 @@ ACTION_SPACE_NODE_ACTION_VALUES = 1 _LOGGER = getLogger(__name__) -class TestService(Service): +class DummyService(Service): """Test Service class""" def describe_state(self) -> Dict: return super().describe_state() def __init__(self, **kwargs): - kwargs["name"] = "TestService" + kwargs["name"] = "DummyService" kwargs["port"] = Port.HTTP kwargs["protocol"] = IPProtocol.TCP super().__init__(**kwargs) @@ -75,15 +75,15 @@ def uc2_network() -> Network: @pytest.fixture(scope="function") -def service(file_system) -> TestService: - return TestService( - name="TestService", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="test_service") +def service(file_system) -> DummyService: + return DummyService( + name="DummyService", port=Port.ARP, file_system=file_system, sys_log=SysLog(hostname="dummy_service") ) @pytest.fixture(scope="function") def service_class(): - return TestService + return DummyService @pytest.fixture(scope="function") diff --git a/tests/integration_tests/component_creation/test_action_integration.py b/tests/integration_tests/component_creation/test_action_integration.py index a6f09436..7bdc80fc 100644 --- a/tests/integration_tests/component_creation/test_action_integration.py +++ b/tests/integration_tests/component_creation/test_action_integration.py @@ -22,8 +22,7 @@ def test_passing_actions_down(monkeypatch) -> None: for n in [pc1, pc2, srv, s1]: sim.network.add_node(n) - database_service = DatabaseService(file_system=srv.file_system) - srv.install_service(database_service) + srv.software_manager.install(DatabaseService) downloads_folder = pc1.file_system.create_folder("downloads") pc1.file_system.create_file("bermuda_triangle.png", folder_name="downloads") diff --git a/tests/integration_tests/system/test_service_on_node.py b/tests/integration_tests/system/test_service_on_node.py index 15dbaf1d..cf9728ce 100644 --- a/tests/integration_tests/system/test_service_on_node.py +++ b/tests/integration_tests/system/test_service_on_node.py @@ -23,7 +23,7 @@ def populated_node( server.power_on() server.software_manager.install(service_class) - service = server.software_manager.software.get("TestService") + service = server.software_manager.software.get("DummyService") service.start() return server, service @@ -42,7 +42,7 @@ def test_service_on_offline_node(service_class): computer.power_on() computer.software_manager.install(service_class) - service: Service = computer.software_manager.software.get("TestService") + service: Service = computer.software_manager.software.get("DummyService") computer.power_off() diff --git a/tests/integration_tests/test_simulation/test_request_response.py b/tests/integration_tests/test_simulation/test_request_response.py index a9f0b58d..95634cf1 100644 --- a/tests/integration_tests/test_simulation/test_request_response.py +++ b/tests/integration_tests/test_simulation/test_request_response.py @@ -13,7 +13,7 @@ from primaite.simulator.network.hardware.node_operating_state import NodeOperati from primaite.simulator.network.hardware.nodes.host.host_node import HostNode from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router from primaite.simulator.network.transmission.transport_layer import Port -from tests.conftest import DummyApplication, TestService +from tests.conftest import DummyApplication, DummyService def test_successful_node_file_system_creation_request(example_network): @@ -61,7 +61,7 @@ def test_successful_application_requests(example_network): def test_successful_service_requests(example_network): net = example_network server_1 = net.get_node_by_hostname("server_1") - server_1.software_manager.install(TestService) + server_1.software_manager.install(DummyService) # Careful: the order here is important, for example we cannot run "stop" unless we run "start" first for verb in [ @@ -77,7 +77,7 @@ def test_successful_service_requests(example_network): "scan", "fix", ]: - resp_1 = net.apply_request(["node", "server_1", "service", "TestService", verb]) + resp_1 = net.apply_request(["node", "server_1", "service", "DummyService", verb]) assert resp_1 == RequestResponse(status="success", data={}) server_1.apply_timestep(timestep=1) server_1.apply_timestep(timestep=1) diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py index 9b37ac80..44c5c781 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_node_actions.py @@ -7,6 +7,7 @@ from primaite.simulator.file_system.folder import Folder from primaite.simulator.network.hardware.base import Node, NodeOperatingState from primaite.simulator.network.hardware.nodes.host.computer import Computer from primaite.simulator.system.software import SoftwareHealthState +from tests.conftest import DummyApplication, DummyService @pytest.fixture @@ -47,7 +48,7 @@ def test_node_shutdown(node): assert node.operating_state == NodeOperatingState.OFF -def test_node_os_scan(node, service, application): +def test_node_os_scan(node): """Test OS Scanning.""" node.operating_state = NodeOperatingState.ON @@ -55,13 +56,15 @@ def test_node_os_scan(node, service, application): # TODO implement processes # add services to node + node.software_manager.install(DummyService) + service = node.software_manager.software.get("DummyService") service.set_health_state(SoftwareHealthState.COMPROMISED) - node.install_service(service=service) assert service.health_state_visible == SoftwareHealthState.UNUSED # add application to node + node.software_manager.install(DummyApplication) + application = node.software_manager.software.get("DummyApplication") application.set_health_state(SoftwareHealthState.COMPROMISED) - node.install_application(application=application) assert application.health_state_visible == SoftwareHealthState.UNUSED # add folder and file to node @@ -91,7 +94,7 @@ def test_node_os_scan(node, service, application): assert file2.visible_health_status == FileSystemItemHealthStatus.CORRUPT -def test_node_red_scan(node, service, application): +def test_node_red_scan(node): """Test revealing to red""" node.operating_state = NodeOperatingState.ON @@ -99,12 +102,14 @@ def test_node_red_scan(node, service, application): # TODO implement processes # add services to node - node.install_service(service=service) + node.software_manager.install(DummyService) + service = node.software_manager.software.get("DummyService") assert service.revealed_to_red is False # add application to node + node.software_manager.install(DummyApplication) + application = node.software_manager.software.get("DummyApplication") application.set_health_state(SoftwareHealthState.COMPROMISED) - node.install_application(application=application) assert application.revealed_to_red is False # add folder and file to node