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
262 lines
8.8 KiB
Python
262 lines
8.8 KiB
Python
"""API 安全依赖：认证、签名校验、防重放、限流、响应签名。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||
|
||
from fastapi import Depends, HTTPException, Request, Response, status
|
||
from fastapi.responses import JSONResponse
|
||
from starlette.middleware.base import BaseHTTPMiddleware
|
||
|
||
from app.config import get_settings
|
||
from app.security import signing
|
||
from app.security.crypto import constant_time_equals
|
||
|
||
_settings = get_settings()
|
||
|
||
# redis 是可选依赖——没装时限流和防重放降级为内存策略（_check_rate_limit / nonce 检查里已兜底）
|
||
try:
|
||
import redis as _redis # type: ignore
|
||
except ModuleNotFoundError:
|
||
_redis = None
|
||
|
||
if TYPE_CHECKING:
|
||
import redis
|
||
|
||
_redis_client: Optional["redis.Redis"] = None
|
||
|
||
# 响应体原始字节暂存键，供响应签名中间件复用
|
||
RAW_BODY_ATTR = "_meshsafe_raw_body"
|
||
CLIENT_ATTR = "_meshsafe_client"
|
||
|
||
|
||
def get_redis():
|
||
"""懒加载 Redis 客户端。未安装 redis 包或连接失败时由调用方兜底。"""
|
||
global _redis_client
|
||
if _redis is None:
|
||
raise RuntimeError("redis 包未安装")
|
||
if _redis_client is None:
|
||
_redis_client = _redis.Redis.from_url(
|
||
_settings.redis_url, decode_responses=True, socket_timeout=5
|
||
)
|
||
return _redis_client
|
||
|
||
|
||
class ClientContext:
|
||
"""已通过认证的调用方上下文。"""
|
||
|
||
def __init__(self, client_id: str, secret: str, role: str = "operator") -> None:
|
||
self.client_id = client_id
|
||
self.secret = secret
|
||
self.role = role
|
||
|
||
|
||
def _lookup_client(client_id: str) -> Optional[ClientContext]:
|
||
for item in _settings.parsed_api_clients:
|
||
if constant_time_equals(item["client_id"], client_id):
|
||
return ClientContext(item["client_id"], item["secret"])
|
||
return None
|
||
|
||
|
||
def _check_rate_limit(client_id: str, path: str) -> None:
|
||
"""基于 Redis 的滑动窗口限流，防止接口被暴力试探。"""
|
||
try:
|
||
r = get_redis()
|
||
bucket = int(time.time() // 60)
|
||
key = f"ms:rl:{client_id}:{path}:{bucket}"
|
||
pipe = r.pipeline()
|
||
pipe.incr(key)
|
||
pipe.expire(key, 120)
|
||
count = pipe.execute()[0]
|
||
limit = 120 if "/scans" in path else 600
|
||
if count > limit:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
||
detail="请求过于频繁，已触发限流",
|
||
)
|
||
except HTTPException:
|
||
raise
|
||
except Exception: # noqa: BLE001 Redis 不可用时不阻断主流程
|
||
return
|
||
|
||
|
||
async def authenticate(request: Request) -> ClientContext:
|
||
"""认证与签名校验依赖。
|
||
|
||
依次执行：
|
||
1. 解析并校验调用方标识
|
||
2. 时间戳窗口校验（防重放）
|
||
3. nonce 一次性校验
|
||
4. 请求体哈希一致性校验
|
||
5. HMAC 签名校验
|
||
"""
|
||
client_id = request.headers.get("X-MS-Client-Id", "")
|
||
timestamp = request.headers.get("X-MS-Timestamp", "")
|
||
nonce = request.headers.get("X-MS-Nonce", "")
|
||
body_hash = request.headers.get("X-MS-Body-SHA256", "")
|
||
signature = request.headers.get("X-MS-Signature", "")
|
||
|
||
if not client_id or not timestamp or not nonce or not signature:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="缺少必要的认证请求头",
|
||
)
|
||
|
||
ctx = _lookup_client(client_id)
|
||
if ctx is None:
|
||
# 不区分"不存在"与"密钥错误"，避免枚举 client_id
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="认证失败")
|
||
|
||
# 1) 时间戳窗口
|
||
try:
|
||
ts_int = int(timestamp)
|
||
except ValueError:
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="认证失败")
|
||
skew = abs(int(time.time()) - ts_int)
|
||
if skew > _settings.signature_window:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail=f"请求时间戳超出允许偏差（{_settings.signature_window} 秒）",
|
||
)
|
||
|
||
# 2) nonce 一次性
|
||
try:
|
||
r = get_redis()
|
||
nonce_key = f"ms:nonce:{client_id}:{nonce}"
|
||
if not r.set(nonce_key, "1", nx=True, ex=_settings.signature_window * 2):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED, detail="重复请求，已拒绝"
|
||
)
|
||
except HTTPException:
|
||
raise
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
|
||
# 3) 读取请求体并校验哈希
|
||
# 注意：multipart/form-data 请求体已被 FastAPI UploadFile 消费，不能再 await request.body()
|
||
# 此时客户端已在 X-MS-Body-SHA256 里提供哈希，直接用即可
|
||
body: bytes = b""
|
||
ct = request.headers.get("content-type", "").lower()
|
||
if not ct.startswith("multipart/"):
|
||
try:
|
||
body = await request.body()
|
||
except RuntimeError:
|
||
# stream 已被消费（比如中间件或 UploadFile），用空 body 兜底
|
||
body = b""
|
||
|
||
skip_hash = ct.startswith("multipart/")
|
||
if not signing.verify_request_signature(
|
||
ctx.secret,
|
||
ctx.client_id,
|
||
request.method,
|
||
request.url.path,
|
||
dict(request.query_params),
|
||
body,
|
||
timestamp,
|
||
nonce,
|
||
body_hash,
|
||
signature,
|
||
skip_body_hash_check=skip_hash,
|
||
):
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="认证失败")
|
||
|
||
_check_rate_limit(ctx.client_id, request.url.path)
|
||
|
||
request.state.client = ctx
|
||
return ctx
|
||
|
||
|
||
async def require_admin(ctx: ClientContext = Depends(authenticate)) -> ClientContext:
|
||
"""要求调用方具备完整 findings 权限。"""
|
||
if ctx.role != "operator":
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail="当前凭证无权查看完整检测明细",
|
||
)
|
||
return ctx
|
||
|
||
|
||
class ResponseSigningMiddleware(BaseHTTPMiddleware):
|
||
"""为所有 API 响应添加签名头。
|
||
|
||
响应体在读取过程中暂存原始字节，签名基于原始字节计算，
|
||
避免因 JSON 序列化差异导致客户端验签失败。
|
||
"""
|
||
|
||
async def dispatch(self, request: Request, call_next):
|
||
response = await call_next(request)
|
||
|
||
if not _settings.sign_response:
|
||
return response
|
||
if request.url.path.startswith("/api/"):
|
||
ctx = getattr(request.state, "client", None)
|
||
if ctx is None:
|
||
return response
|
||
|
||
body = getattr(response, RAW_BODY_ATTR, None)
|
||
if body is None:
|
||
body = b""
|
||
if hasattr(response, "body_iterator"):
|
||
chunks = []
|
||
async for chunk in response.body_iterator:
|
||
chunks.append(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8"))
|
||
body = b"".join(chunks)
|
||
response = Response(
|
||
content=body,
|
||
status_code=response.status_code,
|
||
headers=dict(response.headers),
|
||
media_type=response.media_type,
|
||
)
|
||
# 删除旧 content-length（重建后会自动重新计算）
|
||
try:
|
||
del response.headers["content-length"]
|
||
except KeyError:
|
||
pass
|
||
|
||
headers = signing.sign_response(ctx.secret, ctx.client_id, body)
|
||
for key, value in headers.items():
|
||
response.headers[key] = value
|
||
return response
|
||
|
||
|
||
def build_signed_json(ctx: ClientContext, payload: Dict[str, Any], status_code: int = 200) -> JSONResponse:
|
||
"""构造带签名的 JSON 响应（供直接返回场景使用）。"""
|
||
body = signing.canonical_json_bytes(payload)
|
||
response = JSONResponse(content=payload, status_code=status_code)
|
||
if _settings.sign_response:
|
||
for key, value in signing.sign_response(ctx.secret, ctx.client_id, body).items():
|
||
response.headers[key] = value
|
||
return response
|
||
|
||
|
||
def redact_for_role(payload: Dict[str, Any], role: str) -> Dict[str, Any]:
|
||
"""按角色裁剪响应内容，降低信息泄露面。
|
||
|
||
``basic`` 角色只返回结论摘要，不返回引擎明细与提取文件路径，
|
||
避免通过差异响应反推检测逻辑。
|
||
"""
|
||
if role == "operator":
|
||
return payload
|
||
allowed_top = {
|
||
"task_id",
|
||
"verdict",
|
||
"score",
|
||
"action",
|
||
"reason",
|
||
"status",
|
||
"file",
|
||
"meta",
|
||
}
|
||
trimmed = {k: v for k, v in payload.items() if k in allowed_top}
|
||
file_info = trimmed.get("file")
|
||
if isinstance(file_info, dict):
|
||
trimmed["file"] = {
|
||
"sha256": file_info.get("sha256"),
|
||
"size": file_info.get("size"),
|
||
"real_type": file_info.get("real_type"),
|
||
}
|
||
trimmed["detail_truncated"] = True
|
||
return trimmed
|