oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2448: store last query response for db client

Browse Source
This commit is contained in:
Czar Echavez
2024-04-04 14:17:34 +01:00
parent 5408611c73
commit 383cf051df
3 changed files with 164 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/primaite/simulator/system/applications/database_client.py
View File

@@ -29,6 +29,8 @@ class DatabaseClient(Application):
_query_success_tracker: Dict[str, bool] = {}
_last_connection_successful: Optional[bool] = None
"""Keep track of connections that were established or verified during this step. Used for rewards."""
last_query_response: Optional[Dict] = None
"""Keep track of the latest query response. Used to determine rewards."""
def __init__(self, **kwargs):
kwargs["name"] = "DatabaseClient"
@@ -219,6 +221,9 @@ class DatabaseClient(Application):
if not self._can_perform_action():
return False
# reset last query response
self.last_query_response = None
if connection_id is None:
if self.connections:
connection_id = list(self.connections.keys())[-1]
@@ -252,6 +257,7 @@ class DatabaseClient(Application):
# add connection
self.add_connection(connection_id=payload.get("connection_id"), session_id=session_id)
elif payload["type"] == "sql":
self.last_query_response = payload
query_id = payload.get("uuid")
status_code = payload.get("status_code")
self._query_success_tracker[query_id] = status_code == 200

6
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -204,7 +204,7 @@ class DatabaseService(Service):
if not self.db_file:
self.sys_log.info(f"{self.name}: Failed to run {query} because the database file is missing.")
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
if query == "SELECT":
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
@@ -216,7 +216,7 @@ class DatabaseService(Service):
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
elif query == "DELETE":
self.db_file.health_status = FileSystemItemHealthStatus.COMPROMISED
return {
@@ -236,7 +236,7 @@ class DatabaseService(Service):
"connection_id": connection_id,
}
else:
return {"status_code": 404, "data": False}
return {"status_code": 404, "type": "sql", "data": False}
elif query == "SELECT * FROM pg_stat_activity":
# Check if the connection is active.
if self.health_state_actual == SoftwareHealthState.GOOD: