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

from heartbeat_policy import heartbeat_should_stop

AGENT_VERSION = "2.0.0"
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)

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

        os_release = _os_release()

        system_info = {
            "hostname": platform.node(),
            "platform": platform.platform(),
            "os_release": os_release,
            "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, _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 _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"]


async def send_heartbeat():
    """Send heartbeat to Nexus central — smart: change >10%, over threshold, or every 10min.
    Uses exponential backoff on consecutive failures (max 5 min)."""
    global _last_metrics, _last_full_sync
    _consecutive_failures = 0
    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": AGENT_VERSION,
                    "system_info": {
                        "cpu_usage": cpu,
                        "mem_usage": mem,
                        "disk_usage": disk,
                        "hostname": platform.node(),
                        "platform": platform.platform(),
                        "os_release": _os_release(),
                        # 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},
                    )
                    body = resp.json() if resp.status_code == 200 else None
                    if heartbeat_should_stop(resp.status_code, body):
                        if resp.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 resp.status_code == 200:
                        _consecutive_failures = 0
                        logger.debug("Heartbeat sent")
                    else:
                        _consecutive_failures += 1
                        logger.warning(f"Heartbeat failed: {resp.status_code}")
            else:
                logger.debug("Metrics unchanged, skip heartbeat")

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

        # Exponential backoff: base interval × 2^min(failures, 5), capped at 5 min
        if _consecutive_failures > 0:
            backoff = min(HEARTBEAT_INTERVAL * (2 ** min(_consecutive_failures, 5)), 300)
            logger.info(f"Backing off {backoff}s after {_consecutive_failures} consecutive failures")
            await asyncio.sleep(backoff)
        else:
            await asyncio.sleep(HEARTBEAT_INTERVAL)


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)