diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py index addad35a..69065225 100644 --- a/src/primaite/simulator/system/applications/database_client.py +++ b/src/primaite/simulator/system/applications/database_client.py @@ -43,9 +43,9 @@ class DatabaseClient(Application): def execute(self) -> bool: """Execution definition for db client: perform a select query.""" if self.connections: - can_connect = self.connect(connection_id=list(self.connections.keys())[-1]) + can_connect = self.check_connection(connection_id=list(self.connections.keys())[-1]) else: - can_connect = self.connect() + can_connect = self.check_connection(connection_id=str(uuid4())) self._last_connection_successful = can_connect return can_connect @@ -79,15 +79,17 @@ class DatabaseClient(Application): if not connection_id: connection_id = str(uuid4()) - # if we are reusing a connection_id, remove it from self.connections so that its new status can be populated - # warning: janky - self._connections.pop(connection_id, None) - self.connected = self._connect( server_ip_address=self.server_ip_address, password=self.server_password, connection_id=connection_id ) return self.connected + def check_connection(self, connection_id:str) -> bool: + if not self._can_perform_action(): + return False + print(self.query("SELECT * FROM pg_stat_activity", connection_id=connection_id)) + return self.connected + def _connect( self, server_ip_address: IPv4Address, diff --git a/src/primaite/simulator/system/applications/red_applications/dos_bot.py b/src/primaite/simulator/system/applications/red_applications/dos_bot.py index 9dac6b25..1247bc99 100644 --- a/src/primaite/simulator/system/applications/red_applications/dos_bot.py +++ b/src/primaite/simulator/system/applications/red_applications/dos_bot.py @@ -28,7 +28,7 @@ class DoSAttackStage(IntEnum): "Attack is completed." -class DoSBot(DatabaseClient, Application): +class DoSBot(DatabaseClient): """A bot that simulates a Denial of Service attack.""" target_ip_address: Optional[IPv4Address] = None diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py index 9fdfd5ff..c73132eb 100644 --- a/src/primaite/simulator/system/services/database/database_service.py +++ b/src/primaite/simulator/system/services/database/database_service.py @@ -221,6 +221,18 @@ class DatabaseService(Service): } else: return {"status_code": 404, "data": False} + elif query == "SELECT * FROM pg_stat_activity": + # Check if the connection is active. + if self.health_state_actual == SoftwareHealthState.GOOD: + return { + "status_code": 200, + "type": "sql", + "data": False, + "uuid": query_id, + "connection_id": connection_id, + } + else: + return {"status_code": 401, "data": False} else: # Invalid query self.sys_log.info(f"{self.name}: Invalid {query}")