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, Depends, Header from fastapi.responses import PlainTextResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import Dict, Any, List import uvicorn import base64 import hmac import hashlib import json from datetime import datetime, timedelta 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=["*"], ) # JWT-like HMAC Secure Token Utilities JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY", "seekright_rmm_secret_key_2026_default") def generate_token(username: str) -> str: # Expire in 1 day (24 hours) expiry = (datetime.utcnow() + timedelta(days=1)).isoformat() payload = { "username": username, "expires": expiry } payload_json = json.dumps(payload) payload_b64 = base64.urlsafe_b64encode(payload_json.encode('utf-8')).decode('utf-8').rstrip('=') sig = hmac.new(JWT_SECRET_KEY.encode('utf-8'), payload_b64.encode('utf-8'), hashlib.sha256).hexdigest() return f"{payload_b64}.{sig}" def verify_token(token: str) -> bool: try: parts = token.split(".") if len(parts) != 2: return False payload_b64, sig = parts # Verify signature expected_sig = hmac.new(JWT_SECRET_KEY.encode('utf-8'), payload_b64.encode('utf-8'), hashlib.sha256).hexdigest() if not hmac.compare_digest(sig, expected_sig): return False # Decode and check expiry padding = 4 - (len(payload_b64) % 4) if padding < 4: payload_b64 += "=" * padding payload_json = base64.urlsafe_b64decode(payload_b64.encode('utf-8')).decode('utf-8') payload = json.loads(payload_json) expires_dt = datetime.fromisoformat(payload["expires"]) if datetime.utcnow() > expires_dt: return False # Expired return payload["username"] == "root" except Exception: return False # FastAPI dependency to secure UI endpoints async def get_current_user(authorization: str = Header(None)): if not authorization or not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Missing or invalid authentication credentials") token = authorization.split(" ")[1] if not verify_token(token): raise HTTPException(status_code=401, detail="Authentication token is invalid or has expired") return "root" 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 class LoginPayload(BaseModel): username: str password: str @app.post("/api/login") async def login(payload: LoginPayload): if payload.username == "root" and payload.password == "seekright159@": token = generate_token("root") return {"token": token} raise HTTPException(status_code=401, detail="Invalid username or password") @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, current_user: str = Depends(get_current_user)): """ 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(current_user: str = Depends(get_current_user)): """ 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(current_user: str = Depends(get_current_user)): """ 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(current_user: str = Depends(get_current_user)): """ 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(current_user: str = Depends(get_current_user)): """ 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], current_user: str = Depends(get_current_user)): """ 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)