From 40d3e04e648d442696636152288c52dde200f389 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Tue, 29 Aug 2023 14:33:28 +0100 Subject: [PATCH] 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.