buff2.0/backend-admin/app/middleware/audit.py

"""
异步/批量审计写入器（含采样与轮转）
- errors: 100% 写入
- success: 默认采样率 1%（可通过环境/配置调整）
- 批量写入：达到 batch_size 或超时则刷新
- 日志轮转：按天文件，并支持按大小轮转（max_bytes）
"""
import os
import json
import hashlib
import threading
import queue
import time
from datetime import datetime
from app.utils.timezone import beijing_now_naive
from typing import Any, Dict, List, Optional
from app.core.logger import logger as core_logger
# 后端固定审计参数（硬编码）
# - AUDIT_SUCCESS_SAMPLE_RATE: 成功请求采样率（0.0 - 1.0）
# - AUDIT_BATCH_SIZE: 批量写入阈值
# - AUDIT_FLUSH_INTERVAL: 最长等待秒数（到达则强制刷新批次）
# - AUDIT_MAX_BYTES: 日志文件大小轮转阈值
AUDIT_SUCCESS_SAMPLE_RATE = 0.01  # 1% 成功请求采样
AUDIT_BATCH_SIZE = 50
AUDIT_FLUSH_INTERVAL = 5  # seconds
AUDIT_MAX_BYTES = 10 * 1024 * 1024  # 10 MB

# 日志目录（确保存在）
SECURE_LOG_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "logs", "secure")
SECURE_LOG_DIR = os.path.abspath(SECURE_LOG_DIR)
os.makedirs(SECURE_LOG_DIR, exist_ok=True)

# 要脱敏的字段
REDACT_FIELDS = {
    "password",
    "card_code",
    "shared_secret",
    "identity_secret",
    "aes_key",
    "encrypted_aes_key",
    "signature",
    "access_token",
    "refresh_token",
    "session_token",
    "steamLoginSecure",
}


def _redact(obj: Any, redacted: Optional[List[str]] = None) -> Any:
    """递归脱敏对象中的敏感字段"""
    if redacted is None:
        redacted = []
    if isinstance(obj, dict):
        out = {}
        for k, v in obj.items():
            if k in REDACT_FIELDS:
                out[k] = "<REDACTED>"
                redacted.append(k)
            else:
                out[k] = _redact(v, redacted)
        return out
    if isinstance(obj, list):
        return [_redact(i, redacted) for i in obj]
    return obj


def _payload_hash(payload: Any) -> str:
    try:
        j = json.dumps(payload, sort_keys=True, ensure_ascii=False)
        return hashlib.sha256(j.encode("utf-8")).hexdigest()
    except Exception:
        return ""


class AuditWriter:
    def __init__(self):
        self._q: "queue.Queue[dict]" = queue.Queue()
        self._stop_event = threading.Event()
        self._worker = threading.Thread(target=self._run, daemon=True)
        self._worker.start()

    def enqueue(self, entry: Dict[str, Any]):
        try:
            self._q.put_nowait(entry)
        except Exception:
            core_logger.exception("enqueue 审计条目失败")

    def _current_log_path(self) -> str:
        fname = beijing_now_naive().strftime("secure-%Y-%m-%d.jsonl")
        return os.path.join(SECURE_LOG_DIR, fname)

    def _rotate_if_needed(self, path: str):
        try:
            if os.path.exists(path) and os.path.getsize(path) >= AUDIT_MAX_BYTES:
                ts = beijing_now_naive().strftime("%Y%m%d%H%M%S")
                new_name = f"{path}.{ts}.rot"
                os.rename(path, new_name)
        except Exception:
            core_logger.exception("审计日志轮转失败")

    def _flush_batch(self, batch: List[Dict[str, Any]]):
        if not batch:
            return
        path = self._current_log_path()
        try:
            self._rotate_if_needed(path)
            with open(path, "a", encoding="utf-8") as f:
                for entry in batch:
                    f.write(json.dumps(entry, ensure_ascii=False) + "\n")
        except Exception:
            core_logger.exception("批量写审计日志失败")

    def _run(self):
        buffer: List[Dict[str, Any]] = []
        last_flush = time.time()
        while not self._stop_event.is_set():
            try:
                timeout = max(0.0, AUDIT_FLUSH_INTERVAL - (time.time() - last_flush))
                item = None
                try:
                    item = self._q.get(timeout=timeout)
                except queue.Empty:
                    item = None

                if item is not None:
                    buffer.append(item)

                # flush conditions
                if len(buffer) >= AUDIT_BATCH_SIZE or (time.time() - last_flush) >= AUDIT_FLUSH_INTERVAL:
                    self._flush_batch(buffer)
                    buffer = []
                    last_flush = time.time()
            except Exception:
                core_logger.exception("审计写入工作线程异常")

        # on stop, flush remaining
        try:
            remaining = []
            while not self._q.empty():
                remaining.append(self._q.get_nowait())
            self._flush_batch(remaining)
        except Exception:
            core_logger.exception("停机时刷新审计队列失败")

    def stop(self):
        self._stop_event.set()
        self._worker.join(timeout=2)


# 全局实例
_AUDIT_WRITER = AuditWriter()


def write_audit_entry(
    phase: str,
    endpoint: str,
    payload: Any,
    key_version: Optional[int] = None,
    session_id: Optional[str] = None,
    level: str = "info",
    processing_time_ms: Optional[int] = None,
):
    """
    采样逻辑：
    - 当 level == "error" 时始终写入（errors=100%）
    - 其他（成功）按 AUDIT_SUCCESS_SAMPLE_RATE 采样
    """
    try:
        if level != "error":
            # sample successes
            import random

            if random.random() >= AUDIT_SUCCESS_SAMPLE_RATE:
                return

        timestamp = beijing_now_naive().strftime("%Y-%m-%dT%H:%M:%S") + "+08:00"
        redacted_fields = []
        redacted_payload = _redact(payload, redacted_fields)

        entry = {
            "timestamp": timestamp,
            "phase": phase,
            "endpoint": endpoint,
            "key_version": key_version,
            "session_id": session_id,
            "processing_time_ms": processing_time_ms,
            "redacted_fields": list(set(redacted_fields)),
            "payload_hash": _payload_hash(payload),
            "truncated": False,
            "level": level,
            "payload": redacted_payload if level != "error" else None,
        }

        _AUDIT_WRITER.enqueue(entry)
    except Exception:
        core_logger.exception("写审计日志失败")