Fix problem with checking connection for db admin
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user