feat: implement async command execution, improve client status tracking, and expand takeleap diagnostic commands.

central_api_prototype.py

@@ -123,8 +123,25 @@ def append_to_logs(log_type: str, client_id: str, detail: Any):
     except Exception as e:
         print(f"[!] Error writing log to MongoDB: {e}")
 
+def mark_client_active(client_id: str, clients: Dict[str, Any]):
+    if client_id in clients:
+        was_active = clients[client_id].get("active", False)
+        if not was_active:
+            if client_id in last_down_notification_times:
+                send_rocketchat_notification(
+                    text=f"✅ **System UP:** Client `{client_id}` has recovered and is back online.",
+                    color="#2ecc71"
+                )
+                last_down_notification_times.pop(client_id, None)
+        clients[client_id]["active"] = True
+        clients[client_id]["last_seen"] = datetime.now().isoformat()
+
 def load_clients() -> Dict[str, Any]:
     try:
+        # Clean up any ghost clients with empty or whitespace-only keys
+        clients_collection.delete_many({"_id": {"$in": ["", None]}})
+        clients_collection.delete_many({"_id": {"$regex": "^\\s*$"}})
+
         # If collection is empty, seed it with default registry
         if clients_collection.count_documents({}) == 0:
             default_data = {
@@ -248,6 +265,9 @@ async def get_command(client_id: str):
     Endpoint for clients to poll for pending commands.
     Client Agent hits this endpoint to ask: "Do I have any work to do?"
     """
+    if not client_id or not client_id.strip():
+        raise HTTPException(status_code=400, detail="client_id cannot be empty")
+        
     clients = load_clients()
     
     # 1. Auto-register new client if we haven't seen it yet
@@ -264,8 +284,7 @@ async def get_command(client_id: str):
         
     # 3. Reset the command to "none", mark active, and update last_seen timestamp
     clients[client_id]["pending_command"] = "none"
-    clients[client_id]["active"] = True
-    clients[client_id]["last_seen"] = datetime.now().isoformat()
+    mark_client_active(client_id, clients)
     
     save_clients(clients)
     return {"command": cmd}
@@ -297,6 +316,9 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
     """
     Endpoint for agents to push their live hardware telemetry.
     """
+    if not client_id or not client_id.strip():
+        raise HTTPException(status_code=400, detail="client_id cannot be empty")
+        
     clients = load_clients()
     
     # Auto-register new client if we see it through telemetry first
@@ -305,8 +327,7 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload):
         clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True}
         
     # Update last_seen timestamp, active status, and store the current telemetry details
-    clients[client_id]["last_seen"] = datetime.now().isoformat()
-    clients[client_id]["active"] = True
+    mark_client_active(client_id, clients)
     clients[client_id]["telemetry"] = payload.dict()
     save_clients(clients)
     
@@ -397,6 +418,9 @@ async def receive_command_result(client_id: str, payload: CommandResultPayload):
     """
     Endpoint for remote agents to push execution results back to the central server.
     """
+    if not client_id or not client_id.strip():
+        raise HTTPException(status_code=400, detail="client_id cannot be empty")
+        
     append_to_logs("command_result", client_id, payload.dict())
     return {"status": "success"}