diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index fe3b5b15..41e16936 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -4,7 +4,7 @@ import re import secrets from enum import Enum from ipaddress import IPv4Address, IPv4Network -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from prettytable import PrettyTable @@ -994,6 +994,39 @@ class Node(SimComponent): elif frame.ip.protocol == IPProtocol.ICMP: self.icmp.process_icmp(frame=frame) + 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.uuid} to node {self.uuid}. It's already installed.") + return + service.parent = self + service.install() # Perform any additional setup, such as creating files for this service on the node. + _LOGGER.info(f"Added service {service.uuid} to node {self.uuid}") + + 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.uuid} from node {self.uuid}. 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 + _LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}") + + def __contains__(self, item: Any) -> bool: + if isinstance(item, Service): + return item.uuid in self.services + return None + class Switch(Node): """A class representing a Layer 2 network switch.""" diff --git a/src/primaite/simulator/system/services/database.py b/src/primaite/simulator/system/services/database.py index 720967e7..0d1de15c 100644 --- a/src/primaite/simulator/system/services/database.py +++ b/src/primaite/simulator/system/services/database.py @@ -1,3 +1,5 @@ +from typing import Dict + from primaite.simulator.file_system.file_system_file_type import FileSystemFileType from primaite.simulator.network.hardware.base import Node from primaite.simulator.system.services.service import Service @@ -6,11 +8,17 @@ from primaite.simulator.system.services.service import Service class DatabaseService(Service): """TODO.""" - def __init__(self, parent_node: Node, **kwargs): - super().__init__(**kwargs) - self._setup_files_on_node() + def describe_state(self) -> Dict: + """TODO.""" + return super().describe_state() - def _setup_files_on_node( + def install(self) -> None: + """Perform first time install on a node, creating necessary files.""" + super().install() + assert isinstance(self.parent, Node), "Database install can only happen after the db service is added to a node" + self._setup_files() + + def _setup_files( self, db_size: int = 1000, use_secondary_db_file: bool = False, @@ -30,6 +38,7 @@ class DatabaseService(Service): """ # note that this parent.file_system.create_folder call in the future will be authenticated by using permissions # handler. This permission will be granted based on service account given to the database service. + self.parent: Node folder = self.parent.file_system.create_folder(folder_name) self.parent.file_system.create_file("db_primary_store", db_size, FileSystemFileType.MDF, folder=folder) self.parent.file_system.create_file("db_transaction_log", "1", FileSystemFileType.LDF, folder=folder) @@ -37,7 +46,3 @@ class DatabaseService(Service): self.parent.file_system.create_file( "db_secondary_store", secondary_db_size, FileSystemFileType.NDF, folder=folder ) - - # todo next: - # create session? (maybe not) - # add actions for setting service state to compromised? (probably definitely) diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 8db0b0c4..17eaee3d 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -1,10 +1,13 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, Set +from typing import Any, Dict, Set, TYPE_CHECKING from primaite.simulator.core import Action, ActionManager, SimComponent from primaite.simulator.network.transmission.transport_layer import Port +if TYPE_CHECKING: + from primaite.simulator.network.hardware.base import Node + class SoftwareType(Enum): """ @@ -132,6 +135,20 @@ class Software(SimComponent): """ self.health_state_actual = health_state + @abstractmethod + def install(self) -> None: + """ + Perform first-time setup of this service on a node. + + This is an abstract class that should be overwritten by specific applications or services. It must be called + after the service is already associate with a node. For example, a service may need to authenticate with a + server during installation, or create files in the node's filesystem. + + :param node: Node on which this software runs. + :type node: Node + """ + parent: "Node" = self.parent # noqa + def scan(self) -> None: """Update the observed health status to match the actual health status.""" self.health_state_visible = self.health_state_actual diff --git a/tests/integration_tests/system/__init__.py b/tests/integration_tests/system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py new file mode 100644 index 00000000..f295eaf1 --- /dev/null +++ b/tests/integration_tests/system/test_database_on_node.py @@ -0,0 +1,22 @@ +from primaite.simulator.network.hardware.base import Node +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.services.database import DatabaseService +from primaite.simulator.system.services.service import ServiceOperatingState +from primaite.simulator.system.software import SoftwareCriticality, SoftwareHealthState + + +def test_installing_database(): + db = DatabaseService( + name="SQL-database", + health_state_actual=SoftwareHealthState.GOOD, + health_state_visible=SoftwareHealthState.GOOD, + criticality=SoftwareCriticality.MEDIUM, + ports=[ + Port.SQL_SERVER, + ], + operating_state=ServiceOperatingState.RUNNING, + ) + + node = Node(hostname="db-server") + + node.install_service(db) diff --git a/tests/unit_tests/_primaite/_simulator/_system/__init__.py b/tests/unit_tests/_primaite/_simulator/_system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/__init__.py b/tests/unit_tests/_primaite/_simulator/_system/_services/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py new file mode 100644 index 00000000..ea5c1b83 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py @@ -0,0 +1,17 @@ +from primaite.simulator.network.transmission.transport_layer import Port +from primaite.simulator.system.services.database import DatabaseService +from primaite.simulator.system.services.service import ServiceOperatingState +from primaite.simulator.system.software import SoftwareCriticality, SoftwareHealthState + + +def test_creation(): + db = DatabaseService( + name="SQL-database", + health_state_actual=SoftwareHealthState.GOOD, + health_state_visible=SoftwareHealthState.GOOD, + criticality=SoftwareCriticality.MEDIUM, + ports=[ + Port.SQL_SERVER, + ], + operating_state=ServiceOperatingState.RUNNING, + )