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