Fix problem with checking connection for db admin

This commit is contained in:
Marek Wolan
2024-03-01 16:36:41 +00:00
parent ed01293b86
commit 2a1d99ccce
3 changed files with 21 additions and 7 deletions

View File

@@ -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,

View File

@@ -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

View File

@@ -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}")