""" Nexus Agent — Runs on each managed server Provides: health endpoint, heartbeat, command execution, config reload """ import sys import json import time import asyncio import logging import platform import subprocess import secrets import shlex from datetime import datetime, timezone from pathlib import Path from typing import Optional import ipaddress import socket import httpx from heartbeat_policy import heartbeat_should_stop from cpu_metrics import reset_watch_cpu_bootstrap, sample_cpu_percent, sample_watch_cpu_bundle AGENT_VERSION = "2.1.1" import psutil from fastapi import FastAPI, HTTPException, Header, Request, Depends from fastapi.responses import JSONResponse # --- Configuration --- CONFIG_PATH = Path(__file__).parent / "config.json" def load_config(): if CONFIG_PATH.exists(): with open(CONFIG_PATH, encoding="utf-8") as f: return json.load(f) example = Path(__file__).parent / "config.example.json" if example.exists(): with open(example, encoding="utf-8") as f: return json.load(f) return {} config = load_config() API_KEY = (config.get("api_key") or "").strip() if not API_KEY: print("FATAL: api_key must be set in config.json", file=sys.stderr) sys.exit(1) SERVER_ID = config.get("server_id", 0) CENTRAL_URL = config.get("central", {}).get("url", "") CENTRAL_API_KEY = config.get("central", {}).get("api_key", "") HEARTBEAT_INTERVAL = config.get("heartbeat_interval", 60) HEARTBEAT_KEEPALIVE_SEC = int(config.get("heartbeat_keepalive_sec", 30)) WATCH_INTERVAL_SEC = int(config.get("watch_interval_sec", 5)) # --- IP allowlist (only accept requests from trusted IPs) --- # In config.json: "allowed_ips": ["1.2.3.4", "10.0.0.0/8"] # Empty list = allow all (default for backward compat) _ALLOWED_IPS: list[str] = config.get("allowed_ips") or [] # --- Logging --- log_file = config.get("log_file", "/var/log/nexus-agent.log") log_level = config.get("log_level", "INFO") logging.basicConfig( level=getattr(logging, log_level, logging.INFO), format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.StreamHandler(sys.stdout), logging.FileHandler(log_file, mode="a") if log_file != "stdout" else logging.StreamHandler(sys.stdout), ], ) logger = logging.getLogger("nexus-agent") _heartbeat_task: Optional[asyncio.Task] = None # --- FastAPI App --- from contextlib import asynccontextmanager def _start_heartbeat_task() -> asyncio.Task: """Start heartbeat loop (idempotent if previous task still running).""" global _heartbeat_task if _heartbeat_task is not None and not _heartbeat_task.done(): return _heartbeat_task _heartbeat_task = asyncio.create_task(send_heartbeat()) return _heartbeat_task @asynccontextmanager async def agent_lifespan(app): """Application lifecycle: start heartbeat task on startup, clean up on shutdown.""" global _heartbeat_task logger.info(f"Nexus Agent starting (server_id={SERVER_ID})") _start_heartbeat_task() host = config.get("server", {}).get("host", "127.0.0.1") port = config.get("server", {}).get("port", 8601) logger.info(f"Agent listening on {host}:{port}") yield if _heartbeat_task is not None: _heartbeat_task.cancel() _heartbeat_task = None logger.info("Nexus Agent stopped") app = FastAPI(title="Nexus Agent", version=AGENT_VERSION, lifespan=agent_lifespan) def _ip_allowed(client_ip: str) -> bool: """Check client IP against allowed_ips list (supports CIDR and exact match).""" if not _ALLOWED_IPS: return True # Empty = allow all import ipaddress try: addr = ipaddress.ip_address(client_ip) for entry in _ALLOWED_IPS: entry = entry.strip() try: if "/" in entry: if addr in ipaddress.ip_network(entry, strict=False): return True elif addr == ipaddress.ip_address(entry): return True except ValueError: if client_ip == entry: return True except ValueError: return client_ip in _ALLOWED_IPS return False async def verify_api_key(request: Request, x_api_key: str = Header(default="")): """Dependency: IP allowlist + API key verification.""" client_ip = request.client.host if request.client else "" if not _ip_allowed(client_ip): logger.warning("Blocked request from %s (not in allowed_ips)", client_ip) raise HTTPException(status_code=403, detail="IP not allowed") if not secrets.compare_digest(x_api_key or "", API_KEY): raise HTTPException(status_code=401, detail="Invalid API key") # --- Health Endpoint --- @app.get("/health") async def health_check(request: Request): """Health check endpoint — called by Nexus central server""" client_ip = request.client.host if request.client else "" if not _ip_allowed(client_ip): raise HTTPException(status_code=403, detail="IP not allowed") try: cpu_percent = psutil.cpu_percent(interval=0.5) memory = psutil.virtual_memory() disk = psutil.disk_usage("/") system_info = { "hostname": platform.node(), "platform": platform.platform(), "python_version": platform.python_version(), "cpu_usage": cpu_percent, "mem_usage": memory.percent, "disk_usage": round(disk.percent, 1), "agent_time": datetime.now(timezone.utc).isoformat(), **_hardware_specs_dict(), } return {"status": "healthy", "system_info": system_info} except Exception as e: return {"status": "degraded", "error": str(e)} # --- Command Execution --- _running_procs: dict[int, asyncio.subprocess.Process] = {} def _proc_result(proc: asyncio.subprocess.Process, stdout: bytes, stderr: bytes) -> dict: exit_code = proc.returncode if proc.returncode is not None else -1 return { "status": "success" if exit_code == 0 else "failed", "stdout": stdout.decode(errors="replace")[:10000], "stderr": stderr.decode(errors="replace")[:10000], "exit_code": exit_code, "pid": proc.pid, } @app.post("/exec") async def exec_command(payload: dict, _auth=Depends(verify_api_key)): """Execute a shell command — wait=true (default) blocks; wait=false returns pid for /exec/wait.""" command = payload.get("command", "") timeout = payload.get("timeout", 30) sudo = payload.get("sudo", False) wait = payload.get("wait", True) if not command: raise HTTPException(status_code=400, detail="command is required") logger.warning( "Shell exec via subprocess_shell (server_id=%s, wait=%s, len=%s)", SERVER_ID, wait, len(command), ) if sudo: command = f"sudo -n sh -c {shlex.quote(command)}" proc = None try: proc = await asyncio.create_subprocess_shell( command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True, ) _running_procs[proc.pid] = proc if not wait: return { "status": "started", "pid": proc.pid, "stdout": "", "stderr": "", "exit_code": None, } stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout) _running_procs.pop(proc.pid, None) return _proc_result(proc, stdout, stderr) except asyncio.TimeoutError: if proc: proc.kill() _running_procs.pop(proc.pid, None) return { "status": "timeout", "stdout": "", "stderr": f"Command timed out after {timeout}s", "exit_code": -1, "pid": proc.pid if proc else None, } except Exception as e: return {"status": "error", "stdout": "", "stderr": str(e), "exit_code": -1} @app.post("/exec/kill") async def exec_kill(payload: dict, _auth=Depends(verify_api_key)): """Kill a process started with wait=false.""" pid = int(payload.get("pid", 0)) if pid <= 0: raise HTTPException(status_code=400, detail="pid required") proc = _running_procs.pop(pid, None) if proc and proc.returncode is None: proc.kill() try: await proc.wait() except Exception: pass return {"status": "killed", "pid": pid} raise HTTPException( status_code=404, detail="process not found or not started by this agent", ) @app.post("/exec/wait") async def exec_wait(payload: dict, _auth=Depends(verify_api_key)): """Wait for a process started with wait=false.""" pid = int(payload.get("pid", 0)) timeout = int(payload.get("timeout", 60)) if pid <= 0: raise HTTPException(status_code=400, detail="pid required") proc = _running_procs.get(pid) if not proc: raise HTTPException(status_code=404, detail="unknown pid (finished or invalid)") try: stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout) _running_procs.pop(pid, None) return _proc_result(proc, stdout, stderr) except asyncio.TimeoutError: _running_procs.pop(pid, None) if proc.returncode is None: proc.kill() try: await proc.wait() except Exception: pass return { "status": "timeout", "stdout": "", "stderr": f"Command timed out after {timeout}s", "exit_code": -1, "pid": pid, } # --- Config Reload --- @app.post("/config/reload") async def reload_config(data: dict, _auth=Depends(verify_api_key)): """Nexus pushes new config — update config.json and reload""" try: now = datetime.now(timezone.utc).isoformat() data["updated_at"] = now data["updated_by"] = "nexus" with open(CONFIG_PATH, "w", encoding="utf-8", newline="\n") as f: json.dump(data, f, indent=2, ensure_ascii=False) f.write("\n") global SERVER_ID, CENTRAL_URL, CENTRAL_API_KEY, config config = data SERVER_ID = config.get("server_id", SERVER_ID) CENTRAL_URL = config.get("central", {}).get("url", CENTRAL_URL) CENTRAL_API_KEY = config.get("central", {}).get("api_key", CENTRAL_API_KEY) if CENTRAL_URL and SERVER_ID and CENTRAL_API_KEY: _start_heartbeat_task() logger.info("Heartbeat task (re)started after config reload") return {"status": "ok", "message": "Config updated", "updated_at": now} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) # --- Heartbeat to Central Server --- _last_metrics = {} _last_full_sync = 0 def _cpu_model() -> str: proc = (platform.processor() or "").strip() if proc: return proc[:120] try: with open("/proc/cpuinfo", encoding="utf-8", errors="ignore") as f: for line in f: if line.lower().startswith("model name"): return line.split(":", 1)[1].strip()[:120] except Exception: pass return "" def _hardware_specs_dict() -> dict: """Static-ish capacity fields for heartbeat / health.""" vm = psutil.virtual_memory() du = psutil.disk_usage("/") swap = psutil.swap_memory() logical = psutil.cpu_count(logical=True) or 1 physical = psutil.cpu_count(logical=False) or logical return { "cpu_cores": logical, "cpu_cores_physical": physical, "cpu_model": _cpu_model(), "memory_total_gb": round(vm.total / (1024**3), 1), "memory_used_gb": round(vm.used / (1024**3), 1), "memory_available_gb": round(vm.available / (1024**3), 1), "disk_total_gb": round(du.total / (1024**3), 1), "disk_used_gb": round(du.used / (1024**3), 1), "disk_free_gb": round(du.free / (1024**3), 1), "swap_total_gb": round(swap.total / (1024**3), 1), "swap_used_gb": round(swap.used / (1024**3), 1), "swap_pct": round(swap.percent, 1), "uptime_seconds": int(time.time() - psutil.boot_time()), "os_release": _os_release(), } def _os_release() -> str: """Read PRETTY_NAME from /etc/os-release (e.g. 'Ubuntu 24.04 LTS').""" try: path = Path("/etc/os-release") if path.exists(): for line in path.read_text(encoding="utf-8").splitlines(): if line.startswith("PRETTY_NAME="): return line.split("=", 1)[1].strip().strip('"') except Exception: pass return "" def _alert_thresholds() -> dict: """Align with Nexus center defaults (config.py / settings table). Override in config.json if needed.""" t = config.get("alert_thresholds") or {} return { "cpu": float(t.get("cpu", 80)), "mem": float(t.get("mem", 80)), "disk": float(t.get("disk", 80)), } def _metrics_over_threshold(cpu: float, mem: float, disk: float) -> bool: """Same rule as center _detect_alerts: strictly greater than threshold.""" th = _alert_thresholds() return cpu > th["cpu"] or mem > th["mem"] or disk > th["disk"] def _local_private_ip() -> Optional[str]: """Default-route local IPv4 (for Telegram 内网 display).""" try: with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: sock.connect(("8.8.8.8", 80)) candidate = sock.getsockname()[0] addr = ipaddress.ip_address(candidate) if addr.is_private or addr.is_loopback: return candidate except OSError: pass return None async def _fetch_public_ip() -> Optional[str]: """Outbound probe for real WAN IP (Telegram 外网 display).""" endpoints = ( "https://api.ipify.org", "https://ifconfig.me/ip", ) try: async with httpx.AsyncClient(timeout=4.0) as client: for url in endpoints: try: resp = await client.get(url) if resp.status_code != 200: continue candidate = (resp.text or "").strip() addr = ipaddress.ip_address(candidate) if not (addr.is_private or addr.is_loopback or addr.is_reserved): return candidate except (ValueError, httpx.HTTPError): continue except httpx.HTTPError: pass return None _cached_public_ip: Optional[str] = None _cached_public_ip: Optional[str] = None _watch_active = False _watch_interval_sec = WATCH_INTERVAL_SEC _last_send_ts = 0.0 def _apply_watch_mode_from_response(body: Optional[dict]) -> None: """Switch between 60s idle heartbeat and 5s watch mode per central pin state.""" global _watch_active, _watch_interval_sec prev_active = _watch_active if not isinstance(body, dict): return active = body.get("watch_active") if active is not None: _watch_active = bool(active) interval = body.get("watch_interval_sec") if isinstance(interval, (int, float)) and int(interval) > 0: _watch_interval_sec = int(interval) if _watch_active and not prev_active: reset_watch_cpu_bootstrap() def _collect_load_avg() -> list[float]: try: return [float(x) for x in psutil.getloadavg()] except (AttributeError, OSError, NotImplementedError): return [0.0, 0.0, 0.0] def _build_system_info(*, watch_mode: bool) -> dict: if watch_mode: cpu, per_cpu = sample_watch_cpu_bundle() else: cpu = sample_cpu_percent(watch_mode=False) per_cpu = None vm = psutil.virtual_memory() du = psutil.disk_usage("/") system_info = { "cpu_usage": cpu, "mem_usage": vm.percent, "disk_usage": round(du.percent, 1), "hostname": platform.node(), "platform": platform.platform(), "agent_time": datetime.now(timezone.utc).isoformat(), **_hardware_specs_dict(), } if watch_mode: system_info["watch_mode"] = True system_info["load_avg"] = _collect_load_avg() system_info["per_cpu_pct"] = per_cpu try: net = psutil.net_io_counters() disk_io = psutil.disk_io_counters() if net: system_info["net_io"] = { "bytes_sent": int(net.bytes_sent), "bytes_recv": int(net.bytes_recv), } if disk_io: system_info["disk_io"] = { "read_bytes": int(disk_io.read_bytes), "write_bytes": int(disk_io.write_bytes), } except Exception: logger.debug("Failed to collect net/disk IO for heartbeat", exc_info=True) private_ip = _local_private_ip() if private_ip: system_info["private_ip"] = private_ip return system_info async def _enrich_system_info_network(system_info: dict, *, force_sync: bool) -> None: global _cached_public_ip if force_sync: fetched = await _fetch_public_ip() if fetched: _cached_public_ip = fetched if _cached_public_ip: system_info["public_ip"] = _cached_public_ip async def _post_heartbeat(system_info: dict) -> tuple[int, Optional[dict]]: payload = { "server_id": SERVER_ID, "is_online": True, "agent_version": AGENT_VERSION, "system_info": system_info, } async with httpx.AsyncClient(timeout=10) as client: resp = await client.post( f"{CENTRAL_URL}/api/agent/heartbeat", json=payload, headers={"X-API-Key": CENTRAL_API_KEY}, ) body = resp.json() if resp.status_code == 200 else None return resp.status_code, body async def send_heartbeat(): """Send heartbeat to Nexus central — idle: smart 60s; watch: 5s full metrics when pinned.""" global _last_metrics, _last_full_sync, _cached_public_ip, _last_send_ts _consecutive_failures = 0 while True: sleep_sec = _watch_interval_sec if _watch_active else HEARTBEAT_INTERVAL try: if not CENTRAL_URL or not SERVER_ID: await asyncio.sleep(HEARTBEAT_INTERVAL) continue if _watch_active: system_info = _build_system_info(watch_mode=True) await _enrich_system_info_network(system_info, force_sync=False) status_code, body = await _post_heartbeat(system_info) if heartbeat_should_stop(status_code, body): if status_code == 401: logger.error( "Heartbeat rejected by central (401): invalid API key or " "server_id=%s not configured with per-server agent_api_key. " "Stopping heartbeat — reconfigure agent key via admin UI.", SERVER_ID, ) else: reason = (body or {}).get("reason", "unknown") logger.error( "Heartbeat discarded by central: %s. " "This agent may be misconfigured (wrong server_id=%s). " "Stopping heartbeat to prevent noise.", reason, SERVER_ID, ) return if status_code == 200: _consecutive_failures = 0 _last_send_ts = time.time() _apply_watch_mode_from_response(body) logger.debug("Watch-mode heartbeat sent") else: _consecutive_failures += 1 logger.warning("Watch-mode heartbeat failed: %s", status_code) sleep_sec = _watch_interval_sec if _watch_active else HEARTBEAT_INTERVAL else: cpu = sample_cpu_percent(watch_mode=False) vm = psutil.virtual_memory() du = psutil.disk_usage("/") mem = vm.percent disk = du.percent now_ts = time.time() changed = ( abs(cpu - _last_metrics.get("cpu", 0)) >= 10 or abs(mem - _last_metrics.get("memory", 0)) >= 10 or abs(disk - _last_metrics.get("disk", 0)) >= 10 ) force_sync = (now_ts - _last_full_sync) > 600 over_threshold = _metrics_over_threshold(cpu, mem, disk) keepalive_due = (now_ts - _last_send_ts) >= HEARTBEAT_KEEPALIVE_SEC if changed or force_sync or over_threshold or keepalive_due: _last_metrics = {"cpu": cpu, "memory": mem, "disk": disk} if force_sync: _last_full_sync = now_ts system_info = _build_system_info(watch_mode=False) await _enrich_system_info_network(system_info, force_sync=force_sync) status_code, body = await _post_heartbeat(system_info) if heartbeat_should_stop(status_code, body): if status_code == 401: logger.error( "Heartbeat rejected by central (401): invalid API key or " "server_id=%s not configured with per-server agent_api_key. " "Stopping heartbeat — reconfigure agent key via admin UI.", SERVER_ID, ) else: reason = (body or {}).get("reason", "unknown") logger.error( "Heartbeat discarded by central: %s. " "This agent may be misconfigured (wrong server_id=%s). " "Stopping heartbeat to prevent noise.", reason, SERVER_ID, ) return if status_code == 200: _consecutive_failures = 0 _last_send_ts = now_ts _apply_watch_mode_from_response(body) if _watch_active: logger.info("Central enabled watch mode — switching to %ss interval", _watch_interval_sec) else: logger.debug("Heartbeat sent") else: _consecutive_failures += 1 logger.warning("Heartbeat failed: %s", status_code) else: logger.debug("Metrics unchanged, skip heartbeat") sleep_sec = _watch_interval_sec if _watch_active else HEARTBEAT_INTERVAL except Exception as e: _consecutive_failures += 1 logger.warning("Heartbeat error: %s", e) if _consecutive_failures > 0: backoff = min(HEARTBEAT_INTERVAL * (2 ** min(_consecutive_failures, 5)), 300) logger.info("Backing off %ss after %s consecutive failures", backoff, _consecutive_failures) await asyncio.sleep(backoff) else: await asyncio.sleep(sleep_sec) if __name__ == "__main__": host = config.get("server", {}).get("host", "127.0.0.1") port = config.get("server", {}).get("port", 8601) import uvicorn uvicorn.run(app, host=host, port=port)