Add service actions
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user