zhaolei
16f134e328
核心功能: - 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
456 lines
14 KiB
Python
456 lines
14 KiB
Python
"""诊断套件持久化存储模块。
|
||
|
||
职责：
|
||
- 定时器规则结果 → JSON Lines 追加（rules_history.jsonl）
|
||
- Docker 日志快照 → 存文件（logs/*.log）
|
||
- 自动清理过期数据（基于配置的 history_retention_days）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from app.diagnostic.config import get_diag_config
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _storage_dir() -> Path:
|
||
"""获取存储根目录。"""
|
||
cfg = get_diag_config()
|
||
p = Path(cfg.storage_dir)
|
||
p.mkdir(parents=True, exist_ok=True)
|
||
return p
|
||
|
||
|
||
def _logs_dir() -> Path:
|
||
p = _storage_dir() / "logs"
|
||
p.mkdir(parents=True, exist_ok=True)
|
||
return p
|
||
|
||
|
||
def _rules_history_file() -> Path:
|
||
return _storage_dir() / "rules_history.jsonl"
|
||
|
||
|
||
# ======================================================================
|
||
# 跨进程文件锁写入
|
||
#
|
||
# 多 worker 部署时，每个 worker 都会跑诊断定时器。"a" 追加模式在多进程
|
||
# 下不安全——各进程独立维护文件偏移量，后写的会覆盖先写的内容，把一行
|
||
# JSON 从中间劈开，产出非法 UTF-8 字节，导致后续读取全部报错。
|
||
# 这里用 fcntl 文件锁保证同一时刻只有一个进程在写。
|
||
# ======================================================================
|
||
|
||
try:
|
||
import fcntl # type: ignore
|
||
|
||
_HAS_FCNTL = True
|
||
except ImportError: # pragma: no cover - Windows 平台
|
||
fcntl = None # type: ignore
|
||
_HAS_FCNTL = False
|
||
|
||
|
||
def _append_line_locked(path: Path, line: str) -> None:
|
||
"""向文件追加一行，跨进程互斥。
|
||
|
||
Linux 下用 fcntl.flock 排他锁；无 fcntl 时退化为进程内线程锁。
|
||
"""
|
||
payload = line + "\n"
|
||
with open(path, "a", encoding="utf-8") as fp:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(fp.fileno(), fcntl.LOCK_EX)
|
||
try:
|
||
fp.seek(0, os.SEEK_END)
|
||
fp.write(payload)
|
||
fp.flush()
|
||
os.fsync(fp.fileno())
|
||
finally:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(fp.fileno(), fcntl.LOCK_UN)
|
||
|
||
|
||
def _rewrite_file_locked(path: Path, lines: List[str]) -> None:
|
||
"""整文件重写（清理过期数据用），同样走文件锁。
|
||
|
||
写临时文件再原子替换，避免重写过程中其他进程读到半截内容。
|
||
"""
|
||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||
with open(tmp, "w", encoding="utf-8") as fp:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(fp.fileno(), fcntl.LOCK_EX)
|
||
try:
|
||
for line in lines:
|
||
fp.write(line + "\n")
|
||
fp.flush()
|
||
os.fsync(fp.fileno())
|
||
finally:
|
||
if _HAS_FCNTL:
|
||
fcntl.flock(fp.fileno(), fcntl.LOCK_UN)
|
||
os.replace(tmp, path)
|
||
|
||
|
||
# ======================================================================
|
||
# 规则历史：JSON Lines 追加
|
||
# ======================================================================
|
||
|
||
def append_rule_results(
|
||
rules: List[Any],
|
||
timer_name: str,
|
||
) -> None:
|
||
"""定时器每次跑完追加一行 JSONL。
|
||
|
||
每行格式::
|
||
|
||
{"ts": "2026-10-28T13:12:14", "timer": "error_scan",
|
||
"docker_status": "full",
|
||
"rules": [{"id":"R001","severity":"ok","message":"...","run_at":"...","run_by":"timer:error_scan",...}]}
|
||
"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return
|
||
|
||
try:
|
||
# 北京时间
|
||
bj = timezone(timedelta(hours=8))
|
||
entry: Dict[str, Any] = {
|
||
"ts": datetime.now(bj).strftime("%Y-%m-%d %H:%M:%S"),
|
||
"timer": timer_name,
|
||
"docker_status": None,
|
||
"rules": [
|
||
r.to_dict() if hasattr(r, "to_dict") else dict(r)
|
||
for r in rules
|
||
],
|
||
}
|
||
|
||
f = _rules_history_file()
|
||
_append_line_locked(f, json.dumps(entry, ensure_ascii=False))
|
||
|
||
# 顺带清理过期
|
||
_trim_history()
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("诊断持久化写入失败: %s", e)
|
||
|
||
|
||
def load_rule_history(days: Optional[int] = None) -> List[Dict[str, Any]]:
|
||
"""读取规则历史 JSONL，返回最近 N 天的行。"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return []
|
||
days = days or cfg.history_retention_days
|
||
|
||
f = _rules_history_file()
|
||
if not f.exists():
|
||
return []
|
||
|
||
cutoff = datetime.now() - timedelta(days=days)
|
||
results: List[Dict[str, Any]] = []
|
||
try:
|
||
with open(f, "r", encoding="utf-8", errors="replace") as fp:
|
||
for line in fp:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
# 按 ts 判断是否过期
|
||
ts_str = obj.get("ts", "")
|
||
try:
|
||
ts_dt = datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")
|
||
if ts_dt >= cutoff:
|
||
results.append(obj)
|
||
except ValueError:
|
||
# 时间格式不对就不过滤，保留
|
||
results.append(obj)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("诊断历史读取失败: %s", e)
|
||
|
||
return results
|
||
|
||
|
||
def _trim_history() -> None:
|
||
"""清理过期的规则历史行（保留文件头）。"""
|
||
cfg = get_diag_config()
|
||
days = cfg.history_retention_days
|
||
|
||
f = _rules_history_file()
|
||
if not f.exists():
|
||
return
|
||
|
||
cutoff = datetime.now() - timedelta(days=days)
|
||
kept: List[str] = []
|
||
try:
|
||
with open(f, "r", encoding="utf-8", errors="replace") as fp:
|
||
total = 0
|
||
for line in fp:
|
||
total += 1
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
ts_str = obj.get("ts", "")
|
||
ts_dt = datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")
|
||
if ts_dt >= cutoff:
|
||
kept.append(line)
|
||
except Exception:
|
||
kept.append(line) # 格式异常的保留
|
||
|
||
# 如果删了一些（不是全部），就重写
|
||
if kept and len(kept) < total:
|
||
_rewrite_file_locked(f, kept)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("诊断历史清理失败: %s", e)
|
||
|
||
|
||
# ======================================================================
|
||
# Docker 日志快照
|
||
# ======================================================================
|
||
|
||
def snapshot_container_logs(
|
||
containers: List[str],
|
||
tail: int,
|
||
filename_prefix: Optional[str] = None,
|
||
) -> List[Path]:
|
||
"""把各容器当前 docker logs --tail N 存一份文件。
|
||
|
||
文件名: {container}-{prefix}{timestamp}.log
|
||
prefix 可选（如 timer 存的加前缀 timer-）
|
||
返回写入的文件列表。
|
||
"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return []
|
||
|
||
from app.diagnostic.docker_probe import get_client
|
||
|
||
client = get_client()
|
||
if client is None:
|
||
return []
|
||
|
||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
prefix = f"{filename_prefix}-" if filename_prefix else ""
|
||
out: List[Path] = []
|
||
|
||
for name in containers:
|
||
try:
|
||
c = client.containers.get(name)
|
||
raw = c.logs(tail=tail)
|
||
text = (
|
||
raw.decode("utf-8", errors="replace")
|
||
if isinstance(raw, (bytes, bytearray))
|
||
else str(raw)
|
||
)
|
||
fname = f"{name}-{prefix}{ts}.log"
|
||
fpath = _logs_dir() / fname
|
||
fpath.write_text(text, encoding="utf-8")
|
||
out.append(fpath)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("容器 %s 日志快照失败: %s", name, e)
|
||
|
||
# 顺带清理过期日志
|
||
_trim_logs()
|
||
return out
|
||
|
||
|
||
def load_log_snapshots(containers: Optional[List[str]] = None) -> Dict[str, List[Path]]:
|
||
"""读 logs/ 目录，返回 {container_name: [path, path, ...]}。"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return {}
|
||
|
||
cutoff = datetime.now() - timedelta(days=cfg.history_retention_days)
|
||
result: Dict[str, List[Path]] = {}
|
||
for f in _logs_dir().glob("*.log"):
|
||
stem = f.stem # 去掉 .log
|
||
# 文件名格式: meshsafe-api-timer-20261028-1312 或 meshsafe-api-20261028-1312
|
||
# 容器名是 stem 里第一个 - 前面的部分... 其实容器名本身也可能带 -
|
||
# 简化：直接用 stem 整体当 key（导出时按需匹配）
|
||
# 但我们需要按容器分组，所以用 glob 前缀匹配
|
||
if containers:
|
||
matched = None
|
||
for c in containers:
|
||
if stem.startswith(c + "-"):
|
||
matched = c
|
||
break
|
||
if matched is None:
|
||
continue
|
||
container_key = matched
|
||
else:
|
||
# 没指定容器，尝试从文件名推断
|
||
parts = stem.split("-")
|
||
container_key = parts[0] if parts else stem
|
||
|
||
try:
|
||
mtime = datetime.fromtimestamp(f.stat().st_mtime)
|
||
if mtime >= cutoff:
|
||
result.setdefault(container_key, []).append(f)
|
||
except Exception:
|
||
result.setdefault(container_key, []).append(f)
|
||
|
||
# 按 mtime 排序（新的在前）
|
||
for k in result:
|
||
result[k].sort(key=lambda p: p.stat().st_mtime, reverse=True)
|
||
|
||
return result
|
||
|
||
|
||
def _trim_logs() -> None:
|
||
"""清理过期的日志快照文件。"""
|
||
cfg = get_diag_config()
|
||
cutoff = datetime.now() - timedelta(days=cfg.history_retention_days)
|
||
try:
|
||
for f in _logs_dir().glob("*.log"):
|
||
try:
|
||
mtime = datetime.fromtimestamp(f.stat().st_mtime)
|
||
if mtime < cutoff:
|
||
f.unlink()
|
||
except FileNotFoundError:
|
||
pass
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("日志快照清理失败: %s", e)
|
||
|
||
|
||
# ======================================================================
|
||
# 容器重启事件追踪
|
||
# ======================================================================
|
||
|
||
def _restart_events_file() -> Path:
|
||
return _storage_dir() / "restart_events.jsonl"
|
||
|
||
|
||
def append_restart_event(
|
||
container: str,
|
||
restart_before: int,
|
||
restart_after: int,
|
||
uptime_before: str = "",
|
||
uptime_before_seconds: float = 0,
|
||
trigger: str = "restart_count_increased",
|
||
) -> None:
|
||
"""记录一条容器重启事件到 restart_events.jsonl。"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return
|
||
try:
|
||
bj = timezone(timedelta(hours=8))
|
||
entry = {
|
||
"ts": datetime.now(bj).strftime("%Y-%m-%d %H:%M:%S"),
|
||
"container": container,
|
||
"restart_before": restart_before,
|
||
"restart_after": restart_after,
|
||
"uptime_before": uptime_before,
|
||
"uptime_before_seconds": int(uptime_before_seconds),
|
||
"trigger": trigger,
|
||
}
|
||
f = _restart_events_file()
|
||
_append_line_locked(f, json.dumps(entry, ensure_ascii=False))
|
||
|
||
# 清理过期事件
|
||
_trim_restart_events()
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("重启事件写入失败: %s", e)
|
||
|
||
|
||
def snapshot_on_restart(container: str, tail: int = 300) -> Optional[Path]:
|
||
"""容器刚重启后，抓一份当前 docker logs 片段存到 logs/*-restart-*.log。"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return None
|
||
try:
|
||
from app.diagnostic.docker_probe import get_client
|
||
client = get_client()
|
||
if client is None:
|
||
return None
|
||
c = client.containers.get(container)
|
||
raw = c.logs(tail=tail)
|
||
text = (
|
||
raw.decode("utf-8", errors="replace")
|
||
if isinstance(raw, (bytes, bytearray))
|
||
else str(raw)
|
||
)
|
||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
fname = f"{container}-restart-{ts}.log"
|
||
fpath = _logs_dir() / fname
|
||
fpath.write_text(text, encoding="utf-8")
|
||
return fpath
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("重启日志快照失败 %s: %s", container, e)
|
||
return None
|
||
|
||
|
||
def load_restart_events(days: Optional[int] = None) -> List[Dict[str, Any]]:
|
||
"""读取重启事件 JSONL，返回最近 N 天的行。"""
|
||
cfg = get_diag_config()
|
||
if not cfg.storage_enabled:
|
||
return []
|
||
days = days or cfg.history_retention_days
|
||
|
||
f = _restart_events_file()
|
||
if not f.exists():
|
||
return []
|
||
|
||
cutoff = datetime.now() - timedelta(days=days)
|
||
results: List[Dict[str, Any]] = []
|
||
try:
|
||
with open(f, "r", encoding="utf-8", errors="replace") as fp:
|
||
for line in fp:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
ts_str = obj.get("ts", "")
|
||
try:
|
||
ts_dt = datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")
|
||
if ts_dt >= cutoff:
|
||
results.append(obj)
|
||
except ValueError:
|
||
results.append(obj)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("重启事件读取失败: %s", e)
|
||
|
||
return results
|
||
|
||
|
||
def _trim_restart_events() -> None:
|
||
"""清理过期的重启事件行。"""
|
||
cfg = get_diag_config()
|
||
days = cfg.history_retention_days
|
||
f = _restart_events_file()
|
||
if not f.exists():
|
||
return
|
||
|
||
cutoff = datetime.now() - timedelta(days=days)
|
||
kept: List[str] = []
|
||
try:
|
||
with open(f, "r", encoding="utf-8", errors="replace") as fp:
|
||
total = 0
|
||
for line in fp:
|
||
total += 1
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
obj = json.loads(line)
|
||
ts_str = obj.get("ts", "")
|
||
ts_dt = datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")
|
||
if ts_dt >= cutoff:
|
||
kept.append(line)
|
||
except Exception:
|
||
kept.append(line)
|
||
|
||
if kept and len(kept) < total:
|
||
_rewrite_file_locked(f, kept)
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("重启事件清理失败: %s", e)
|