Fix backup issues and align with Yak

This commit is contained in:
Marek Wolan
2024-01-10 18:04:48 +00:00
parent 9d06284edb
commit 1505d08721
7 changed files with 39 additions and 21 deletions

View File

@@ -112,10 +112,8 @@ agents:
- service_ref: domain_controller_dns_server
- node_ref: web_server
services:
- service_ref: web_server_database_client
- service_ref: web_server_web_service
- node_ref: database_server
services:
- service_ref: database_service
folders:
- folder_name: database
files:

View File

@@ -555,7 +555,7 @@ class NodeObservation(AbstractObservation):
folder_configs = config.get("folders", {})
folders = [
FolderObservation.from_config(
config=c, game=game, parent_where=where, num_files_per_folder=num_files_per_folder
config=c, game=game, parent_where=where + ["file_system"], num_files_per_folder=num_files_per_folder
)
for c in folder_configs
]

View File

@@ -23,7 +23,6 @@ class PrimaiteGymEnv(gymnasium.Env):
super().__init__()
self.game: "PrimaiteGame" = game
self.agent: ProxyAgent = self.game.rl_agents[0]
self.flatten_obs: bool = False
def step(self, action: ActType) -> Tuple[ObsType, SupportsFloat, bool, bool, Dict[str, Any]]:
"""Perform a step in the environment."""

View File

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union
from primaite import getLogger
from primaite.simulator.file_system.file_system import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
from primaite.simulator.file_system.folder import Folder
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.core.software_manager import SoftwareManager
@@ -39,7 +40,6 @@ class DatabaseService(Service):
kwargs["port"] = Port.POSTGRES_SERVER
kwargs["protocol"] = IPProtocol.TCP
super().__init__(**kwargs)
self._db_file: File
self._create_db_file()
def set_original_state(self):
@@ -49,7 +49,7 @@ class DatabaseService(Service):
vals_to_include = {
"password",
"connections",
"backup_server",
"backup_server_ip",
"latest_backup_directory",
"latest_backup_file_name",
}
@@ -86,8 +86,8 @@ class DatabaseService(Service):
# send backup copy of database file to FTP server
response = ftp_client_service.send_file(
dest_ip_address=self.backup_server_ip,
src_file_name=self._db_file.name,
src_folder_name=self.folder.name,
src_file_name=self.db_file.name,
src_folder_name="database",
dest_folder_name=str(self.uuid),
dest_file_name="database.db",
)
@@ -121,13 +121,10 @@ class DatabaseService(Service):
return False
# replace db file
self.file_system.delete_file(folder_name=self.folder.name, file_name="downloads.db")
self.file_system.copy_file(
src_folder_name="downloads", src_file_name="database.db", dst_folder_name=self.folder.name
)
self._db_file = self.file_system.get_file(folder_name=self.folder.name, file_name="database.db")
self.file_system.delete_file(folder_name="database", file_name="downloads.db")
self.file_system.copy_file(src_folder_name="downloads", src_file_name="database.db", dst_folder_name="database")
if self._db_file is None:
if self.db_file is None:
self.sys_log.error("Copying database backup failed.")
return False
@@ -137,8 +134,17 @@ class DatabaseService(Service):
def _create_db_file(self):
"""Creates the Simulation File and sqlite file in the file system."""
self._db_file: File = self.file_system.create_file(folder_name="database", file_name="database.db")
self.folder = self.file_system.get_folder_by_id(self._db_file.folder_id)
self.file_system.create_file(folder_name="database", file_name="database.db")
@property
def db_file(self) -> File:
"""Returns the database file."""
return self.file_system.get_file(folder_name="database", file_name="database.db")
@property
def folder(self) -> Folder:
"""Returns the database folder."""
return self.file_system.get_folder_by_id(self.db_file.folder_id)
def _process_connect(
self, session_id: str, password: Optional[str] = None
@@ -171,12 +177,12 @@ class DatabaseService(Service):
"""
self.sys_log.info(f"{self.name}: Running {query}")
if query == "SELECT":
if self._db_file.health_status == FileSystemItemHealthStatus.GOOD:
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {"status_code": 200, "type": "sql", "data": True, "uuid": query_id}
else:
return {"status_code": 404, "data": False}
elif query == "DELETE":
self._db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
return {"status_code": 200, "type": "sql", "data": False, "uuid": query_id}
else:
# Invalid query
@@ -231,3 +237,13 @@ class DatabaseService(Service):
software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
return payload["status_code"] == 200
def apply_timestep(self, timestep: int) -> None:
"""
Apply a single timestep of simulation dynamics to this service.
Here at the first step, the database backup is created, in addition to normal service update logic.
"""
if timestep == 1:
self.backup_database()
return super().apply_timestep(timestep)

View File

@@ -106,5 +106,6 @@ class FTPServer(FTPServiceABC):
if payload.status_code is not None:
return False
self.send(self._process_ftp_command(payload=payload, session_id=session_id), session_id)
# self.send(self._process_ftp_command(payload=payload, session_id=session_id), session_id)
self._process_ftp_command(payload=payload, session_id=session_id)
return True

View File

@@ -84,7 +84,7 @@ class DataManipulationBot(DatabaseClient):
payload: Optional[str] = None,
port_scan_p_of_success: float = 0.1,
data_manipulation_p_of_success: float = 0.1,
repeat: bool = False,
repeat: bool = True,
):
"""
Configure the DataManipulatorBot to communicate with a DatabaseService.

View File

@@ -13,6 +13,7 @@ from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.database_client import DatabaseClient
from primaite.simulator.system.services.service import Service
from primaite.simulator.system.software import SoftwareHealthState
_LOGGER = getLogger(__name__)
@@ -123,7 +124,10 @@ class WebServer(Service):
# get all users
if db_client.query("SELECT"):
# query succeeded
self.set_health_state(SoftwareHealthState.GOOD)
response.status_code = HttpStatusCode.OK
else:
self.set_health_state(SoftwareHealthState.COMPROMISED)
return response
except Exception: