95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
"""Agent 上报接收路由（数据写入 metrics.db）"""
|
||
|
||
import threading
|
||
import time
|
||
|
||
from fastapi import APIRouter, Depends, HTTPException, status
|
||
from sqlmodel import Session, select
|
||
|
||
from app.core.config import settings
|
||
from app.core.security import require_agent_key
|
||
from app.core.timeutils import utcnow
|
||
from app.database import assets_engine, get_metrics_session
|
||
from app.models.asset import Asset
|
||
from app.models.monitor import EventLog, MetricPoint, SecurityCheck, ServerInfo
|
||
from app.schemas.agent import AgentReport
|
||
|
||
router = APIRouter(prefix="/api/agent", tags=["agent"])
|
||
|
||
# 内存级频率限制：{asset_id: 上次上报的 monotonic 时间戳}
|
||
# 同步路由运行在线程池，多线程并发读写需加锁保证“读-判断-写”原子性
|
||
_last_report_ts: dict = {}
|
||
_rate_limit_lock = threading.Lock()
|
||
|
||
|
||
def _check_rate_limit(asset_id: int) -> None:
|
||
"""限制同一资产的上报频率，防止配置错误的 Agent 高频写入填满数据库"""
|
||
interval = settings.AGENT_REPORT_MIN_INTERVAL
|
||
if interval <= 0:
|
||
return
|
||
now = time.monotonic()
|
||
with _rate_limit_lock:
|
||
last = _last_report_ts.get(asset_id)
|
||
if last is not None and (now - last) < interval:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||
detail=f"上报过于频繁，同一资产最小间隔 {interval} 秒",
|
||
)
|
||
_last_report_ts[asset_id] = now
|
||
|
||
|
||
def _validate_asset_id(asset_id: int) -> None:
|
||
"""校验上报的 asset_id 在资产库中真实存在，防止脏数据写入"""
|
||
with Session(assets_engine) as session:
|
||
asset = session.get(Asset, asset_id)
|
||
if not asset:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail=f"asset_id={asset_id} 不存在，请先在资产列表中登记该服务器",
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/report",
|
||
summary="Agent 上报资源/信息/安全状态",
|
||
dependencies=[Depends(require_agent_key)],
|
||
)
|
||
def report(data: AgentReport, session: Session = Depends(get_metrics_session)) -> dict:
|
||
_validate_asset_id(data.asset_id)
|
||
_check_rate_limit(data.asset_id)
|
||
|
||
# 资源监控时序
|
||
if data.metrics:
|
||
session.add(MetricPoint(asset_id=data.asset_id, **data.metrics.model_dump()))
|
||
|
||
# 服务器信息快照（每资产一条，覆盖更新）
|
||
if data.server_info:
|
||
info = data.server_info.model_dump()
|
||
info["last_seen"] = utcnow()
|
||
existing = session.exec(
|
||
select(ServerInfo).where(ServerInfo.asset_id == data.asset_id)
|
||
).first()
|
||
if existing:
|
||
for key, value in info.items():
|
||
setattr(existing, key, value)
|
||
session.add(existing)
|
||
else:
|
||
session.add(ServerInfo(asset_id=data.asset_id, **info))
|
||
|
||
# 安全检查项
|
||
if data.security:
|
||
for check in data.security:
|
||
session.add(SecurityCheck(asset_id=data.asset_id, **check.model_dump()))
|
||
|
||
# 事件日志
|
||
session.add(
|
||
EventLog(
|
||
level="info",
|
||
source="agent",
|
||
asset_id=data.asset_id,
|
||
message=f"Agent 上报 asset_id={data.asset_id}",
|
||
)
|
||
)
|
||
session.commit()
|
||
return {"status": "ok", "asset_id": data.asset_id}
|