import os from dotenv import load_dotenv # Dynamically select configuration file based on APP_ENV app_env = os.getenv("APP_ENV", "development") if app_env == "production": load_dotenv(".env.production") else: load_dotenv(".env.development") # Fallback to load default .env if any variables are not yet defined load_dotenv() from fastapi import FastAPI, HTTPException from fastapi.responses import PlainTextResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Dict, Any, List import uvicorn app = FastAPI(title="RMM Central API") # Enable CORS for frontend dashboard queries (React dev server runs on a separate port) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) import os import json from datetime import datetime # In-memory storage for client hardware telemetry client_telemetry = {} from pymongo import MongoClient # Database Connection: Configured via environment variable with local fallback MONGODB_URI = os.getenv("MONGODB_URI", "mongodb://localhost:27017") DB_NAME = os.getenv("MONGODB_DB", "rmm_db") mongo_client = MongoClient(MONGODB_URI) db = mongo_client[DB_NAME] clients_collection = db["clients"] logs_collection = db["logs"] MAX_ENTRIES_PER_CLIENT = 100 # Keep the latest 100 historical logs per client node ROCKETCHAT_WEBHOOK_URL = os.getenv("ROCKETCHAT_WEBHOOK_URL", "") # Alert states are now stored directly in the MongoDB client documents # to support multi-worker environments safely. def send_rocketchat_notification(text: str, color: str = "#808080"): """ Dispatches a formatted notification payload to the Rocket.Chat incoming webhook. """ if not ROCKETCHAT_WEBHOOK_URL: try: print(f"[Rocket.Chat Simulation] {text}") except UnicodeEncodeError: # Fallback for Windows consoles that do not support printing unicode emojis sanitized_text = text.encode('ascii', errors='backslashreplace').decode('ascii') print(f"[Rocket.Chat Simulation] {sanitized_text}") return payload = { "text": text, "attachments": [ { "color": color, "ts": datetime.now().isoformat() } ] } try: import urllib.request import json req = urllib.request.Request( ROCKETCHAT_WEBHOOK_URL, data=json.dumps(payload).encode('utf-8'), headers={'Content-Type': 'application/json'} ) with urllib.request.urlopen(req, timeout=3.0) as response: pass except Exception as e: print(f"[!] Error sending Rocket.Chat notification: {e}") def append_to_logs(log_type: str, client_id: str, detail: Any): """ Appends a new event log grouped by client_id in MongoDB logs collection. Ensures fair-share log limits per system so noisy clients never overwrite others. """ try: entry = { "client_id": client_id, "timestamp": datetime.now().isoformat(), "type": log_type, "detail": detail } logs_collection.insert_one(entry) # Enforce fair-share logging limits per system (FIFO circular buffer) count = logs_collection.count_documents({"client_id": client_id}) if count > MAX_ENTRIES_PER_CLIENT: oldest_docs = logs_collection.find( {"client_id": client_id}, {"_id": 1} ).sort("timestamp", 1).limit(count - MAX_ENTRIES_PER_CLIENT) ids_to_delete = [doc["_id"] for doc in oldest_docs] if ids_to_delete: logs_collection.delete_many({"_id": {"$in": ids_to_delete}}) 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"): 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() 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*$"}}) # 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 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 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) 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}") class TelemetryPayload(BaseModel): cpu_percent: float memory_percent: float memory_total_gb: float memory_free_gb: float disks: List[Dict[str, Any]] gpus: List[Dict[str, Any]] class CommandResultPayload(BaseModel): command: str returncode: int stdout: str stderr: str @app.get("/api/get-commands") async def get_whitelisted_commands(platform: str): """ Endpoint for remote agents to fetch their OS-specific whitelisted commands. """ commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json") try: with open(commands_file, "r") as f: data = json.load(f) return data.get(platform, {}) 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): """ 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 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} # 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) return {"command": cmd} @app.post("/api/schedule-command") async def schedule_command(client_id: str, command: str): """ 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 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") 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 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() # 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 for disk in payload.disks: mount = disk.get("mount", "/") percent = disk.get("percent", 0.0) if percent >= 80.0: if mount not in storage_alerts: storage_alerts.append(mount) storage_modified = True send_rocketchat_notification( text=f"⚠️ **Storage Warning:** Client `{client_id}` disk `{mount}` is at **{percent}%** capacity.", color="#f39c12" ) else: if mount in storage_alerts: storage_alerts.remove(mount) storage_modified = True send_rocketchat_notification( text=f"✅ **Storage Recovered:** Client `{client_id}` disk `{mount}` has cleared warning state and is at **{percent}%**.", 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 gpu_modified = True send_rocketchat_notification( text=f"🚨 **GPU Failure:** Client `{client_id}` is not reporting any GPU data! (nvidia-smi has failed or is missing)", color="#e74c3c" ) else: if gpu_alert_sent: gpu_alert_sent = False gpu_modified = True send_rocketchat_notification( text=f"✅ **GPU Recovered:** Client `{client_id}` is reporting GPU data successfully again.", color="#2ecc71" ) if gpu_modified: clients[client_id]["gpu_alert_sent"] = gpu_alert_sent if storage_modified or gpu_modified: save_clients(clients) data = payload.dict() client_telemetry[client_id] = data # Temporarily print the data to the console for the user to see! 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"} @app.get("/api/telemetry") async def get_all_telemetry(): """ Endpoint to view the live dashboard data of all clients. """ clients = load_clients() return {cid: info.get("telemetry", {}) for cid, info in clients.items() if info.get("telemetry")} @app.get("/api/clients") async def get_clients_api(): """ Endpoint for the React UI to fetch the live active client registry with computed statuses. """ return load_clients() @app.get("/api/logs") async def get_logs_history(): """ Endpoint to retrieve logs history grouped by client_id. """ try: cursor = logs_collection.find().sort("timestamp", 1) grouped_logs = {} for doc in cursor: cid = doc.get("client_id") if not cid: continue if cid not in grouped_logs: grouped_logs[cid] = [] entry = { "timestamp": doc.get("timestamp"), "type": doc.get("type"), "detail": doc.get("detail") } grouped_logs[cid].append(entry) return grouped_logs except Exception as e: print(f"[!] Error loading logs from MongoDB: {e}") return {} @app.post("/api/command-result") 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"} @app.get("/api/get-raw-commands") async def get_raw_commands(): """ Endpoint for the UI to load the complete commands.json whitelist file. """ commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json") try: with open(commands_file, "r") as f: return json.load(f) except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to read commands: {e}") @app.post("/api/save-commands") async def save_commands(commands: Dict[str, Any]): """ Endpoint for the UI to save changes back to commands.json. """ commands_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "commands.json") try: with open(commands_file, "w") as f: json.dump(commands, f, indent=2) return {"status": "success", "message": "commands.json saved successfully"} except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to save commands: {e}") @app.get("/metrics", response_class=PlainTextResponse) async def get_prometheus_metrics(): """ Exposes the latest client telemetry in standard Prometheus text format. """ clients = load_clients() metrics_lines = [] for client_id, data in clients.items(): # Active status (1.0 for active/online, 0.0 for inactive/offline) active_val = 1.0 if data.get("active", False) else 0.0 metrics_lines.append(f'rmm_client_active{{client_id="{client_id}"}} {active_val}') # Telemetry stats telemetry = data.get("telemetry") if telemetry: cpu = telemetry.get("cpu_percent", 0.0) ram = telemetry.get("memory_percent", 0.0) ram_free = telemetry.get("memory_free_gb", 0.0) ram_total = telemetry.get("memory_total_gb", 0.0) metrics_lines.append(f'rmm_cpu_utilization{{client_id="{client_id}"}} {cpu}') metrics_lines.append(f'rmm_memory_utilization{{client_id="{client_id}"}} {ram}') metrics_lines.append(f'rmm_memory_free_bytes{{client_id="{client_id}"}} {ram_free * (1024**3)}') metrics_lines.append(f'rmm_memory_total_bytes{{client_id="{client_id}"}} {ram_total * (1024**3)}') # Disk mounts for disk in telemetry.get("disks", []): mount = disk.get("mount", "/") disk_percent = disk.get("percent", 0.0) metrics_lines.append(f'rmm_disk_utilization{{client_id="{client_id}",mount="{mount}"}} {disk_percent}') # GPU utilization and temps for gpu in telemetry.get("gpus", []): gpu_name = gpu.get("name", "GPU") # Utilization percent parsing gpu_util_str = str(gpu.get("utilization", "0")) gpu_util = float(gpu_util_str.replace("%", "").strip()) # Temperature parsing gpu_temp_str = str(gpu.get("temp", "0")) gpu_temp = float(gpu_temp_str.replace("C", "").strip()) metrics_lines.append(f'rmm_gpu_utilization{{client_id="{client_id}",gpu_name="{gpu_name}"}} {gpu_util}') metrics_lines.append(f'rmm_gpu_temperature{{client_id="{client_id}",gpu_name="{gpu_name}"}} {gpu_temp}') return "\n".join(metrics_lines) + "\n" if __name__ == "__main__": print("[*] Starting FastAPI Central Server...") # uvicorn runs the FastAPI app on port 8000 uvicorn.run("central_api_prototype:app", host="0.0.0.0", port=8000, reload=True)