loom/backend/app/core/middleware.py

"""请求审计中间件（纯 ASGI 实现）。

不使用 `BaseHTTPMiddleware`，因为它内部借助 anyio task group 执行 call_next，
会与 asyncpg 的底层 future 跨 loop（TestClient 场景）或出现 request body 重放坑。

对变更接口（POST/PATCH/PUT/DELETE）成功响应后异步落 OpLog。
"""
from __future__ import annotations

import asyncio
import json
import time
from typing import Any

from sqlalchemy.exc import SQLAlchemyError
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from app.core.logging import logger
from app.core.security import JWTError, decode_token
from app.db.session import AsyncSessionLocal
from app.models import OpLog, User

MUTATION_METHODS = {"POST", "PATCH", "PUT", "DELETE"}
SKIP_PATHS = {
    "/api/v1/auth/login",
    "/api/v1/auth/refresh",
    "/api/v1/healthz",
}
SENSITIVE_KEYS = {"password", "google_key", "password_hash", "refresh_token", "access_token"}


def _sanitize(obj: Any) -> Any:
    if isinstance(obj, dict):
        return {
            k: ("***" if k in SENSITIVE_KEYS and v else _sanitize(v))
            for k, v in obj.items()
        }
    if isinstance(obj, list):
        return [_sanitize(i) for i in obj]
    return obj


def _infer_module(path: str) -> str:
    parts = path.strip("/").split("/")
    if len(parts) >= 3 and parts[0] == "api" and parts[1] == "v1":
        return parts[2]
    return "unknown"


def _infer_action(method: str, path: str) -> str:
    if path.endswith("/reset-password"):
        return "reset_password"
    if path.endswith("/test"):
        return "test"
    return {"POST": "create", "PATCH": "update", "PUT": "update", "DELETE": "delete"}.get(
        method, method.lower()
    )


class OpLogMiddleware:
    """纯 ASGI 中间件：只在请求/响应 message 流动的边界记录必要字段。"""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        method: str = scope.get("method", "GET").upper()
        path: str = scope.get("path", "")
        need_audit = (
            method in MUTATION_METHODS
            and path.startswith("/api/v1/")
            and path not in SKIP_PATHS
        )

        if not need_audit:
            await self.app(scope, receive, send)
            return

        # —— 收集请求体（一次性读取，再用缓存回放给下游）
        body_chunks: list[bytes] = []
        more = True
        while more:
            msg = await receive()
            if msg["type"] == "http.request":
                body_chunks.append(msg.get("body", b""))
                more = msg.get("more_body", False)
            else:
                more = False
        raw = b"".join(body_chunks)
        payload: Any = None
        if raw:
            try:
                payload = _sanitize(json.loads(raw))
            except json.JSONDecodeError:
                payload = {"_raw_bytes": len(raw)}

        replayed = False

        async def replay_receive() -> Message:
            nonlocal replayed
            if not replayed:
                replayed = True
                return {"type": "http.request", "body": raw, "more_body": False}
            return {"type": "http.disconnect"}

        # —— 拦响应的 start，拿 status
        status_code: int = 500
        start = time.perf_counter()

        async def send_wrapper(message: Message) -> None:
            nonlocal status_code
            if message["type"] == "http.response.start":
                status_code = int(message.get("status", 500))
            await send(message)

        await self.app(scope, replay_receive, send_wrapper)

        cost_ms = int((time.perf_counter() - start) * 1000)

        if 200 <= status_code < 400:
            headers = {
                k.decode("latin-1").lower(): v.decode("latin-1")
                for k, v in scope.get("headers", [])
            }
            client = scope.get("client") or (None, None)
            try:
                asyncio.create_task(
                    _write_oplog(
                        method=method,
                        path=path,
                        status_code=status_code,
                        payload=payload,
                        ip=client[0] if client else None,
                        user_agent=headers.get("user-agent"),
                        authorization=headers.get("authorization"),
                        cost_ms=cost_ms,
                    )
                )
            except RuntimeError:
                # 测试环境下 loop 可能已关闭；忽略，保证主路径不受影响
                pass


async def _write_oplog(
    *,
    method: str,
    path: str,
    status_code: int,
    payload: Any,
    ip: str | None,
    user_agent: str | None,
    authorization: str | None,
    cost_ms: int,
) -> None:
    operator_id: int | None = None
    operator_name: str | None = None
    if authorization and authorization.lower().startswith("bearer "):
        token = authorization.split(" ", 1)[1].strip()
        try:
            sub = decode_token(token).get("sub")
            operator_id = int(sub) if sub is not None else None
        except (JWTError, ValueError):
            operator_id = None

    try:
        async with AsyncSessionLocal() as db:
            if operator_id:
                u = await db.get(User, operator_id)
                operator_name = u.username if u else None
            entry = OpLog(
                operator_id=operator_id,
                operator_name=operator_name,
                module=_infer_module(path),
                action=_infer_action(method, path),
                method=method,
                path=path,
                status_code=status_code,
                ip=ip,
                user_agent=(user_agent or "")[:255] or None,
                payload=payload,
                cost_ms=cost_ms,
            )
            db.add(entry)
            await db.commit()
    except SQLAlchemyError:
        logger.exception("oplog_write_failed", path=path)
    except Exception:  # noqa: BLE001
        logger.exception("oplog_write_unhandled", path=path)