Files
rmm-backend/central_api_prototype.py

639 lines
24 KiB
Python

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):
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"
)
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:
# 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"]
# 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 = False
except Exception:
pass
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 {}
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
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")
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
)
# 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")
if cmd != "none":
append_to_logs("command_polled", client_id, {"command": cmd})
mark_client_active(client_id)
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.
"""
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:
clients_collection.update_one(
{"_id": cid},
{"$set": {"pending_command": command}},
upsert=True
)
append_to_logs("command_scheduled", cid, {"command": command})
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")
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
)
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)
storage_modified = False
gpu_modified = False
# Storage warning checker
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"
)
# GPU Failure Monitor
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"
)
update_fields = {
"telemetry": payload.dict()
}
if storage_modified:
update_fields["active_storage_alerts"] = storage_alerts
if gpu_modified:
update_fields["gpu_alert_sent"] = gpu_alert_sent
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
# 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)