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:
113
app/security/ratelimit.py
Normal file
113
app/security/ratelimit.py
Normal file
@@ -0,0 +1,113 @@
|
||||
"""接口限流(Token bucket 滑窗)。
|
||||
|
||||
设计:
|
||||
- 客户端维度(按 client_id 或 IP),每分钟请求配额
|
||||
- 优先用 Redis(多 worker 共享计数),Redis 不可用时降级为内存字典
|
||||
- 返回 429 + Retry-After header
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import Depends, HTTPException, Request, status
|
||||
|
||||
# ---------- 内存降级模式 ----------
|
||||
# key -> (window_start_timestamp, count)
|
||||
_MEM_BUCKET: dict[str, tuple[float, int]] = {}
|
||||
_MEM_LOCK = threading.Lock()
|
||||
_MEM_WINDOW = 60 # 秒
|
||||
|
||||
|
||||
def _get_redis():
|
||||
"""懒加载 Redis(避免模块 import 时没装直接崩)。"""
|
||||
try:
|
||||
import redis # type: ignore
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
from app.config import get_settings
|
||||
try:
|
||||
return redis.from_url(get_settings().redis_url, decode_responses=True)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
|
||||
|
||||
def _memory_check(key: str, limit: int) -> bool:
|
||||
"""内存滑窗:返回 True 表示允许,False 表示超限。"""
|
||||
now = time.time()
|
||||
with _MEM_LOCK:
|
||||
start_ts, count = _MEM_BUCKET.get(key, (now, 0.0))
|
||||
if now - start_ts > _MEM_WINDOW:
|
||||
# 窗口外,重建
|
||||
_MEM_BUCKET[key] = (now, 1)
|
||||
return True
|
||||
if count >= limit:
|
||||
return False
|
||||
_MEM_BUCKET[key] = (start_ts, count + 1)
|
||||
return True
|
||||
|
||||
|
||||
def _redis_check(client_key: str, limit: int) -> bool:
|
||||
"""Redis 滑窗(INCR + EXPIRE)。"""
|
||||
r = _get_redis()
|
||||
if r is None:
|
||||
return _memory_check(client_key, limit)
|
||||
try:
|
||||
key = f"meshsafe:rl:{client_key}"
|
||||
pipe = r.pipeline()
|
||||
pipe.incr(key)
|
||||
pipe.expire(key, _MEM_WINDOW)
|
||||
results = pipe.execute()
|
||||
count = results[0]
|
||||
return count <= limit
|
||||
except Exception: # noqa: BLE001
|
||||
return _memory_check(client_key, limit)
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
"""限流依赖:返回一个 callable 给 Depends。
|
||||
|
||||
用法:
|
||||
@app.post("/scan/sync")
|
||||
async def scan_sync(..., limiter=Depends(RateLimiter(max_requests_per_min=60))):
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, max_requests_per_min: int = 60, label: str = "scan") -> None:
|
||||
self.max_requests = max_requests_per_min
|
||||
self.label = label
|
||||
|
||||
async def __call__(self, request: Request) -> None:
|
||||
# 优先用 client_id(客户端签名带上的),退而用 IP
|
||||
ctx = getattr(request.state, "client", None)
|
||||
client_id = (ctx.client_id if ctx else None) or getattr(request.state, "client_id", None) or request.client.host
|
||||
key = f"{self.label}:{client_id}"
|
||||
if not _redis_check(key, self.max_requests):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail=f"Rate limit exceeded ({self.label}: {self.max_requests}/min)",
|
||||
headers={"Retry-After": str(_MEM_WINDOW)},
|
||||
)
|
||||
|
||||
|
||||
def rate_limit(max_requests_per_min: int = 60, label: str = "scan"):
|
||||
"""更简洁的限流装饰器(函数依赖)。
|
||||
|
||||
用法:
|
||||
@app.post("/scan/sync")
|
||||
async def scan_sync(..., _ = Depends(rate_limit(60, "scans"))):
|
||||
...
|
||||
"""
|
||||
async def _check(request: Request) -> None:
|
||||
ctx = getattr(request.state, "client", None)
|
||||
client_id = (ctx.client_id if ctx else None) or getattr(request.state, "client_id", None) or request.client.host
|
||||
key = f"{label}:{client_id}"
|
||||
if not _redis_check(key, max_requests_per_min):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||||
detail=f"Rate limit exceeded ({label}: {max_requests_per_min}/min)",
|
||||
headers={"Retry-After": str(_MEM_WINDOW)},
|
||||
)
|
||||
return _check
|
||||
Reference in New Issue
Block a user