diff --git a/central_api_prototype.py b/central_api_prototype.py index 5cb1f50..64507d1 100644 --- a/central_api_prototype.py +++ b/central_api_prototype.py @@ -176,18 +176,31 @@ 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 clients[client_id].get("last_down_alert_time"): +def mark_client_active(client_id: str): + try: + now_str = datetime.now().isoformat() + # Atomically check if transitioned from offline to online (active is not True) + old_doc = clients_collection.find_one_and_update( + {"_id": client_id, "active": {"$ne": True}}, + {"$set": {"active": True, "last_seen": now_str, "last_down_alert_time": None}}, + return_document=False + ) + if old_doc: + # Client transitioned from inactive/None to active! + # Send UP recovery alert only if they had a down alert time logged + if old_doc.get("last_down_alert_time"): send_rocketchat_notification( text=f"✅ **System UP:** Client `{client_id}` has recovered and is back online.", color="#2ecc71" ) - clients[client_id]["last_down_alert_time"] = None - clients[client_id]["active"] = True - clients[client_id]["last_seen"] = datetime.now().isoformat() + else: + # Client is already active, just update their last_seen timestamp + clients_collection.update_one( + {"_id": client_id}, + {"$set": {"last_seen": now_str}} + ) + except Exception as e: + print(f"[!] Error marking client active: {e}") def load_clients() -> Dict[str, Any]: try: @@ -195,92 +208,93 @@ def load_clients() -> Dict[str, Any]: clients_collection.delete_many({"_id": {"$in": ["", None]}}) clients_collection.delete_many({"_id": {"$regex": "^\\s*$"}}) - - # Load all documents cursor = clients_collection.find() data = {} for doc in cursor: cid = doc["_id"] - data[cid] = {k: v for k, v in doc.items() if k != "_id"} - - # Compute active status dynamically based on 30-second heartbeat check-ins - modified = False - for cid, info in data.items(): - last_seen_str = info.get("last_seen") - was_active = info.get("active", False) - is_active = False + # Compute active status dynamically for reads, but do NOT write back to database or alert + last_seen_str = doc.get("last_seen") + is_active = doc.get("active", False) if last_seen_str: try: last_seen_dt = datetime.fromisoformat(last_seen_str) - if (datetime.now() - last_seen_dt).total_seconds() < 30: - is_active = True + if (datetime.now() - last_seen_dt).total_seconds() >= 30: + is_active = False except Exception: pass - if was_active != is_active: - info["active"] = is_active - modified = True - - # Active status transitions alerts - if is_active: - # Transitions from offline (DOWN) to online (UP) - if info.get("last_down_alert_time"): - send_rocketchat_notification( - text=f"✅ **System UP:** Client `{cid}` has recovered and is back online.", - color="#2ecc71" - ) - info["last_down_alert_time"] = None - else: - # Transitions from online (UP) to offline (DOWN) for the first time - now = datetime.now() - info["last_down_alert_time"] = now.isoformat() - send_rocketchat_notification( - text=f"🚨 **System DOWN:** Client `{cid}` has missed heartbeats for over 30 seconds.", - color="#e74c3c" - ) - else: - # No transition (state remains the same) - if not is_active: - # Client remains offline. Check if we should send a repeating reminder. - # We ONLY send a reminder if it is already tracked in last_down_alert_time. - # This prevents sending alerts for historically offline systems at server boot! - last_down_str = info.get("last_down_alert_time") - if last_down_str: - now = datetime.now() - try: - last_down_dt = datetime.fromisoformat(last_down_str) - elapsed_seconds = (now - last_down_dt).total_seconds() - if elapsed_seconds >= 3600: # 1 hour (3600 seconds) - info["last_down_alert_time"] = now.isoformat() - send_rocketchat_notification( - text=f"🚨 **System STILL DOWN:** Client `{cid}` remains offline (reminder sent every hour).", - color="#e74c3c" - ) - modified = True - except Exception: - pass - - if modified: - save_clients(data) + info = {k: v for k, v in doc.items() if k != "_id"} + info["active"] = is_active + data[cid] = info return data except Exception as e: print(f"[!] Error loading clients from MongoDB: {e}") return {} -def save_clients(data: Dict[str, Any]): - try: - for cid, info in data.items(): - # Use upsert to update existing client or insert new client automatically - clients_collection.update_one( - {"_id": cid}, - {"$set": info}, - upsert=True - ) - except Exception as e: - print(f"[!] Error saving clients to MongoDB: {e}") +import asyncio + +async def monitor_heartbeats_loop(): + while True: + try: + now = datetime.now() + timeout_time = (now - timedelta(seconds=30)).isoformat() + + # Find clients that are active in DB but haven't checked in for 30s + cursor = clients_collection.find({ + "active": True, + "last_seen": {"$lt": timeout_time} + }) + for doc in cursor: + cid = doc["_id"] + now_str = now.isoformat() + + # Atomically update to active = False + # Ensures only one request/worker triggers the state transition and sends the alert + old_doc = clients_collection.find_one_and_update( + {"_id": cid, "active": True, "last_seen": doc["last_seen"]}, + {"$set": {"active": False, "last_down_alert_time": now_str}}, + return_document=False + ) + if old_doc: + send_rocketchat_notification( + text=f"🚨 **System DOWN:** Client `{cid}` has missed heartbeats for over 30 seconds.", + color="#e74c3c" + ) + + # Hourly reminders for STILL DOWN clients + reminder_time = (now - timedelta(hours=1)).isoformat() + cursor_still_down = clients_collection.find({ + "active": False, + "last_down_alert_time": {"$lt": reminder_time} + }) + for doc in cursor_still_down: + cid = doc["_id"] + old_last_down = doc["last_down_alert_time"] + now_str = now.isoformat() + + # Atomically update reminder timestamp + old_doc = clients_collection.find_one_and_update( + {"_id": cid, "active": False, "last_down_alert_time": old_last_down}, + {"$set": {"last_down_alert_time": now_str}}, + return_document=False + ) + if old_doc: + send_rocketchat_notification( + text=f"🚨 **System STILL DOWN:** Client `{cid}` remains offline (reminder sent every hour).", + color="#e74c3c" + ) + except Exception as e: + print(f"[!] Error in heartbeat monitoring loop: {e}") + + await asyncio.sleep(10) + +@app.on_event("startup") +async def startup_event(): + asyncio.create_task(monitor_heartbeats_loop()) + print("[*] Heartbeat monitoring background task initialized.") class TelemetryPayload(BaseModel): cpu_percent: float @@ -320,7 +334,6 @@ async def get_whitelisted_commands(platform: str): except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to load commands: {e}") - @app.get("/api/get-command") async def get_command(client_id: str): """ @@ -330,25 +343,29 @@ async def get_command(client_id: str): if not client_id or not client_id.strip(): raise HTTPException(status_code=400, detail="client_id cannot be empty") - clients = load_clients() + now_str = datetime.now().isoformat() + # Dynamic fleet auto-registration + clients_collection.update_one( + {"_id": client_id}, + {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}}, + upsert=True + ) - # 1. Auto-register new client if we haven't seen it yet - if client_id not in clients: - print(f"[+] Dynamic fleet auto-registration: Registered new agent '{client_id}'") - clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True} + # Get and reset command atomically + updated_doc = clients_collection.find_one_and_update( + {"_id": client_id}, + {"$set": {"pending_command": "none"}}, + return_document=False # Returns the state before update + ) + + cmd = "none" + if updated_doc: + cmd = updated_doc.get("pending_command", "none") - # 2. Fetch the command - cmd = clients[client_id].get("pending_command", "none") - - # Log if a valid pending command was polled if cmd != "none": append_to_logs("command_polled", client_id, {"command": cmd}) - # 3. Reset the command to "none", mark active, and update last_seen timestamp - clients[client_id]["pending_command"] = "none" - mark_client_active(client_id, clients) - - save_clients(clients) + mark_client_active(client_id) return {"command": cmd} @app.post("/api/schedule-command") @@ -357,20 +374,18 @@ async def schedule_command(client_id: str, command: str, current_user: str = Dep Endpoint for your Dashboard/UI to schedule a new command. Supports a comma-separated list of client IDs for batch fleet updates. """ - clients = load_clients() target_ids = [cid.strip() for cid in client_id.split(",") if cid.strip()] - if not target_ids: raise HTTPException(status_code=400, detail="No target client IDs specified") for cid in target_ids: - if cid not in clients: - clients[cid] = {"pending_command": "none", "last_seen": None, "active": False} - clients[cid]["pending_command"] = command - # Log command scheduling per system + clients_collection.update_one( + {"_id": cid}, + {"$set": {"pending_command": command}}, + upsert=True + ) append_to_logs("command_scheduled", cid, {"command": command}) - save_clients(clients) return {"status": "success", "message": f"Command '{command}' scheduled for {', '.join(target_ids)}"} @app.post("/api/telemetry") @@ -381,24 +396,25 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload): if not client_id or not client_id.strip(): raise HTTPException(status_code=400, detail="client_id cannot be empty") - clients = load_clients() + now_str = datetime.now().isoformat() + # Auto-register client if we see it through telemetry first + clients_collection.update_one( + {"_id": client_id}, + {"$setOnInsert": {"pending_command": "none", "active": True, "last_seen": now_str}}, + upsert=True + ) - # Auto-register new client if we see it through telemetry first - if client_id not in clients: - print(f"[+] Dynamic fleet auto-registration: Registered new agent '{client_id}' through telemetry") - clients[client_id] = {"pending_command": "none", "last_seen": None, "active": True} - - # Update last_seen timestamp, active status, and store the current telemetry details - mark_client_active(client_id, clients) - clients[client_id]["telemetry"] = payload.dict() + mark_client_active(client_id) + + # Get current alert states from database to prevent duplicate alerts + client_doc = clients_collection.find_one({"_id": client_id}) or {} + storage_alerts = client_doc.get("active_storage_alerts", []) + gpu_alert_sent = client_doc.get("gpu_alert_sent", False) - # Log telemetry history event - append_to_logs("telemetry", client_id, payload.dict()) - - # Storage limits warning checker - storage_alerts = clients[client_id].get("active_storage_alerts", []) storage_modified = False + gpu_modified = False + # Storage warning checker for disk in payload.disks: mount = disk.get("mount", "/") percent = disk.get("percent", 0.0) @@ -420,13 +436,7 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload): color="#2ecc71" ) - if storage_modified: - clients[client_id]["active_storage_alerts"] = storage_alerts - # GPU Failure Monitor - gpu_alert_sent = clients[client_id].get("gpu_alert_sent", False) - gpu_modified = False - if len(payload.gpus) == 0: if not gpu_alert_sent: gpu_alert_sent = True @@ -444,11 +454,39 @@ async def receive_telemetry(client_id: str, payload: TelemetryPayload): color="#2ecc71" ) + update_fields = { + "telemetry": payload.dict() + } + if storage_modified: + update_fields["active_storage_alerts"] = storage_alerts if gpu_modified: - clients[client_id]["gpu_alert_sent"] = gpu_alert_sent + update_fields["gpu_alert_sent"] = gpu_alert_sent - if storage_modified or gpu_modified: - save_clients(clients) + clients_collection.update_one( + {"_id": client_id}, + {"$set": update_fields} + ) + + # Log telemetry history event + append_to_logs("telemetry", client_id, payload.dict()) + + data = payload.dict() + client_telemetry[client_id] = data + + # Log received stats + print(f"\n[TELEMETRY RECEIVED] from {client_id}:") + print(f" CPU: {data['cpu_percent']}%") + print(f" RAM: {data['memory_percent']}% ({data['memory_free_gb']} GB Free / {data['memory_total_gb']} GB Total)") + + for disk in data['disks']: + print(f" Drive {disk['mount']}: {disk['percent']}% Used ({disk['free_gb']} GB Free / {disk['total_gb']} GB Total)") + + for gpu in data['gpus']: + power_str = f"Power: {gpu.get('power_draw', 'N/A')} / {gpu.get('power_limit', 'N/A')}" + fan_str = f"Fan: {gpu.get('fan_speed', 'N/A')}" + print(f" GPU [{gpu['name']}]: Core: {gpu['utilization']}, Temp: {gpu['temp']}, VRAM: {gpu['memory_used']}/{gpu['memory_total']}, {power_str}, {fan_str}") + + return {"status": "success"} data = payload.dict() client_telemetry[client_id] = data