feat: MeshSafe 文件摆渡安全卫士 v1.0 初始提交
核心功能: - 6 层检测管线: L0 结构分析 → L1 ClamAV → L2 文档深度检测 → L3 关键词匹配 → L4 YARA → L5 PE/二进制 - FastAPI + PostgreSQL + Redis + Docker Compose 架构 - 客户端 HMAC-SHA256 签名鉴权 + 限流 + 幂等性 - 控制台 UI (登录/Dashboard/诊断中心/策略管理/任务列表/测试中心/产品介绍页) - 三层诊断定时器 (error_scan 30s / db_aggregate 30s / container_check 5min) 并发与稳定性修复: - 诊断定时器 Redis 时间槽位锁 (多 worker 同一槽位只有一个执行) - JSONL 存储 fcntl 跨进程排他锁 + 容错解码 + 原子重写 - 移除 SQLite fallback, 统一 PostgreSQL - API_CLIENTS 脱敏为环境变量, .env 已 .gitignore
This commit is contained in:
268
app/diagnostic/scheduler.py
Normal file
268
app/diagnostic/scheduler.py
Normal file
@@ -0,0 +1,268 @@
|
||||
"""三层定时器循环。
|
||||
|
||||
在 FastAPI startup 事件中用 asyncio.create_task() 启动。
|
||||
定时器在后台跑，每次迭代独立 try/except + timeout，不会卡死主 event loop。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.config import get_settings
|
||||
from app.db import SessionLocal
|
||||
from app.diagnostic.config import DiagConfig, get_diag_config
|
||||
from app.diagnostic.engine import DiagnosticEngine, RuleResult
|
||||
from app.services.scanner import Scanner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# 诊断套件全局运行时状态（单例 dict）
|
||||
_DIAG_STATE: Dict[str, Any] = {
|
||||
"enabled": False,
|
||||
"docker_status": "none",
|
||||
"docker_info": {},
|
||||
"started_at": None,
|
||||
"rules_cache": {}, # {rule_id: {result, ts, run_by}}
|
||||
"latest": {"rules": [], "generated_at": None, "docker_status": "none"},
|
||||
}
|
||||
|
||||
# 后台引用（防止 GC 收走）
|
||||
_TIMER_TASKS: List[asyncio.Task] = []
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 外部访问入口
|
||||
# ======================================================================
|
||||
|
||||
def get_diag_state() -> Dict[str, Any]:
|
||||
"""获取诊断套件运行时状态（给页面/API 查询用）。"""
|
||||
cache = {k: {"result": v["result"].to_dict() if isinstance(v.get("result"), RuleResult) else v.get("result"),
|
||||
"ts": v.get("ts"), "run_by": v.get("run_by")}
|
||||
for k, v in _DIAG_STATE["rules_cache"].items()}
|
||||
return {
|
||||
"enabled": _DIAG_STATE["enabled"],
|
||||
"docker_status": _DIAG_STATE["docker_status"],
|
||||
"docker_info": _DIAG_STATE["docker_info"],
|
||||
"started_at": _DIAG_STATE["started_at"],
|
||||
"rules_cache": cache,
|
||||
"latest": _DIAG_STATE["latest"],
|
||||
}
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 后台入口（在 app.main.startup 中调用）
|
||||
# ======================================================================
|
||||
|
||||
def start_timers(docker_status: str, docker_info: Dict[str, Any]) -> None:
|
||||
"""启动三个定时器循环（后台 daemon task）。"""
|
||||
cfg = get_diag_config()
|
||||
_DIAG_STATE["enabled"] = cfg.enabled
|
||||
_DIAG_STATE["docker_status"] = docker_status
|
||||
_DIAG_STATE["docker_info"] = docker_info
|
||||
_DIAG_STATE["started_at"] = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
if not cfg.enabled:
|
||||
logger.info("MeshSafe Diagnostics: disabled")
|
||||
return
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
if cfg.error_scan_interval > 0:
|
||||
_TIMER_TASKS.append(loop.create_task(
|
||||
_timer_loop("error_scan", cfg.error_scan_interval,
|
||||
lambda eng: eng.run_timer_error_scan(),
|
||||
docker_status, cfg.docker_log_tail),
|
||||
name="diag-error-scan",
|
||||
))
|
||||
logger.info("诊断定时器启动: error_scan (%ss)", cfg.error_scan_interval)
|
||||
|
||||
if cfg.db_aggregate_interval > 0:
|
||||
_TIMER_TASKS.append(loop.create_task(
|
||||
_timer_loop("db_aggregate", cfg.db_aggregate_interval,
|
||||
lambda eng: eng.run_timer_db_aggregate(),
|
||||
docker_status, cfg.docker_log_tail),
|
||||
name="diag-db-aggregate",
|
||||
))
|
||||
logger.info("诊断定时器启动: db_aggregate (%ss)", cfg.db_aggregate_interval)
|
||||
|
||||
if cfg.container_check_interval > 0:
|
||||
_TIMER_TASKS.append(loop.create_task(
|
||||
_timer_loop("container_check", cfg.container_check_interval,
|
||||
lambda eng: eng.run_timer_container(),
|
||||
docker_status, cfg.docker_log_tail),
|
||||
name="diag-container-check",
|
||||
))
|
||||
logger.info("诊断定时器启动: container_check (%ss)", cfg.container_check_interval)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 定时器循环（async）
|
||||
# ======================================================================
|
||||
|
||||
async def _timer_loop(
|
||||
name: str,
|
||||
interval: int,
|
||||
runner_fn, # 接受 DiagnosticEngine 返回 List[RuleResult]
|
||||
docker_status: str,
|
||||
docker_log_tail: int,
|
||||
) -> None:
|
||||
"""通用定时器循环。每次迭代独立 try/except + 超时保护。"""
|
||||
|
||||
# 等一个 interval 再开始第一次（避免 startup 时 DB 还没 ready）
|
||||
await asyncio.sleep(min(interval, 15))
|
||||
|
||||
while True:
|
||||
try:
|
||||
# 分布式锁（时间槽位）：多 worker 部署时，每个时间槽只有
|
||||
# 一个 worker 能拿到锁并执行，保证执行频率精确等于 interval。
|
||||
if not _acquire_timer_lock(name, interval):
|
||||
logger.debug("诊断定时器 %s 本轮跳过（其他 worker 持锁）", name)
|
||||
await asyncio.sleep(interval)
|
||||
continue
|
||||
|
||||
results = await asyncio.wait_for(
|
||||
asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
lambda: _run_sync(name, runner_fn, docker_status, docker_log_tail),
|
||||
),
|
||||
timeout=30.0,
|
||||
)
|
||||
_cache_results(name, results, docker_status)
|
||||
logger.debug("诊断定时器 %s 跑完: %d 条规则", name, len(results))
|
||||
except asyncio.CancelledError:
|
||||
logger.info("诊断定时器 %s 被取消", name)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("诊断定时器 %s 超时（30s），本轮跳过", name)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.exception("诊断定时器 %s 异常: %s，%ss 后重试", name, e, interval)
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 定时器分布式锁（Redis）
|
||||
#
|
||||
# 多 worker 部署（UVICORN_WORKERS > 1）时，每个 worker 都会启动定时器，
|
||||
# 导致：同一轮规则被跑 N 次、历史数据量 ×N、并发写 JSONL 造成文件损坏。
|
||||
# 用 Redis SET NX 抢锁，保证同一时刻只有一个 worker 在跑。
|
||||
# ======================================================================
|
||||
|
||||
_TIMER_LOCK_PREFIX = "meshsafe:diag:timer_lock:"
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _lock_redis():
|
||||
"""定时器锁专用 Redis 客户端（进程内复用，避免每次调用重建连接）。"""
|
||||
import redis # type: ignore
|
||||
|
||||
from app.config import get_settings
|
||||
|
||||
return redis.from_url(
|
||||
get_settings().redis_url, decode_responses=True, socket_timeout=3
|
||||
)
|
||||
|
||||
|
||||
def _acquire_timer_lock(timer_name: str, interval: int) -> bool:
|
||||
"""尝试获取定时器锁（时间槽位设计）。
|
||||
|
||||
key 里带上当前时间槽位号：``meshsafe:diag:timer_lock:<timer>:<slot>``，
|
||||
slot = 当前时间 // interval。这样无论多少 worker、谁先谁后，
|
||||
同一个槽位内最多只有一个 worker 能拿到锁并执行，
|
||||
下一个槽位自然放行 —— 执行频率精确等于 interval，不多不少。
|
||||
|
||||
Redis 不可用时返回 True（降级为单 worker 行为，不阻塞主流程）。
|
||||
"""
|
||||
try:
|
||||
client = _lock_redis()
|
||||
slot = int(time.time() // interval)
|
||||
key = f"{_TIMER_LOCK_PREFIX}{timer_name}:{slot}"
|
||||
# EX 只做兜底（槽位 key 不会互相干扰），设 2 倍 interval 足够
|
||||
return bool(client.set(key, "1", nx=True, ex=interval * 2))
|
||||
except Exception as e: # noqa: BLE001
|
||||
# Redis 不可用 → 不加锁（单机单 worker 场景行为不变）
|
||||
logger.debug("定时器锁获取失败，降级为无锁执行: %s", e)
|
||||
return True
|
||||
|
||||
|
||||
def _run_sync(
|
||||
name: str,
|
||||
runner_fn,
|
||||
docker_status: str,
|
||||
docker_log_tail: int,
|
||||
) -> List[RuleResult]:
|
||||
"""在线程池里同步跑规则（避免阻塞 event loop）。"""
|
||||
db = SessionLocal()
|
||||
try:
|
||||
# 延迟 import 避免循环依赖
|
||||
from app.main import scanner, policy_store
|
||||
eng = DiagnosticEngine(
|
||||
db=db, scanner=scanner, policy_store=policy_store,
|
||||
docker_status=docker_status, docker_log_tail=docker_log_tail,
|
||||
)
|
||||
results = runner_fn(eng)
|
||||
return results
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _cache_results(
|
||||
timer_name: str,
|
||||
results: List[RuleResult],
|
||||
docker_status: str,
|
||||
) -> None:
|
||||
"""把本轮规则结果写入内存缓存（dict）+ 持久化存储。"""
|
||||
now = time.time()
|
||||
for r in results:
|
||||
_DIAG_STATE["rules_cache"][r.id] = {
|
||||
"result": r,
|
||||
"ts": now,
|
||||
"run_by": r.run_by or f"timer:{timer_name}",
|
||||
}
|
||||
|
||||
# 构建 latest 快照（合并所有缓存 + 本轮结果）
|
||||
all_results = []
|
||||
seen = set()
|
||||
for rid, entry in _DIAG_STATE["rules_cache"].items():
|
||||
if rid not in seen and entry.get("result"):
|
||||
all_results.append(entry["result"])
|
||||
seen.add(rid)
|
||||
|
||||
_DIAG_STATE["latest"] = {
|
||||
"rules": sorted(all_results, key=lambda x: x.id),
|
||||
"generated_at": now,
|
||||
"docker_status": docker_status,
|
||||
}
|
||||
|
||||
# ---------- 持久化 ----------
|
||||
try:
|
||||
from app.diagnostic.config import get_diag_config
|
||||
from app.diagnostic.storage import append_rule_results
|
||||
|
||||
cfg = get_diag_config()
|
||||
if cfg.storage_enabled:
|
||||
append_rule_results(results, timer_name)
|
||||
|
||||
# 可选：定时器也存 docker 日志快照（省磁盘所以用较小 tail）
|
||||
if cfg.log_snapshot_on_timer and docker_status != "none":
|
||||
from app.diagnostic.storage import snapshot_container_logs
|
||||
containers = [
|
||||
"meshsafe-api", "meshsafe-clamav", "meshsafe-postgres",
|
||||
"meshsafe-redis", "meshsafe-nginx",
|
||||
]
|
||||
# 不要每次都存 — 只在 container_check（5 分钟周期）里存一次
|
||||
if timer_name == "container_check":
|
||||
snapshot_container_logs(
|
||||
containers, cfg.log_snapshot_tail,
|
||||
filename_prefix="timer",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.debug("诊断持久化跳过（非致命）: %s", e)
|
||||
Reference in New Issue
Block a user