Nexus/server/api/agent.py

"""Nexus — Agent API Routes (Heartbeat receiver + exec endpoint)
Presentation layer — Agent servers call these endpoints to report status and receive commands.

Heartbeat flow:
  1. Agent sends heartbeat → write to Redis (real-time data for frontend)
  2. Alert detection: CPU/mem/disk > threshold → WebSocket push + Telegram
  3. Recovery detection: previously alerted metrics now normal → WebSocket + Telegram
  4. Background task: Redis→MySQL batch flush every 10min
"""

import json
import logging
import secrets
from datetime import datetime, timezone
from typing import Optional

from fastapi import APIRouter, Depends, Header, HTTPException, Request
from sqlalchemy.ext.asyncio import AsyncSession

from server.api.dependencies import get_db, get_server_service
from server.api.schemas import AgentHeartbeat, AgentScriptCallback
from server.api.websocket import broadcast_alert, broadcast_recovery
from server.application.services.server_service import ServerService
from server.domain.models import Server
from server.config import settings
from server.infrastructure.redis.client import get_redis
from server.infrastructure.telegram import send_telegram_system_alert

logger = logging.getLogger("nexus.agent")

router = APIRouter(prefix="/api/agent", tags=["agent"])

# Redis key patterns (TTL must exceed flush interval so keys survive until MySQL flush)
from server.background.heartbeat_flush import FLUSH_INTERVAL

REDIS_KEY_PREFIX = "heartbeat:"
REDIS_KEY_EXPIRE = int(FLUSH_INTERVAL * 1.5)  # 900s when flush is 600s
REDIS_ALERT_KEY_PREFIX = "alerts:"


def _verify_api_key(x_api_key: str = Header(...)):
    """Extract Agent API key header. Actual auth is per-server in the handler.

    server_id lives in the payload (not the header), so per-server verification
    runs inside the route via _verify_server_api_key(). Here we only ensure the
    header is present; an empty key is rejected early.
    """
    if not x_api_key:
        raise HTTPException(status_code=401, detail="Missing API key")
    return x_api_key


async def _verify_server_api_key(
    server_id: int, provided_key: str, service: ServerService
) -> Optional[Server]:
    """Load server and verify per-server agent_api_key (single DB read).

    Returns the Server on success; None if missing, misconfigured, wrong key, or DB error.
    No global API_KEY fallback (BL-06: per-server enforced).
    """
    try:
        server = await service.get_server(server_id)
    except Exception:
        logger.warning(
            "Failed to look up server %s for per-server API key check — rejecting",
            server_id,
            exc_info=True,
        )
        return None

    if not server:
        return None

    if not server.agent_api_key:
        logger.warning(
            "Server %s (%s) rejected: no agent_api_key set (global fallback removed). "
            "Generate via POST /api/servers/%s/agent-key",
            server_id,
            server.name,
            server_id,
        )
        return None

    if not secrets.compare_digest(provided_key, server.agent_api_key):
        return None

    return server


def _threshold_for(thresholds: dict, metric: str) -> float:
    defaults = {
        "cpu": settings.CPU_ALERT_THRESHOLD,
        "mem": settings.MEM_ALERT_THRESHOLD,
        "disk": settings.DISK_ALERT_THRESHOLD,
    }
    return float(thresholds.get(metric, defaults.get(metric, 80)))


def _safe_float(value, name: str) -> Optional[float]:
    """Convert agent-supplied metric value to float, discarding malformed values."""
    if value is None:
        return None
    try:
        return float(value)
    except (ValueError, TypeError):
        logger.warning(f"Agent sent non-numeric {name}: {value!r:.50} — skipped")
        return None


def _detect_alerts(system_info: dict, thresholds: dict) -> list:
    """Detect metrics exceeding alert thresholds

    Returns list of (alert_type, value) tuples.
    E.g. [("cpu", 85.3), ("mem", 92.1)]
    """
    alerts = []
    cpu = _safe_float(system_info.get("cpu_usage") or system_info.get("cpu"), "cpu")
    mem = _safe_float(system_info.get("mem_usage") or system_info.get("mem"), "mem")
    disk = _safe_float(system_info.get("disk_usage") or system_info.get("disk"), "disk")

    if cpu is not None and cpu > _threshold_for(thresholds, "cpu"):
        alerts.append(("cpu", cpu))
    if mem is not None and mem > _threshold_for(thresholds, "mem"):
        alerts.append(("mem", mem))
    if disk is not None and disk > _threshold_for(thresholds, "disk"):
        alerts.append(("disk", disk))

    return alerts


def _alert_metric_key(entry: str) -> str | None:
    """Parse Redis alert set member — supports legacy ``metric:value`` and ``metric``."""
    if not entry:
        return None
    metric = entry.split(":", 1)[0].strip()
    return metric or None


def _detect_recovery(system_info: dict, prev_alerts: set, thresholds: dict) -> list:
    """Detect previously alerted metrics that have returned to normal

    Returns list of (metric, current_value) tuples — at most one entry per metric
    even when Redis still holds legacy ``metric:value`` duplicates.
    """
    recoveries: list[tuple[str, float]] = []
    seen_metrics: set[str] = set()
    for prev in prev_alerts:
        metric = _alert_metric_key(prev)
        if not metric or metric in seen_metrics:
            continue
        current = _safe_float(
            system_info.get(f"{metric}_usage") or system_info.get(metric), metric
        )
        if current is not None and current < _threshold_for(thresholds, metric):
            seen_metrics.add(metric)
            recoveries.append((metric, current))

    return recoveries


@router.post("/heartbeat", response_model=dict)
async def receive_heartbeat(
    payload: AgentHeartbeat,
    api_key: str = Depends(_verify_api_key),
    service: ServerService = Depends(get_server_service),
):
    """Receive Agent heartbeat data

    Flow:
    1. Write heartbeat to Redis (frontend reads Redis for live status)
    2. Detect alerts (CPU/mem/disk > threshold) → WebSocket + Telegram push
    3. Detect recoveries (previously alerted metrics now normal) → push
    4. MySQL history via heartbeat_flush (10min), not per-request writes
    """
    server_id = payload.server_id

    # ── Unregistered / legacy agent: no server_id → silently acknowledge ──
    # Old MultiSync agents and misconfigured agents send heartbeats without
    # server_id.  Return 200 so they don't retry endlessly; log for visibility.
    if server_id is None:
        logger.warning(
            "Heartbeat discarded: missing server_id (legacy/unregistered agent). "
            "agent_version=%s system_info_keys=%s",
            payload.agent_version,
            list((payload.system_info or {}).keys()),
        )
        return {"status": "discarded", "reason": "missing server_id"}

    is_online = payload.is_online
    system_info = payload.system_info or {}
    agent_version = payload.agent_version or ""

    # ── Per-server API key verification (single get_server) ──
    server = await _verify_server_api_key(server_id, api_key, service)
    if not server:
        raise HTTPException(status_code=401, detail="Invalid API key for this server")

    # ── 1. Write to Redis (real-time data) ──

    # Time drift detection: compare agent_time with central server time
    TIME_DRIFT_WARN_SECONDS = 30    # ≥30s: warning (yellow)
    TIME_DRIFT_CRIT_SECONDS = 60    # ≥60s: critical (red, affects TOTP / cron)

    time_drift_seconds: float = 0.0
    drift_level: str = "ok"  # ok / warn / crit
    agent_time_str = system_info.get("agent_time", "") if system_info else ""
    if agent_time_str:
        try:
            agent_ts = datetime.fromisoformat(agent_time_str)
            central_ts = datetime.now(timezone.utc)
            # Normalise to aware for comparison
            if agent_ts.tzinfo is None:
                agent_ts = agent_ts.replace(tzinfo=timezone.utc)
            drift = abs((central_ts - agent_ts).total_seconds())
            time_drift_seconds = round(drift, 1)
            if drift >= TIME_DRIFT_CRIT_SECONDS:
                drift_level = "crit"
                logger.warning(
                    "Clock drift CRITICAL: server %s drift=%.1fs (agent_time=%s)",
                    server_id, drift, agent_time_str,
                )
            elif drift >= TIME_DRIFT_WARN_SECONDS:
                drift_level = "warn"
                logger.warning(
                    "Clock drift WARNING: server %s drift=%.1fs",
                    server_id, drift,
                )
        except Exception as e:
            logger.debug(f"Failed to parse agent_time for drift check: {e}")

    redis = get_redis()
    try:
        redis_key = f"{REDIS_KEY_PREFIX}{server_id}"
        await redis.hset(redis_key, mapping={
            "is_online": str(is_online),
            "system_info": json.dumps(system_info),
            "last_heartbeat": datetime.now(timezone.utc).isoformat(),
            "agent_version": agent_version,
            "time_drift_seconds": str(time_drift_seconds),
            "drift_level": drift_level,
        })
        # Set TTL = flush interval × 1.5 (prevent stale online status after key expires)
        await redis.expire(redis_key, REDIS_KEY_EXPIRE)
    except Exception as e:
        logger.error(f"Redis heartbeat write failed for server {server_id}: {e}")
        # Redis write failed — still continue with MySQL update

    # ── 1b. Server name for alert broadcasts ──
    server_name = server.name or f"server-{server_id}"

    # ── 1c. Time drift Telegram alert (critical only, 5-min cooldown like metric alerts) ──
    if drift_level == "crit":
        from server.api.websocket import _should_push_telegram
        drift_alert_key = f"alert:{server_id}:time_drift"
        if _should_push_telegram(drift_alert_key):
            try:
                await send_telegram_system_alert(
                    summary=f"子机与 Nexus 主站时钟偏差 {time_drift_seconds} 秒，已达严重级别",
                    title="子机时钟偏差告警",
                    notify_key="NOTIFY_TIME_DRIFT",
                    icon="🕐",
                    details=[
                        f"子机名称：{server_name}",
                        f"子机 ID：{server_id}",
                        "偏差等级：严重（≥ 60 秒）",
                        "检测来源：Agent 心跳上报（Layer 2）",
                        "可能影响：TOTP 双因素登录、定时任务、cron 与日志时间对齐",
                    ],
                    action="请在子机检查 NTP/chrony/systemd-timesyncd 是否正常同步",
                )
            except Exception:
                logger.debug("Failed to send drift alert via Telegram")

    # ── 2. Alert detection ──
    alert_thresholds = {
        "cpu": settings.CPU_ALERT_THRESHOLD,
        "mem": settings.MEM_ALERT_THRESHOLD,
        "disk": settings.DISK_ALERT_THRESHOLD,
    }

    if system_info:
        alerts = _detect_alerts(system_info, alert_thresholds)
        for alert_type, alert_value in alerts:
            logger.warning(f"Alert: server {server_id} {alert_type}={alert_value}%")
            await broadcast_alert(server_id, alert_type, alert_value, server_name)

            # Track active alert in Redis
            try:
                redis = get_redis()
                # Store metric name only — value in member caused duplicate recovery pushes
                await redis.sadd(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", alert_type)
                await redis.expire(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", 3600)  # 1h TTL
            except Exception:
                logger.debug(f"Failed to track alert in Redis for server {server_id}")

    # ── 3. Recovery detection ──
    if system_info:
        try:
            redis = get_redis()
            prev_alerts = await redis.smembers(f"{REDIS_ALERT_KEY_PREFIX}{server_id}")
            if prev_alerts:
                recoveries = _detect_recovery(system_info, prev_alerts, alert_thresholds)
                for metric, value in recoveries:
                    logger.info(f"Recovery: server {server_id} {metric}={value}%")
                    await broadcast_recovery(server_id, metric, value, server_name)

                    # Remove recovered metric (legacy metric:value members included)
                    for prev in prev_alerts:
                        if _alert_metric_key(prev) == metric:
                            await redis.srem(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", prev)
        except Exception as e:
            logger.error(f"Recovery detection failed for server {server_id}: {e}")

    # MySQL history is flushed every 10min by heartbeat_flush (Redis → MySQL).

    return {"status": "ok", "server_id": server_id}

@router.post("/script-callback", response_model=dict)
async def script_job_callback(
    payload: AgentScriptCallback,
    request: Request,
    api_key: str = Depends(_verify_api_key),
    service: ServerService = Depends(get_server_service),
    db: AsyncSession = Depends(get_db),
):
    """Receive long-task completion from child host (curl after nohup script exits).

    Authenticated by X-API-Key (per-server agent_api_key) + one-time job_id + secret.
    """
    if not await _verify_server_api_key(payload.server_id, api_key, service):
        raise HTTPException(status_code=401, detail="Invalid API key for this server")
    from server.application.services.script_job_callback import apply_script_job_callback
    from server.api.websocket import broadcast_system_event
    from server.infrastructure.database.script_repo import ScriptExecutionRepositoryImpl
    from server.infrastructure.redis.script_callback_rate import check_script_callback_rate

    client_ip = request.client.host if request.client else ""
    await check_script_callback_rate(payload.job_id, client_ip)

    repo = ScriptExecutionRepositoryImpl(db)
    result = await apply_script_job_callback(
        repo,
        job_id=payload.job_id,
        server_id=payload.server_id,
        secret=payload.secret,
        exit_code=payload.exit_code,
        message=payload.message,
        log_tail=payload.log_tail,
        operator="agent",
    )

    if not result:
        raise HTTPException(status_code=404, detail="Invalid or expired job")

    from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
    from server.domain.models import AuditLog

    audit_repo = AuditLogRepositoryImpl(db)
    await audit_repo.create(AuditLog(
        admin_username="agent",
        action="script_job_callback",
        target_type="script_execution",
        target_id=result["execution_id"],
        detail=(
            f"服务器ID={payload.server_id} 任务ID={payload.job_id} "
            f"退出码={payload.exit_code} 状态={result.get('status')}"
        ),
    ))

    await broadcast_system_event(
        "script_job_done",
        f"脚本执行 #{result['execution_id']} 服务器 {result['server_id']} "
        f"{'成功' if result['exit_code'] == 0 else '失败'}",
    )
    logger.info(
        "Script job callback: execution_id=%s server_id=%s exit_code=%s",
        result["execution_id"],
        result["server_id"],
        result["exit_code"],
    )
    return {"status": "ok", **result}


# Remote command execution lives on each managed host: web/agent/agent.py POST /exec
# Do NOT expose shell exec on the Nexus control plane (was P0 RCE via global API key).