Nexus/web/agent/agent.py

"""
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 httpx
import psutil
from fastapi import FastAPI, HTTPException, Header, Request
from fastapi.responses import JSONResponse

# --- Configuration ---
CONFIG_PATH = Path(__file__).parent / "config.json"


def load_config():
    if CONFIG_PATH.exists():
        with open(CONFIG_PATH) as f:
            return json.load(f)
    example = Path(__file__).parent / "config.example.json"
    if example.exists():
        with open(example) 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)

# --- 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")

# --- FastAPI App ---

app = FastAPI(title="Nexus Agent", version="2.0.0")


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


def verify_api_key(request: Request, x_api_key: str = Header(default="")):
    # IP allowlist check first
    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),
            "memory_total_gb": round(memory.total / (1024**3), 2),
            "memory_used_gb": round(memory.used / (1024**3), 2),
            "disk_total_gb": round(disk.total / (1024**3), 2),
            "disk_used_gb": round(disk.used / (1024**3), 2),
            "uptime_seconds": int(time.time() - psutil.boot_time()),
            "agent_time": datetime.now(timezone.utc).isoformat(),
        }

        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, x_api_key: str = Header(default="")):
    """Execute a shell command — wait=true (default) blocks; wait=false returns pid for /exec/wait."""
    verify_api_key(x_api_key)

    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, x_api_key: str = Header(default="")):
    """Kill a process started with wait=false."""
    verify_api_key(x_api_key)
    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, x_api_key: str = Header(default="")):
    """Wait for a process started with wait=false."""
    verify_api_key(x_api_key)
    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, x_api_key: str = Header(default="")):
    """Nexus pushes new config — update config.json and reload"""
    verify_api_key(x_api_key)
    try:
        now = datetime.now(timezone.utc).isoformat()
        data["updated_at"] = now
        data["updated_by"] = "nexus"
        with open(CONFIG_PATH, "w") as f:
            json.dump(data, f, indent=2)

        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)

        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 _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"]


async def send_heartbeat():
    """Send heartbeat to Nexus central — smart: change >10%, over threshold, or every 10min."""
    global _last_metrics, _last_full_sync
    while True:
        try:
            if not CENTRAL_URL or not SERVER_ID:
                await asyncio.sleep(HEARTBEAT_INTERVAL)
                continue

            cpu = psutil.cpu_percent()
            mem = psutil.virtual_memory().percent
            disk = psutil.disk_usage("/").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)

            if changed or force_sync or over_threshold:
                _last_metrics = {"cpu": cpu, "memory": mem, "disk": disk}
                if force_sync:
                    _last_full_sync = now_ts

                payload = {
                    "server_id": SERVER_ID,
                    "is_online": True,
                    "agent_version": "2.0.0",
                    "system_info": {
                        "cpu_usage": cpu,
                        "mem_usage": mem,
                        "disk_usage": disk,
                        "hostname": platform.node(),
                        "platform": platform.platform(),
                        # agent_time: central uses this to detect clock drift
                        "agent_time": datetime.now(timezone.utc).isoformat(),
                    },
                }

                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},
                    )
                    if resp.status_code == 200:
                        logger.debug("Heartbeat sent")
                    else:
                        logger.warning(f"Heartbeat failed: {resp.status_code}")
            else:
                logger.debug("Metrics unchanged, skip heartbeat")

        except Exception as e:
            logger.warning(f"Heartbeat error: {e}")

        await asyncio.sleep(HEARTBEAT_INTERVAL)


# --- App Lifecycle ---

@app.on_event("startup")
async def startup():
    logger.info(f"Nexus Agent starting (server_id={SERVER_ID})")
    asyncio.create_task(send_heartbeat())

    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}")


@app.on_event("shutdown")
async def shutdown():
    logger.info("Nexus Agent stopped")


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)