From a9e969aa13cc5d824356a451f32fb7ae9d8d4de6 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 25 Aug 2023 15:29:53 +0100 Subject: [PATCH 01/14] init work on database --- .../file_system/file_system_file_type.py | 8 ++++ .../simulator/system/services/database.py | 41 +++++++++++++++++++ .../simulator/system/services/service.py | 12 ++++++ 3 files changed, 61 insertions(+) create mode 100644 src/primaite/simulator/system/services/database.py diff --git a/src/primaite/simulator/file_system/file_system_file_type.py b/src/primaite/simulator/file_system/file_system_file_type.py index 7e2d8706..88aeb430 100644 --- a/src/primaite/simulator/file_system/file_system_file_type.py +++ b/src/primaite/simulator/file_system/file_system_file_type.py @@ -87,6 +87,14 @@ class FileSystemFileType(str, Enum): GZ = 31 "Gzip compressed file." + # Database file types + MDF = 32 + "MS SQL Server primary database file" + NDF = 33 + "MS SQL Server secondary database file" + LDF = 34 + "MS SQL Server transaction log" + file_type_sizes_KB = { FileSystemFileType.UNKNOWN: 0, diff --git a/src/primaite/simulator/system/services/database.py b/src/primaite/simulator/system/services/database.py new file mode 100644 index 00000000..29e3f242 --- /dev/null +++ b/src/primaite/simulator/system/services/database.py @@ -0,0 +1,41 @@ +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 + + +class DatabaseService(Service): + """TODO.""" + + def __init__(self, parent_node: Node, **kwargs): + super().__init__(**kwargs) + self._setup_files_on_node() + + def _setup_files_on_node( + self, + db_size: int = 1000, + use_secondary_db_file: bool = False, + secondary_db_size: int = 300, + folder_name: str = "database", + ): + """Set up files that are required by the database on the parent host. + + :param db_size: Initial file size of the main database file, defaults to 1000 + :type db_size: int, optional + :param use_secondary_db_file: Whether to use a secondary database file, defaults to False + :type use_secondary_db_file: bool, optional + :param secondary_db_size: Size of the secondary db file, defaults to None + :type secondary_db_size: int, optional + :param folder_name: Name of the folder which will be setup to hold the db files, defaults to "database" + :type folder_name: str, optional + """ + 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) + if use_secondary_db_file: + 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/services/service.py b/src/primaite/simulator/system/services/service.py index eafff3f0..ed2aa23b 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -2,6 +2,7 @@ from abc import abstractmethod from enum import Enum from typing import Any, Dict, List +from primaite.simulator.network.hardware.base import Node from primaite.simulator.system.software import IOSoftware @@ -32,6 +33,17 @@ class Service(IOSoftware): operating_state: ServiceOperatingState "The current operating state of the Service." + @abstractmethod + def __init__(self, parent_node: Node, **kwargs): + """Create the service on a node. + + :param parent_node: The node on which this service runs. + :type parent_node: Node + """ + super().__init__(**kwargs) + self.parent: Node = parent_node + self.parent.software_manager.add_service(self) + @abstractmethod def describe_state(self) -> Dict: """ From 319e87d200e81da1b33437c5f545f395ae063028 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 28 Aug 2023 22:34:20 +0100 Subject: [PATCH 02/14] Make changes to the way actions work --- src/primaite/simulator/core.py | 24 ++++++++++++- src/primaite/simulator/domain/controller.py | 24 +++++++------ src/primaite/simulator/network/container.py | 21 +++++++----- src/primaite/simulator/sim_container.py | 34 ++++++++++--------- .../system/applications/application.py | 10 +----- .../simulator/system/services/database.py | 2 ++ .../simulator/system/services/service.py | 12 ------- src/primaite/simulator/system/software.py | 26 +++++++------- 8 files changed, 83 insertions(+), 70 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index 63120ecf..7d8999e8 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -136,7 +136,7 @@ class SimComponent(BaseModel): if not kwargs.get("uuid"): kwargs["uuid"] = str(uuid4()) super().__init__(**kwargs) - self.action_manager: Optional[ActionManager] = None + self._action_manager: ActionManager = self._init_action_manager() self._parent: Optional["SimComponent"] = None @abstractmethod @@ -153,6 +153,28 @@ class SimComponent(BaseModel): } return state + def _init_action_manager(self) -> ActionManager: + """ + Initialise the action manager for this component. + + When using a hierarchy of components, the child classes should call the parent class's _init_action_manager and + add additional actions on top of the existing generic ones. + + Example usage for inherited classes: + + ..code::python + + class WebBrowser(Application): + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() # all actions generic to any Application get initialised + am.add_action(...) # initialise any actions specific to the web browser + return am + + :return: Actiona manager object belonging to this sim component. + :rtype: ActionManager + """ + return ActionManager() + def apply_action(self, action: List[str], context: Dict = {}) -> None: """ Apply an action to a simulation component. Action data is passed in as a 'namespaced' list of strings. diff --git a/src/primaite/simulator/domain/controller.py b/src/primaite/simulator/domain/controller.py index f772ab22..961ef550 100644 --- a/src/primaite/simulator/domain/controller.py +++ b/src/primaite/simulator/domain/controller.py @@ -85,17 +85,6 @@ class DomainController(SimComponent): def __init__(self, **kwargs): super().__init__(**kwargs) - self.action_manager = ActionManager() - # Action 'account' matches requests like: - # ['account', '', *account_action] - self.action_manager.add_action( - "account", - Action( - func=lambda request, context: self.accounts[request.pop(0)].apply_action(request, context), - validator=GroupMembershipValidator([AccountGroup.DOMAIN_ADMIN]), - ), - ) - def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -109,6 +98,19 @@ class DomainController(SimComponent): state.update({"accounts": {uuid: acct.describe_state() for uuid, acct in self.accounts.items()}}) return state + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + # Action 'account' matches requests like: + # ['account', '', *account_action] + am.add_action( + "account", + Action( + func=lambda request, context: self.accounts[request.pop(0)].apply_action(request, context), + validator=GroupMembershipValidator([AccountGroup.DOMAIN_ADMIN]), + ), + ) + return am + def _register_account(self, account: Account) -> None: """TODO.""" ... diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 1c03358c..d04da987 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -17,15 +17,6 @@ class Network(SimComponent): """Initialise the network.""" super().__init__(**kwargs) - self.action_manager = ActionManager() - self.action_manager.add_action( - "node", - Action( - func=lambda request, context: self.nodes[request.pop(0)].apply_action(request, context), - validator=AllowAllValidator(), - ), - ) - def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -44,6 +35,18 @@ class Network(SimComponent): ) return state + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + + am.add_action( + "node", + Action( + func=lambda request, context: self.nodes[request.pop(0)].apply_action(request, context), + validator=AllowAllValidator(), + ), + ) + return am + def add_node(self, node: Node) -> None: """ Add an existing node to the network. diff --git a/src/primaite/simulator/sim_container.py b/src/primaite/simulator/sim_container.py index 319defe4..8f676e6f 100644 --- a/src/primaite/simulator/sim_container.py +++ b/src/primaite/simulator/sim_container.py @@ -21,22 +21,6 @@ class Simulation(SimComponent): super().__init__(**kwargs) - self.action_manager = ActionManager() - # pass through network actions to the network objects - self.action_manager.add_action( - "network", - Action( - func=lambda request, context: self.network.apply_action(request, context), validator=AllowAllValidator() - ), - ) - # pass through domain actions to the domain object - self.action_manager.add_action( - "domain", - Action( - func=lambda request, context: self.domain.apply_action(request, context), validator=AllowAllValidator() - ), - ) - def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -54,3 +38,21 @@ class Simulation(SimComponent): } ) return state + + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + # pass through network actions to the network objects + am.add_action( + "network", + Action( + func=lambda request, context: self.network.apply_action(request, context), validator=AllowAllValidator() + ), + ) + # pass through domain actions to the domain object + am.add_action( + "domain", + Action( + func=lambda request, context: self.domain.apply_action(request, context), validator=AllowAllValidator() + ), + ) + return am diff --git a/src/primaite/simulator/system/applications/application.py b/src/primaite/simulator/system/applications/application.py index 37748560..6a07f00f 100644 --- a/src/primaite/simulator/system/applications/application.py +++ b/src/primaite/simulator/system/applications/application.py @@ -1,6 +1,6 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, List, Set +from typing import Any, Dict, Set from primaite.simulator.system.software import IOSoftware @@ -53,14 +53,6 @@ class Application(IOSoftware): ) return state - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the Application. - - :param action: A list of actions to apply. - """ - pass - def reset_component_for_episode(self, episode: int): """ Resets the Application component for a new episode. diff --git a/src/primaite/simulator/system/services/database.py b/src/primaite/simulator/system/services/database.py index 29e3f242..720967e7 100644 --- a/src/primaite/simulator/system/services/database.py +++ b/src/primaite/simulator/system/services/database.py @@ -28,6 +28,8 @@ class DatabaseService(Service): :param folder_name: Name of the folder which will be setup to hold the db files, defaults to "database" :type folder_name: str, optional """ + # 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. 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) diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index ed2aa23b..eafff3f0 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -2,7 +2,6 @@ from abc import abstractmethod from enum import Enum from typing import Any, Dict, List -from primaite.simulator.network.hardware.base import Node from primaite.simulator.system.software import IOSoftware @@ -33,17 +32,6 @@ class Service(IOSoftware): operating_state: ServiceOperatingState "The current operating state of the Service." - @abstractmethod - def __init__(self, parent_node: Node, **kwargs): - """Create the service on a node. - - :param parent_node: The node on which this service runs. - :type parent_node: Node - """ - super().__init__(**kwargs) - self.parent: Node = parent_node - self.parent.software_manager.add_service(self) - @abstractmethod def describe_state(self) -> Dict: """ diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index a2acd9fb..8e931cad 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -1,6 +1,6 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, List, Set +from typing import Any, Dict, Set from primaite.simulator.core import SimComponent from primaite.simulator.network.transmission.transport_layer import Port @@ -98,17 +98,6 @@ class Software(SimComponent): ) return state - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the software. - - The specifics of how these actions are applied should be implemented in subclasses. - - :param action: A list of actions to apply. - :type action: List[str] - """ - pass - def reset_component_for_episode(self, episode: int): """ Resets the software component for a new episode. @@ -119,6 +108,19 @@ class Software(SimComponent): """ pass + def set_health_state(self, health_state: SoftwareHealthState) -> None: + """ + Assign a new health state to this software. + + Note: this should only be possible when the software is currently running, but the software base class has no + operating state, only subclasses do. So subclasses will need to implement this check. TODO: check if this should + be changed so that the base Software class has a running attr. + + :param health_state: New health state to assign to the software + :type health_state: SoftwareHealthState + """ + self.health_state_actual = health_state + class IOSoftware(Software): """ From 1eff41c7861cb6aa53fee877b44762e56c16826f Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 11:10:29 +0100 Subject: [PATCH 03/14] Update docs based on new action options --- docs/source/simulation_structure.rst | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/docs/source/simulation_structure.rst b/docs/source/simulation_structure.rst index 7630ae0f..f3ef866c 100644 --- a/docs/source/simulation_structure.rst +++ b/docs/source/simulation_structure.rst @@ -49,16 +49,14 @@ snippet demonstrates usage of the ``ActionPermissionValidator``. name: str apps = [] - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.action_manager = ActionManager() - - self.action_manager.add_action( + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + am.add_action( "reset_factory_settings", Action( func = lambda request, context: self.reset_factory_settings(), validator = GroupMembershipValidator([AccountGroup.DOMAIN_ADMIN]), - ), + ) ) def reset_factory_settings(self): From 7b61322e704b06da9ef797162ba0021df25ed116 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 12:34:41 +0100 Subject: [PATCH 04/14] Add service actions --- src/primaite/simulator/core.py | 7 +- .../simulator/system/services/service.py | 79 +++++++++++++++++-- src/primaite/simulator/system/software.py | 17 +++- 3 files changed, 92 insertions(+), 11 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index 7d8999e8..90abb675 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -41,7 +41,9 @@ class Action: the action can be performed or not. """ - def __init__(self, func: Callable[[List[str], Dict], None], validator: ActionPermissionValidator) -> None: + def __init__( + self, func: Callable[[List[str], Dict], None], validator: ActionPermissionValidator = AllowAllValidator() + ) -> None: """ Save the functions that are for this action. @@ -58,7 +60,8 @@ class Action: :param func: Function that performs the request. :type func: Callable[[List[str], Dict], None] - :param validator: Function that checks if the request is authenticated given the context. + :param validator: Function that checks if the request is authenticated given the context. By default, if no + validator is provided, an 'allow all' validator is added which permits all requests. :type validator: ActionPermissionValidator """ self.func: Callable[[List[str], Dict], None] = func diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index eafff3f0..1a36589f 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -1,7 +1,8 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, List +from typing import Any, Dict +from primaite.simulator.core import Action, ActionManager from primaite.simulator.system.software import IOSoftware @@ -46,13 +47,16 @@ class Service(IOSoftware): state.update({"operating_state": self.operating_state.name}) return state - def apply_action(self, action: List[str]) -> None: - """ - Applies a list of actions to the Service. - - :param action: A list of actions to apply. - """ - pass + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + am.add_action("stop", Action(func=lambda request, context: self.stop())) + am.add_action("start", Action(func=lambda request, context: self.start())) + am.add_action("pause", Action(func=lambda request, context: self.pause())) + am.add_action("resume", Action(func=lambda request, context: self.resume())) + am.add_action("restart", Action(func=lambda request, context: self.restart())) + am.add_action("disable", Action(func=lambda request, context: self.disable())) + am.add_action("enable", Action(func=lambda request, context: self.enable())) + return am def reset_component_for_episode(self, episode: int): """ @@ -86,3 +90,62 @@ class Service(IOSoftware): :return: True if successful, False otherwise. """ pass + + # TODO: validate this state transition model. + # Possibly state transition could be defined more succinctly than a separate function with lots of if statements. + + def stop(self) -> None: + """Stop the service.""" + if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: + self.operating_state = ServiceOperatingState.STOPPED + + def start(self) -> None: + """Start the service.""" + if self.operating_state == ServiceOperatingState.STOPPED: + self.operating_state = ServiceOperatingState.RUNNING + + def pause(self) -> None: + """Pause the service.""" + if self.operating_state == ServiceOperatingState.RUNNING: + self.operating_state = ServiceOperatingState.PAUSED + + def resume(self) -> None: + """Resume paused service.""" + if self.operating_state == ServiceOperatingState.PAUSED: + self.operating_state = ServiceOperatingState.RUNNING + + def restart(self) -> None: + """Restart running service.""" + if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: + self.operating_state = ServiceOperatingState.RESTARTING + self.restart_countdown = 5 # TODO: implement restart duration + + def disable(self) -> None: + """Disable the service.""" + if self.operating_state in [ + ServiceOperatingState.RUNNING, + ServiceOperatingState.STOPPED, + ServiceOperatingState.PAUSED, + ]: + self.operating_state = ServiceOperatingState.DISABLED + + def enable(self) -> None: + """Enable the disabled service.""" + if self.operating_state == ServiceOperatingState.DISABLED: + self.operating_state = ServiceOperatingState.STOPPED + + def apply_timestep(self, timestep: int) -> None: + """ + Apply a single timestep of simulation dynamics to this service. + + In this instance, if any multi-timestep processes are currently occurring (such as restarting or installation), + then they are brought one step closer to being finished. + + :param timestep: The current timestep number. (Amount of time since simulation episode began) + :type timestep: int + """ + super().apply_timestep(timestep) + if self.operating_state == ServiceOperatingState.RESTARTING: + self.restart_countdown -= 1 + if self.restart_countdown <= 0: + self.operating_state = ServiceOperatingState.RUNNING diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 8e931cad..8db0b0c4 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -2,7 +2,7 @@ from abc import abstractmethod from enum import Enum from typing import Any, Dict, Set -from primaite.simulator.core import SimComponent +from primaite.simulator.core import Action, ActionManager, SimComponent from primaite.simulator.network.transmission.transport_layer import Port @@ -98,6 +98,17 @@ class Software(SimComponent): ) return state + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + am.add_action( + "compromise", + Action( + func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED), + ), + ) + am.add_action("scan", Action(func=lambda request, context: self.scan())) + return am + def reset_component_for_episode(self, episode: int): """ Resets the software component for a new episode. @@ -121,6 +132,10 @@ class Software(SimComponent): """ self.health_state_actual = health_state + def scan(self) -> None: + """Update the observed health status to match the actual health status.""" + self.health_state_visible = self.health_state_actual + class IOSoftware(Software): """ From 94325d1fde9d8da023836ed437252f03cdaae889 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 13:21:34 +0100 Subject: [PATCH 05/14] Add Install method to software. --- .../simulator/network/hardware/base.py | 35 ++++++++++++++++++- .../simulator/system/services/database.py | 21 ++++++----- src/primaite/simulator/system/software.py | 19 +++++++++- tests/integration_tests/system/__init__.py | 0 .../system/test_database_on_node.py | 22 ++++++++++++ .../_primaite/_simulator/_system/__init__.py | 0 .../_simulator/_system/_services/__init__.py | 0 .../_system/_services/test_database.py | 17 +++++++++ 8 files changed, 104 insertions(+), 10 deletions(-) create mode 100644 tests/integration_tests/system/__init__.py create mode 100644 tests/integration_tests/system/test_database_on_node.py create mode 100644 tests/unit_tests/_primaite/_simulator/_system/__init__.py create mode 100644 tests/unit_tests/_primaite/_simulator/_system/_services/__init__.py create mode 100644 tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py 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, + ) From f0b82cbdfba6070217eee4011c9c14aaab5f6f38 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 14:15:49 +0100 Subject: [PATCH 06/14] Add ability to uninstall service --- .../simulator/file_system/file_system.py | 4 +-- .../simulator/network/hardware/base.py | 1 + .../simulator/system/services/database.py | 31 ++++++++++++++--- src/primaite/simulator/system/software.py | 19 ++++++----- .../system/test_database_on_node.py | 34 +++++++++++++++++++ 5 files changed, 74 insertions(+), 15 deletions(-) diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 440b7dc5..1346d3e0 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -1,5 +1,5 @@ from random import choice -from typing import Dict, Optional +from typing import Dict, Optional, Union from primaite import getLogger from primaite.simulator.core import SimComponent @@ -211,7 +211,7 @@ class FileSystem(SimComponent): if file is not None: return file - def get_folder_by_name(self, folder_name: str) -> FileSystemFolder: + def get_folder_by_name(self, folder_name: str) -> Union[FileSystemFolder, None]: """ Returns a the first folder with a matching name. diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index a68ff480..e3e38f86 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1004,6 +1004,7 @@ class Node(SimComponent): if service in self: _LOGGER.warning(f"Can't add service {service.uuid} to node {self.uuid}. 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. _LOGGER.info(f"Added service {service.uuid} to node {self.uuid}") diff --git a/src/primaite/simulator/system/services/database.py b/src/primaite/simulator/system/services/database.py index 0d1de15c..554455b8 100644 --- a/src/primaite/simulator/system/services/database.py +++ b/src/primaite/simulator/system/services/database.py @@ -12,6 +12,21 @@ class DatabaseService(Service): """TODO.""" return super().describe_state() + def uninstall(self) -> None: + """ + Undo installation procedure. + + This method deletes files created when installing the database, and the database folder if it is empty. + """ + super().uninstall() + node: Node = self.parent + node.file_system.delete_file(self.primary_store) + node.file_system.delete_file(self.transaction_log) + if self.secondary_store: + node.file_system.delete_file(self.secondary_store) + if len(self.folder.files) == 0: + node.file_system.delete_folder(self.folder) + def install(self) -> None: """Perform first time install on a node, creating necessary files.""" super().install() @@ -39,10 +54,16 @@ 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) + self.folder = self.parent.file_system.create_folder(folder_name) + self.primary_store = self.parent.file_system.create_file( + "db_primary_store", db_size, FileSystemFileType.MDF, folder=self.folder + ) + self.transaction_log = self.parent.file_system.create_file( + "db_transaction_log", "1", FileSystemFileType.LDF, folder=self.folder + ) if use_secondary_db_file: - self.parent.file_system.create_file( - "db_secondary_store", secondary_db_size, FileSystemFileType.NDF, folder=folder + self.secondary_store = self.parent.file_system.create_file( + "db_secondary_store", secondary_db_size, FileSystemFileType.NDF, folder=self.folder ) + else: + self.secondary_store = None diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 17eaee3d..1fcdb522 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -1,13 +1,10 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict, Set, TYPE_CHECKING +from typing import Any, Dict, Set 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): """ @@ -143,11 +140,17 @@ class Software(SimComponent): 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 + pass + + def uninstall(self) -> None: + """Uninstall this service from a node. + + This is an abstract class that should be overwritten by applications or services. It must be called after the + `install` method has already been run on that node. It should undo any installation steps, for example by + deleting files, or contacting a server. + """ + pass def scan(self) -> None: """Update the observed health status to match the actual health status.""" diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py index f295eaf1..73d19339 100644 --- a/tests/integration_tests/system/test_database_on_node.py +++ b/tests/integration_tests/system/test_database_on_node.py @@ -20,3 +20,37 @@ def test_installing_database(): node = Node(hostname="db-server") node.install_service(db) + + assert db in node + + file_exists = False + for folder in node.file_system.folders.values(): + for file in folder.files.values(): + if file.name == "db_primary_store": + file_exists = True + break + if file_exists: + break + assert file_exists + + +def test_uninstalling_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) + + node.uninstall_service(db) + + assert db not in node + assert node.file_system.get_folder_by_name("database") is None From 40d3e04e648d442696636152288c52dde200f389 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 14:33:28 +0100 Subject: [PATCH 07/14] Move init action manager function to the top --- src/primaite/simulator/core.py | 28 +++++++-------- src/primaite/simulator/domain/controller.py | 26 +++++++------- src/primaite/simulator/network/container.py | 24 ++++++------- src/primaite/simulator/sim_container.py | 36 +++++++++---------- .../simulator/system/services/service.py | 22 ++++++------ src/primaite/simulator/system/software.py | 22 ++++++------ 6 files changed, 79 insertions(+), 79 deletions(-) diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index c12b1ad5..69edd8db 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -142,20 +142,6 @@ class SimComponent(BaseModel): self._action_manager: ActionManager = self._init_action_manager() self._parent: Optional["SimComponent"] = None - @abstractmethod - def describe_state(self) -> Dict: - """ - Return a dictionary describing the state of this object and any objects managed by it. - - This is similar to pydantic ``model_dump()``, but it only outputs information about the objects owned by this - object. If there are objects referenced by this object that are owned by something else, it is not included in - this output. - """ - state = { - "uuid": self.uuid, - } - return state - def _init_action_manager(self) -> ActionManager: """ Initialise the action manager for this component. @@ -178,6 +164,20 @@ class SimComponent(BaseModel): """ return ActionManager() + @abstractmethod + def describe_state(self) -> Dict: + """ + Return a dictionary describing the state of this object and any objects managed by it. + + This is similar to pydantic ``model_dump()``, but it only outputs information about the objects owned by this + object. If there are objects referenced by this object that are owned by something else, it is not included in + this output. + """ + state = { + "uuid": self.uuid, + } + return state + def apply_action(self, action: List[str], context: Dict = {}) -> None: """ Apply an action to a simulation component. Action data is passed in as a 'namespaced' list of strings. diff --git a/src/primaite/simulator/domain/controller.py b/src/primaite/simulator/domain/controller.py index 961ef550..b436ca79 100644 --- a/src/primaite/simulator/domain/controller.py +++ b/src/primaite/simulator/domain/controller.py @@ -85,19 +85,6 @@ class DomainController(SimComponent): def __init__(self, **kwargs): super().__init__(**kwargs) - 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.update({"accounts": {uuid: acct.describe_state() for uuid, acct in self.accounts.items()}}) - return state - def _init_action_manager(self) -> ActionManager: am = super()._init_action_manager() # Action 'account' matches requests like: @@ -111,6 +98,19 @@ class DomainController(SimComponent): ) return am + 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.update({"accounts": {uuid: acct.describe_state() for uuid, acct in self.accounts.items()}}) + return state + def _register_account(self, account: Account) -> None: """TODO.""" ... diff --git a/src/primaite/simulator/network/container.py b/src/primaite/simulator/network/container.py index 95eaeb0c..e0226e6c 100644 --- a/src/primaite/simulator/network/container.py +++ b/src/primaite/simulator/network/container.py @@ -17,6 +17,18 @@ class Network(SimComponent): """Initialise the network.""" super().__init__(**kwargs) + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + + am.add_action( + "node", + Action( + func=lambda request, context: self.nodes[request.pop(0)].apply_action(request, context), + validator=AllowAllValidator(), + ), + ) + return am + def describe_state(self) -> Dict: """ Produce a dictionary describing the current state of this object. @@ -35,18 +47,6 @@ class Network(SimComponent): ) return state - def _init_action_manager(self) -> ActionManager: - am = super()._init_action_manager() - - am.add_action( - "node", - Action( - func=lambda request, context: self.nodes[request.pop(0)].apply_action(request, context), - validator=AllowAllValidator(), - ), - ) - return am - def add_node(self, node: Node) -> None: """ Add an existing node to the network. diff --git a/src/primaite/simulator/sim_container.py b/src/primaite/simulator/sim_container.py index 8f676e6f..2a5123f3 100644 --- a/src/primaite/simulator/sim_container.py +++ b/src/primaite/simulator/sim_container.py @@ -21,24 +21,6 @@ class Simulation(SimComponent): super().__init__(**kwargs) - 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.update( - { - "network": self.network.describe_state(), - "domain": self.domain.describe_state(), - } - ) - return state - def _init_action_manager(self) -> ActionManager: am = super()._init_action_manager() # pass through network actions to the network objects @@ -56,3 +38,21 @@ class Simulation(SimComponent): ), ) return am + + 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.update( + { + "network": self.network.describe_state(), + "domain": self.domain.describe_state(), + } + ) + return state diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 1a36589f..7e67d05f 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -33,6 +33,17 @@ class Service(IOSoftware): operating_state: ServiceOperatingState "The current operating state of the Service." + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + am.add_action("stop", Action(func=lambda request, context: self.stop())) + am.add_action("start", Action(func=lambda request, context: self.start())) + am.add_action("pause", Action(func=lambda request, context: self.pause())) + am.add_action("resume", Action(func=lambda request, context: self.resume())) + am.add_action("restart", Action(func=lambda request, context: self.restart())) + am.add_action("disable", Action(func=lambda request, context: self.disable())) + am.add_action("enable", Action(func=lambda request, context: self.enable())) + return am + @abstractmethod def describe_state(self) -> Dict: """ @@ -47,17 +58,6 @@ class Service(IOSoftware): state.update({"operating_state": self.operating_state.name}) return state - def _init_action_manager(self) -> ActionManager: - am = super()._init_action_manager() - am.add_action("stop", Action(func=lambda request, context: self.stop())) - am.add_action("start", Action(func=lambda request, context: self.start())) - am.add_action("pause", Action(func=lambda request, context: self.pause())) - am.add_action("resume", Action(func=lambda request, context: self.resume())) - am.add_action("restart", Action(func=lambda request, context: self.restart())) - am.add_action("disable", Action(func=lambda request, context: self.disable())) - am.add_action("enable", Action(func=lambda request, context: self.enable())) - return am - def reset_component_for_episode(self, episode: int): """ Resets the Service component for a new episode. diff --git a/src/primaite/simulator/system/software.py b/src/primaite/simulator/system/software.py index 1fcdb522..605a062b 100644 --- a/src/primaite/simulator/system/software.py +++ b/src/primaite/simulator/system/software.py @@ -75,6 +75,17 @@ class Software(SimComponent): revealed_to_red: bool = False "Indicates if the software has been revealed to red agent, defaults is False." + def _init_action_manager(self) -> ActionManager: + am = super()._init_action_manager() + am.add_action( + "compromise", + Action( + func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED), + ), + ) + am.add_action("scan", Action(func=lambda request, context: self.scan())) + return am + @abstractmethod def describe_state(self) -> Dict: """ @@ -98,17 +109,6 @@ class Software(SimComponent): ) return state - def _init_action_manager(self) -> ActionManager: - am = super()._init_action_manager() - am.add_action( - "compromise", - Action( - func=lambda request, context: self.set_health_state(SoftwareHealthState.COMPROMISED), - ), - ) - am.add_action("scan", Action(func=lambda request, context: self.scan())) - return am - def reset_component_for_episode(self, episode: int): """ Resets the software component for a new episode. From 62be66205cdbb91e43b3438eae396b660dff199c Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 31 Aug 2023 10:57:45 +0100 Subject: [PATCH 08/14] Fix unit tests --- .../_primaite/_simulator/_network/_hardware/test_nic.py | 2 +- tests/unit_tests/_primaite/_simulator/test_core.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py index dc508508..c417b5b9 100644 --- a/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py +++ b/tests/unit_tests/_primaite/_simulator/_network/_hardware/test_nic.py @@ -48,7 +48,7 @@ def test_nic_deserialize(): nic_json = nic.model_dump_json() deserialized_nic = NIC.model_validate_json(nic_json) - assert nic == deserialized_nic + assert nic_json == deserialized_nic.model_dump_json() def test_nic_ip_address_as_gateway_fails(): diff --git a/tests/unit_tests/_primaite/_simulator/test_core.py b/tests/unit_tests/_primaite/_simulator/test_core.py index 0d227633..bbb1298f 100644 --- a/tests/unit_tests/_primaite/_simulator/test_core.py +++ b/tests/unit_tests/_primaite/_simulator/test_core.py @@ -43,4 +43,5 @@ class TestIsolatedSimComponent: comp = TestComponent(name="computer", size=(5, 10)) dump = comp.model_dump_json() - assert comp == TestComponent.model_validate_json(dump) + reconstructed = TestComponent.model_validate_json(dump) + assert comp == reconstructed.model_dump_json() From 7759c178bbe68c63430181c637a96a06c3c61d48 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 31 Aug 2023 11:20:16 +0100 Subject: [PATCH 09/14] Add logging and service restarting --- .../simulator/network/hardware/base.py | 2 ++ .../simulator/system/services/service.py | 24 ++++++++++++------- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index e3e38f86..2bdb4b55 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/src/primaite/simulator/network/hardware/base.py @@ -1007,6 +1007,7 @@ class Node(SimComponent): 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.info(f"Added service {service.uuid} to node {self.uuid}") def uninstall_service(self, service: Service) -> None: @@ -1021,6 +1022,7 @@ class Node(SimComponent): 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}") _LOGGER.info(f"Removed service {service.uuid} from node {self.uuid}") def __contains__(self, item: Any) -> bool: diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 7e67d05f..6932ce4c 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -1,6 +1,6 @@ from abc import abstractmethod from enum import Enum -from typing import Any, Dict +from typing import Any, Dict, Optional from primaite.simulator.core import Action, ActionManager from primaite.simulator.system.software import IOSoftware @@ -32,6 +32,10 @@ class Service(IOSoftware): operating_state: ServiceOperatingState "The current operating state of the Service." + restart_duration: int = 5 + "How many timesteps does it take to restart this service." + _restart_countdown: Optional[int] = None + "If currently restarting, how many timesteps remain until the restart is finished." def _init_action_manager(self) -> ActionManager: am = super()._init_action_manager() @@ -97,41 +101,43 @@ class Service(IOSoftware): def stop(self) -> None: """Stop the service.""" if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: + self.parent.sys_log.info(f"Stopping service {self.name}") self.operating_state = ServiceOperatingState.STOPPED def start(self) -> None: """Start the service.""" if self.operating_state == ServiceOperatingState.STOPPED: + self.parent.sys_log.info(f"Starting service {self.name}") self.operating_state = ServiceOperatingState.RUNNING def pause(self) -> None: """Pause the service.""" if self.operating_state == ServiceOperatingState.RUNNING: + self.parent.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.PAUSED def resume(self) -> None: """Resume paused service.""" if self.operating_state == ServiceOperatingState.PAUSED: + self.parent.sys_log.info(f"Resuming service {self.name}") self.operating_state = ServiceOperatingState.RUNNING def restart(self) -> None: """Restart running service.""" if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: + self.parent.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.RESTARTING - self.restart_countdown = 5 # TODO: implement restart duration + self.restart_countdown = self.restarting_duration # TODO: implement restart duration def disable(self) -> None: """Disable the service.""" - if self.operating_state in [ - ServiceOperatingState.RUNNING, - ServiceOperatingState.STOPPED, - ServiceOperatingState.PAUSED, - ]: - self.operating_state = ServiceOperatingState.DISABLED + self.parent.sys_log.info(f"Disabling Application {self.name}") + self.operating_state = ServiceOperatingState.DISABLED def enable(self) -> None: """Enable the disabled service.""" if self.operating_state == ServiceOperatingState.DISABLED: + self.parent.sys_log.info(f"Enabling Application {self.name}") self.operating_state = ServiceOperatingState.STOPPED def apply_timestep(self, timestep: int) -> None: @@ -146,6 +152,6 @@ class Service(IOSoftware): """ super().apply_timestep(timestep) if self.operating_state == ServiceOperatingState.RESTARTING: - self.restart_countdown -= 1 if self.restart_countdown <= 0: self.operating_state = ServiceOperatingState.RUNNING + self.restart_countdown -= 1 From f60f775f03a9a4ef44d4c3e491d0bf5e50a7111d Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 31 Aug 2023 11:27:52 +0100 Subject: [PATCH 10/14] Improve logging --- src/primaite/simulator/system/services/service.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 6932ce4c..756f723d 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -2,9 +2,12 @@ from abc import abstractmethod from enum import Enum from typing import Any, Dict, Optional +from primaite import getLogger from primaite.simulator.core import Action, ActionManager from primaite.simulator.system.software import IOSoftware +_LOGGER = getLogger(__name__) + class ServiceOperatingState(Enum): """Enumeration of Service Operating States.""" @@ -95,35 +98,37 @@ class Service(IOSoftware): """ pass - # TODO: validate this state transition model. - # Possibly state transition could be defined more succinctly than a separate function with lots of if statements. - def stop(self) -> None: """Stop the service.""" + _LOGGER.debug(f"Stopping service {self.name}") if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: self.parent.sys_log.info(f"Stopping service {self.name}") self.operating_state = ServiceOperatingState.STOPPED def start(self) -> None: """Start the service.""" + _LOGGER.debug(f"Starting service {self.name}") if self.operating_state == ServiceOperatingState.STOPPED: self.parent.sys_log.info(f"Starting service {self.name}") self.operating_state = ServiceOperatingState.RUNNING def pause(self) -> None: """Pause the service.""" + _LOGGER.debug(f"Pausing service {self.name}") if self.operating_state == ServiceOperatingState.RUNNING: self.parent.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.PAUSED def resume(self) -> None: """Resume paused service.""" + _LOGGER.debug(f"Resuming service {self.name}") if self.operating_state == ServiceOperatingState.PAUSED: self.parent.sys_log.info(f"Resuming service {self.name}") self.operating_state = ServiceOperatingState.RUNNING def restart(self) -> None: """Restart running service.""" + _LOGGER.debug(f"Restarting service {self.name}") if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: self.parent.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.RESTARTING @@ -131,11 +136,13 @@ class Service(IOSoftware): def disable(self) -> None: """Disable the service.""" + _LOGGER.debug(f"Disabling service {self.name}") self.parent.sys_log.info(f"Disabling Application {self.name}") self.operating_state = ServiceOperatingState.DISABLED def enable(self) -> None: """Enable the disabled service.""" + _LOGGER.debug(f"Enabling service {self.name}") if self.operating_state == ServiceOperatingState.DISABLED: self.parent.sys_log.info(f"Enabling Application {self.name}") self.operating_state = ServiceOperatingState.STOPPED @@ -153,5 +160,6 @@ class Service(IOSoftware): super().apply_timestep(timestep) if self.operating_state == ServiceOperatingState.RESTARTING: if self.restart_countdown <= 0: + _LOGGER.debug(f"Restarting finished for service {self.name}") self.operating_state = ServiceOperatingState.RUNNING self.restart_countdown -= 1 From bd5aacaf0c33ee56650cf6e11e60f0d7b9d9641e Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 31 Aug 2023 11:32:11 +0100 Subject: [PATCH 11/14] Remove todo comments --- src/primaite/simulator/system/services/database.py | 11 +++++++++-- src/primaite/simulator/system/services/service.py | 2 +- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/primaite/simulator/system/services/database.py b/src/primaite/simulator/system/services/database.py index 554455b8..23b856f7 100644 --- a/src/primaite/simulator/system/services/database.py +++ b/src/primaite/simulator/system/services/database.py @@ -6,10 +6,17 @@ from primaite.simulator.system.services.service import Service class DatabaseService(Service): - """TODO.""" + """Service loosely modelled on Microsoft SQL Server.""" def describe_state(self) -> Dict: - """TODO.""" + """ + 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 + """ return super().describe_state() def uninstall(self) -> None: diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py index 756f723d..f9cc784d 100644 --- a/src/primaite/simulator/system/services/service.py +++ b/src/primaite/simulator/system/services/service.py @@ -132,7 +132,7 @@ class Service(IOSoftware): if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]: self.parent.sys_log.info(f"Pausing service {self.name}") self.operating_state = ServiceOperatingState.RESTARTING - self.restart_countdown = self.restarting_duration # TODO: implement restart duration + self.restart_countdown = self.restarting_duration def disable(self) -> None: """Disable the service.""" From 61fa83a00d81e8967f13966bab99e898b93552f3 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Thu, 31 Aug 2023 14:55:14 +0100 Subject: [PATCH 12/14] Fix failing test --- tests/unit_tests/_primaite/_simulator/test_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/_primaite/_simulator/test_core.py b/tests/unit_tests/_primaite/_simulator/test_core.py index bbb1298f..069e6ea2 100644 --- a/tests/unit_tests/_primaite/_simulator/test_core.py +++ b/tests/unit_tests/_primaite/_simulator/test_core.py @@ -44,4 +44,4 @@ class TestIsolatedSimComponent: comp = TestComponent(name="computer", size=(5, 10)) dump = comp.model_dump_json() reconstructed = TestComponent.model_validate_json(dump) - assert comp == reconstructed.model_dump_json() + assert dump == reconstructed.model_dump_json() From 289f81826637c4760d7f98dca8e17cff70d60204 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Mon, 4 Sep 2023 16:11:47 +0000 Subject: [PATCH 13/14] Apply suggestions from code review --- src/primaite/simulator/file_system/file_system.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index 1346d3e0..e9385644 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -211,7 +211,7 @@ class FileSystem(SimComponent): if file is not None: return file - def get_folder_by_name(self, folder_name: str) -> Union[FileSystemFolder, None]: + def get_folder_by_name(self, folder_name: str) -> Optional[FileSystemFolder]: """ Returns a the first folder with a matching name. From 0892a976fd47b6f49b1ec7ddd008068ae821c3d8 Mon Sep 17 00:00:00 2001 From: Christopher McCarthy Date: Mon, 4 Sep 2023 18:37:05 +0000 Subject: [PATCH 14/14] Apply suggestions from code review --- src/primaite/simulator/file_system/file_system.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py index e9385644..79159e60 100644 --- a/src/primaite/simulator/file_system/file_system.py +++ b/src/primaite/simulator/file_system/file_system.py @@ -1,5 +1,5 @@ from random import choice -from typing import Dict, Optional, Union +from typing import Dict, Optional from primaite import getLogger from primaite.simulator.core import SimComponent