diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index 2ac23661..ee0eb7ff 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -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: diff --git a/src/primaite/game/agent/observations.py b/src/primaite/game/agent/observations.py index eecf4163..0cb3e8f6 100644 --- a/src/primaite/game/agent/observations.py +++ b/src/primaite/game/agent/observations.py @@ -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 ] diff --git a/src/primaite/session/environment.py b/src/primaite/session/environment.py index 36ab3f58..6701f183 100644 --- a/src/primaite/session/environment.py +++ b/src/primaite/session/environment.py @@ -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.""" diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 89329a17..7c665b9a 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -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) diff --git a/src/primaite/simulator/system/services/ftp/ftp_server.py b/src/primaite/simulator/system/services/ftp/ftp_server.py index 0278b616..87f38597 100644 --- a/src/primaite/simulator/system/services/ftp/ftp_server.py +++ b/src/primaite/simulator/system/services/ftp/ftp_server.py @@ -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 diff --git a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py index fcd9a3cc..48a05a67 100644 --- a/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py +++ b/src/primaite/simulator/system/services/red_services/data_manipulation_bot.py @@ -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. diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py index afd6cb74..eaea6bb1 100644 --- a/src/primaite/simulator/system/services/web_server/web_server.py +++ b/src/primaite/simulator/system/services/web_server/web_server.py @@ -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: