YG_FT/backend/app/core/audit.py

"""
审计日志装饰器模块

提供 @audit_log 装饰器，用于自动记录关键业务操作的审计日志。

使用示例：
    from app.core.audit import audit_log

    @audit_log(action="create_dataset", target_type="dataset")
    async def create_dataset(request: Request, ...):
        ...
"""

from __future__ import annotations

import asyncio
import functools
import time
from datetime import datetime, timezone
from typing import Any, Callable, Optional, TypeVar

from fastapi import Request

from app.core.logging import get_client_ip, get_logger, mask_sensitive_string, request_id_var

logger = get_logger("app.audit")

F = TypeVar("F", bound=Callable[..., Any])


def audit_log(
    action: str,
    target_type: str = "",
    *,
    detail_template: str = "",
    extract_target_id: Optional[Callable[[Any], str]] = None,
) -> Callable[[F], F]:
    """
    审计日志装饰器

    Args:
        action: 操作类型，如 create_dataset、update_model 等
        target_type: 目标资源类型，如 dataset、model 等
        detail_template: 日志详情模板（支持 format 参数）
        extract_target_id: 从返回值中提取目标 ID 的函数

    Returns:
        装饰后的函数
    """

    def decorator(func: F) -> F:
        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                started_at = time.perf_counter()
                trace_id = request_id_var.get("-")
                try:
                    result = await func(*args, **kwargs)
                    elapsed_ms = (time.perf_counter() - started_at) * 1000
                    target_id = _extract_target_id(result, kwargs, extract_target_id)
                    detail = _build_detail(detail_template, kwargs)
                    _record_audit(
                        action=action,
                        actor_id=_extract_actor_id(kwargs),
                        target_type=target_type,
                        target_id=target_id,
                        detail=detail,
                        trace_id=trace_id,
                        duration_ms=elapsed_ms,
                        kwargs=kwargs,
                        args=args,
                    )
                    return result
                except Exception as exc:
                    _record_audit(
                        action=action,
                        actor_id=_extract_actor_id(kwargs),
                        target_type=target_type,
                        target_id=_extract_target_id(None, kwargs, extract_target_id),
                        detail=_build_detail(detail_template, kwargs),
                        trace_id=trace_id,
                        duration_ms=(time.perf_counter() - started_at) * 1000,
                        result="failure",
                        reason=_safe_exception_reason(exc),
                        kwargs=kwargs,
                        args=args,
                    )
                    logger.warning("业务操作失败 action=%s reason=%s", action, _safe_exception_reason(exc))
                    raise

            return async_wrapper  # type: ignore
        else:
            @functools.wraps(func)
            def sync_wrapper(*args, **kwargs):
                started_at = time.perf_counter()
                trace_id = request_id_var.get("-")
                try:
                    result = func(*args, **kwargs)
                    elapsed_ms = (time.perf_counter() - started_at) * 1000
                    target_id = _extract_target_id(result, kwargs, extract_target_id)
                    detail = _build_detail(detail_template, kwargs)
                    _record_audit(
                        action=action,
                        actor_id=_extract_actor_id(kwargs),
                        target_type=target_type,
                        target_id=target_id,
                        detail=detail,
                        trace_id=trace_id,
                        duration_ms=elapsed_ms,
                        kwargs=kwargs,
                        args=args,
                    )
                    return result
                except Exception as exc:
                    _record_audit(
                        action=action,
                        actor_id=_extract_actor_id(kwargs),
                        target_type=target_type,
                        target_id=_extract_target_id(None, kwargs, extract_target_id),
                        detail=_build_detail(detail_template, kwargs),
                        trace_id=trace_id,
                        duration_ms=(time.perf_counter() - started_at) * 1000,
                        result="failure",
                        reason=_safe_exception_reason(exc),
                        kwargs=kwargs,
                        args=args,
                    )
                    logger.warning("业务操作失败 action=%s reason=%s", action, _safe_exception_reason(exc))
                    raise

            return sync_wrapper  # type: ignore

    return decorator


def _extract_target_id(
    result: Any, kwargs: dict, extractor: Optional[Callable[[Any], str]]
) -> Optional[str]:
    """从返回值或 kwargs 中提取目标 ID"""
    if extractor:
        try:
            return extractor(result)
        except Exception:
            pass
    if isinstance(result, dict):
        return result.get("id")
    # 尝试从路径参数中提取
    for key in ("dataset_id", "model_id", "task_id", "resource_id"):
        val = kwargs.get(key)
        if val:
            return str(val)
    return None


def _build_detail(template: str, kwargs: dict) -> str:
    """构建审计详情"""
    if not template:
        return ""
    try:
        return template.format(**kwargs)
    except (KeyError, IndexError):
        return template


def _record_audit(
    action: str,
    actor_id: Optional[str],
    target_type: str,
    target_id: Optional[str],
    detail: str,
    trace_id: str,
    duration_ms: float,
    *,
    kwargs: dict[str, Any] | None = None,
    args: tuple[Any, ...] = (),
    result: str = "success",
    reason: str | None = None,
) -> None:
    """通过已有的 record_audit 方法写入审计日志"""
    try:
        from app.db.platform_store import get_platform_store

        store = get_platform_store()
        kwargs = kwargs or {}
        request = _extract_request(args, kwargs)
        current_user = kwargs.get("current_user") or kwargs.get("user") or {}
        request_id = request.headers.get("X-Request-ID") if request else None
        request_id = request_id or trace_id
        client_ip = get_client_ip(request) or None
        detail_text = f"{detail} trace_id={trace_id} duration_ms={duration_ms:.1f}" if detail else f"trace_id={trace_id} duration_ms={duration_ms:.1f}"
        store.record_audit(
            action=action,
            actor_id=actor_id,
            target_type=target_type or None,
            target_id=target_id,
            tenant_id=str(current_user.get("tenant_id") or "") or None,
            detail=mask_sensitive_string(detail_text),
            result=result,
            reason=mask_sensitive_string(reason or "") or None,
            request_id=request_id,
            session_id=str(current_user.get("session_id") or "") or None,
            ip=client_ip,
        )
    except Exception:
        logger.error("写入审计日志失败 action=%s", action, exc_info=True)


def _extract_actor_id(kwargs: dict) -> Optional[str]:
    """从 FastAPI 注入的当前用户中提取操作人 ID。"""
    for key in ("current_user", "user"):
        value = kwargs.get(key)
        if isinstance(value, dict) and value.get("id"):
            return str(value["id"])
    return None


def _extract_request(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Request | None:
    for value in tuple(kwargs.values()) + tuple(args):
        if isinstance(value, Request):
            return value
    return None


def _safe_exception_reason(exc: Exception) -> str:
    """Keep audit failures useful without recording credentials or tokens."""
    value = getattr(exc, "detail", None) or str(exc) or exc.__class__.__name__
    return mask_sensitive_string(str(value))[:500]


# ==================== 预定义的审计操作常量 ====================

class AuditActions:
    """预定义的审计操作类型"""
    # 数据集操作
    CREATE_DATASET = "create_dataset"
    UPDATE_DATASET = "update_dataset"
    DELETE_DATASET = "delete_dataset"

    # 模型操作
    CREATE_MODEL = "create_model"
    UPDATE_MODEL = "update_model"
    DELETE_MODEL = "delete_model"

    # 微调任务
    CREATE_FINE_TUNE = "create_fine_tune"
    UPDATE_FINE_TUNE = "update_fine_tune"
    DELETE_FINE_TUNE = "delete_fine_tune"

    # 推理任务
    CREATE_INFERENCE = "create_inference"
    UPDATE_INFERENCE = "update_inference"
    DELETE_INFERENCE = "delete_inference"

    # 用户管理
    CREATE_USER = "create_user"
    UPDATE_USER = "update_user"
    DELETE_USER = "delete_user"

    # 租户管理
    CREATE_TENANT = "create_tenant"
    UPDATE_TENANT = "update_tenant"
    DELETE_TENANT = "delete_tenant"

    # 权限授权
    GRANT_ACL = "grant_acl"
    REVOKE_ACL = "revoke_acl"

    # 系统配置
    UPDATE_CONFIG = "update_config"