Nexus Deploy
b866e944ab
Apply sync/install/auth/schedule/retry/agent/settings fixes from full code review; document accepted WS and Agent legacy risks for solo ops. Co-authored-by: Cursor <cursoragent@cursor.com>
379 lines
15 KiB
Python
379 lines
15 KiB
Python
"""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
|
||
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(...)):
|
||
"""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, with limited global fallback.
|
||
|
||
Authentication priority:
|
||
1. Server must exist
|
||
2. If server has agent_api_key → must match exactly
|
||
3. Legacy only: no per-server key → global API_KEY (logged for migration)
|
||
"""
|
||
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 False
|
||
|
||
if not server:
|
||
return False
|
||
|
||
if server.agent_api_key:
|
||
return secrets.compare_digest(provided_key, server.agent_api_key)
|
||
|
||
if secrets.compare_digest(provided_key, settings.API_KEY):
|
||
logger.warning(
|
||
"Server %s (%s) used global API_KEY for agent auth — "
|
||
"assign per-server key via POST /api/servers/%s/agent-key",
|
||
server_id,
|
||
server.name,
|
||
server_id,
|
||
)
|
||
return True
|
||
return False
|
||
|
||
|
||
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 _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 = parts[0]
|
||
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):
|
||
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 ──
|
||
if not await _verify_server_api_key(server_id, api_key, service):
|
||
raise HTTPException(status_code=401, detail="Invalid API key for this server")
|
||
|
||
server = await service.get_server(server_id)
|
||
if not server:
|
||
logger.warning(
|
||
"Heartbeat discarded: unknown server_id=%s (unregistered host)",
|
||
server_id,
|
||
)
|
||
return {"status": "discarded", "reason": "unknown_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(
|
||
f"🕐 时钟偏差过大 — 服务器 #{server_id} ({server_name})\n"
|
||
f"偏差 {time_drift_seconds}s（≥60s 影响 TOTP / cron 精度）\n"
|
||
f"请检查 NTP 配置",
|
||
notify_key="NOTIFY_TIME_DRIFT",
|
||
)
|
||
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()
|
||
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}")
|
||
|
||
# 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 or global) + 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). |