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_logger, 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,
                        target_type=target_type,
                        target_id=target_id,
                        detail=detail,
                        trace_id=trace_id,
                        duration_ms=elapsed_ms,
                    )
                    return result
                except Exception:
                    logger.error(
                        "审计日志记录失败 action=%s", action, exc_info=True
                    )
                    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,
                        target_type=target_type,
                        target_id=target_id,
                        detail=detail,
                        trace_id=trace_id,
                        duration_ms=elapsed_ms,
                    )
                    return result
                except Exception:
                    logger.error(
                        "审计日志记录失败 action=%s", action, exc_info=True
                    )
                    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,
    target_type: str,
    target_id: Optional[str],
    detail: str,
    trace_id: str,
    duration_ms: float,
) -> None:
    """通过已有的 record_audit 方法写入审计日志"""
    try:
        from app.db.platform_store import get_platform_store

        store = get_platform_store()
        store.record_audit(
            action=action,
            target_type=target_type or None,
            target_id=target_id,
            detail=f"{detail} trace_id={trace_id} duration_ms={duration_ms:.1f}" if detail else f"trace_id={trace_id} duration_ms={duration_ms:.1f}",
        )
    except Exception:
        logger.error("写入审计日志失败 action=%s", action, exc_info=True)


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

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"