Merge branch 'dev' into feature/1752-dns-server-and-client
This commit is contained in:
Czar Echavez
76
src/primaite/simulator/system/services/database.py
Normal file
76
src/primaite/simulator/system/services/database.py
Normal file
@@ -0,0 +1,76 @@
|
||||
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
|
||||
|
||||
|
||||
class DatabaseService(Service):
|
||||
"""Service loosely modelled on Microsoft SQL Server."""
|
||||
|
||||
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
|
||||
"""
|
||||
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()
|
||||
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,
|
||||
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
|
||||
"""
|
||||
# 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
|
||||
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.secondary_store = self.parent.file_system.create_file(
|
||||
"db_secondary_store", secondary_db_size, FileSystemFileType.NDF, folder=self.folder
|
||||
)
|
||||
else:
|
||||
self.secondary_store = None
|
||||
@@ -0,0 +1,34 @@
|
||||
from ipaddress import IPv4Address
|
||||
from typing import Any, Optional
|
||||
|
||||
from primaite.simulator.network.transmission.network_layer import IPProtocol
|
||||
from primaite.simulator.network.transmission.transport_layer import Port
|
||||
from primaite.simulator.system.services.service import Service
|
||||
|
||||
|
||||
class DataManipulatorService(Service):
|
||||
"""
|
||||
Red Agent Data Integration Service.
|
||||
|
||||
The Service represents a bot that causes files/folders in the File System to
|
||||
become corrupted.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
kwargs["name"] = "DataManipulatorBot"
|
||||
kwargs["port"] = Port.POSTGRES_SERVER
|
||||
kwargs["protocol"] = IPProtocol.TCP
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def start(self, target_ip_address: IPv4Address, payload: Optional[Any] = "DELETE TABLE users", **kwargs):
|
||||
"""
|
||||
Run the DataManipulatorService actions.
|
||||
|
||||
:param: target_ip_address: The IP address of the target machine to attack
|
||||
:param: payload: The payload to send to the target machine
|
||||
"""
|
||||
super().start()
|
||||
|
||||
self.software_manager.send_payload_to_session_manager(
|
||||
payload=payload, dest_ip_address=target_ip_address, dest_port=self.port
|
||||
)
|
||||
@@ -1,9 +1,12 @@
|
||||
from abc import abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List
|
||||
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."""
|
||||
@@ -29,29 +32,36 @@ class Service(IOSoftware):
|
||||
Services are programs that run in the background and may perform input/output operations.
|
||||
"""
|
||||
|
||||
operating_state: ServiceOperatingState
|
||||
operating_state: ServiceOperatingState = ServiceOperatingState.STOPPED
|
||||
"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()
|
||||
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:
|
||||
"""
|
||||
Describes the current state of the software.
|
||||
Produce a dictionary describing the current state of this object.
|
||||
|
||||
The specifics of the software's state, including its health, criticality,
|
||||
and any other pertinent information, should be implemented in subclasses.
|
||||
Please see :py:meth:`primaite.simulator.core.SimComponent.describe_state` for a more detailed explanation.
|
||||
|
||||
:return: A dictionary containing key-value pairs representing the current state of the software.
|
||||
:return: Current state of this object and child objects.
|
||||
:rtype: Dict
|
||||
"""
|
||||
pass
|
||||
|
||||
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
|
||||
state = super().describe_state()
|
||||
state.update({"operating_state": self.operating_state.name})
|
||||
return state
|
||||
|
||||
def reset_component_for_episode(self, episode: int):
|
||||
"""
|
||||
@@ -85,3 +95,69 @@ class Service(IOSoftware):
|
||||
:return: True if successful, False otherwise.
|
||||
"""
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the service."""
|
||||
_LOGGER.debug(f"Stopping service {self.name}")
|
||||
if self.operating_state in [ServiceOperatingState.RUNNING, ServiceOperatingState.PAUSED]:
|
||||
self.sys_log.info(f"Stopping service {self.name}")
|
||||
self.operating_state = ServiceOperatingState.STOPPED
|
||||
|
||||
def start(self, **kwargs) -> None:
|
||||
"""Start the service."""
|
||||
_LOGGER.debug(f"Starting service {self.name}")
|
||||
if self.operating_state == ServiceOperatingState.STOPPED:
|
||||
self.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.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.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.sys_log.info(f"Pausing service {self.name}")
|
||||
self.operating_state = ServiceOperatingState.RESTARTING
|
||||
self.restart_countdown = self.restarting_duration
|
||||
|
||||
def disable(self) -> None:
|
||||
"""Disable the service."""
|
||||
_LOGGER.debug(f"Disabling service {self.name}")
|
||||
self.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.sys_log.info(f"Enabling Application {self.name}")
|
||||
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:
|
||||
if self.restart_countdown <= 0:
|
||||
_LOGGER.debug(f"Restarting finished for service {self.name}")
|
||||
self.operating_state = ServiceOperatingState.RUNNING
|
||||
self.restart_countdown -= 1
|
||||
|
||||
Reference in New Issue
Block a user