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.config import settings
from server.infrastructure.redis.client import get_redis, get_redis_sync
from server.infrastructure.telegram import send_telegram_system_alert

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

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

# Redis key patterns
REDIS_KEY_PREFIX = "heartbeat:"
REDIS_KEY_EXPIRE = 600  # 10 min TTL — matches flush interval
REDIS_ALERT_KEY_PREFIX = "alerts:"


def _verify_api_key(x_api_key: str = Header(...)):
    """Verify Agent API key — accepts either global or per-server key.

    For heartbeat endpoints, the global key is checked first; if it doesn't
    match, the request is still allowed through so the handler can verify
    the per-server key (which requires server_id from the payload).

    """
    if not x_api_key:
        raise HTTPException(status_code=401, detail="Missing API key")
    # If the global key matches, allow immediately
    if secrets.compare_digest(x_api_key, settings.API_KEY):
        return x_api_key
    # Non-matching key — allow through for per-server verification in handler.
    # The handler MUST call _verify_server_api_key() before processing.
    # This two-step approach is needed because server_id is in the payload,
    # not the header, so per-server verification can't happen in a Depends().
    return x_api_key


async def _verify_server_api_key(server_id: int, provided_key: str, service: ServerService) -> bool:
    """Verify API key against per-server agent_api_key, falling back to global API key.

    Authentication priority:
    1. If server has agent_api_key set → must match exactly
    2. Otherwise → must match global API_KEY
    """
    # Try to get server's per-server key
    try:
        server = await service.get_server(server_id)
        if server and server.agent_api_key:
            return secrets.compare_digest(provided_key, server.agent_api_key)
    except Exception:
        logger.debug(f"Failed to look up server {server_id} for per-server API key check, falling back to global key")

    # Fallback: global API key
    return secrets.compare_digest(provided_key, settings.API_KEY)


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 _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 = system_info.get("cpu_usage") or system_info.get("cpu")
    mem = system_info.get("mem_usage") or system_info.get("mem")
    disk = system_info.get("disk_usage") or system_info.get("disk")

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

    return alerts


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.
    """
    recoveries = []
    for prev in prev_alerts:
        parts = prev.split(":")
        if len(parts) != 2:
            continue
        metric, prev_val = parts[0], float(parts[1])
        current = system_info.get(f"{metric}_usage") or system_info.get(metric)

        if current is not None and float(current) < _threshold_for(thresholds, metric):
            recoveries.append((metric, float(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. Also update MySQL via ServerService (for persistent record)
    """
    server_id = payload.server_id
    is_online = payload.is_online
    system_info = payload.system_info or {}
    agent_version = payload.agent_version or ""

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

    # ── 1. Write to Redis (real-time data) ──
    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,
        })
        # 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. Look up server name for alert broadcasts ──
    server_name = f"server-{server_id}"
    try:
        server = await service.get_server(server_id)
        if server:
            server_name = server.name
    except Exception:
        logger.debug(f"Failed to look up server name for server {server_id}")

    # ── 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()
                await redis.sadd(f"{REDIS_ALERT_KEY_PREFIX}{server_id}", f"{alert_type}:{alert_value}")
                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 alert from Redis tracking
                    for prev in prev_alerts:
                        if prev.startswith(f"{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}")

    # ── 4. Also update MySQL (persistent record for non-real-time queries) ──
    await service.update_heartbeat(server_id, payload.model_dump())

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

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

    Authenticated by one-time job_id + secret stored in Redis at exec time.
    """
    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"server_id={payload.server_id} job_id={payload.job_id} "
            f"exit_code={payload.exit_code} status={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).