buff3.0/backend-client-bak/app/middleware/logging.py

"""
API日志中间件
包含增强的API日志功能，记录入参密文、入参明文、出参明文、出参密文
"""
import time
import json
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.api_log import APILog
from app.core.logger import logger, get_logger_by_path
from app.core.security import verify_token
from app.models.user import User
from datetime import datetime
from app.utils.timezone import beijing_now_naive
from typing import Any

# 导入增强的API日志功能
from app.utils.api_context import (
    set_request_context,
    get_decrypted_request,
    get_decrypted_response,
)
from app.utils.api_logger import api_logger, redact_for_db

class LoggingResponse(Response):
    """用于捕获响应内容的响应类"""

    def __init__(self, content: Any, **kwargs):
        super().__init__(content, **kwargs)
        self.response_content = content

class APILoggingMiddleware(BaseHTTPMiddleware):
    """API日志中间件"""

    async def dispatch(self, request: Request, call_next):
        start_time = time.time()

        # 获取请求信息
        method = request.method
        url = str(request.url)
        client_ip = request.client.host if request.client else None
        user_agent = request.headers.get("user-agent", "")

        # 获取用户ID
        user_id = None
        try:
            auth_header = request.headers.get("authorization", "")
            if auth_header.startswith("Bearer "):
                token = auth_header[7:]  # 移除 "Bearer " 前缀
                payload = verify_token(token)
                if payload:
                    username = payload.get("sub")
                    if username:
                        # 从数据库获取用户ID（优先查 users 表；admin 用户不应写入 user_id）
                        db_temp = SessionLocal()
                        try:
                            user = db_temp.query(User).filter(User.username == username).first()
                            if user:
                                user_id = user.id
                        finally:
                            db_temp.close()
        except Exception as e:
            logger.warning("获取用户ID失败", exc_info=True)

        # 获取请求体（如果是POST/PUT/PATCH），其它方法不记录请求体
        # 安全端点 /api/v1/ 需在路由内自行读取 body 解密，此处不消费 body 以免导致 422/400
        request_body = None
        if method in ["POST", "PUT", "PATCH"] and not request.url.path.startswith("/api/v1/"):
            try:
                body = await request.body()
                if body:  # 只在有内容时记录
                    request_body = body.decode('utf-8')[:1000]  # 限制长度
            except Exception:
                pass

        # 执行请求
        response = await call_next(request)

        # 计算执行时间
        execution_time = int((time.time() - start_time) * 1000)

        # 获取响应状态码
        status_code = response.status_code

        # 获取响应数据：尽量读取完整响应体（若为 streaming response 则消费 body_iterator 并重建响应）
        response_body = None
        try:
            # 1) 如果 response.body 可用，直接读取
            if hasattr(response, "body") and response.body is not None:
                body = response.body
                if isinstance(body, bytes):
                    response_text = body.decode('utf-8', errors='ignore')
                else:
                    response_text = str(body)
                response_body = response_text
            # 2) 否则如果有 body_iterator，消费它并重建响应对象以便后续正确返回给客户端
            elif hasattr(response, "body_iterator") and response.body_iterator is not None:
                body_bytes = b""
                try:
                    async for chunk in response.body_iterator:
                        if isinstance(chunk, bytes):
                            body_bytes += chunk
                        else:
                            try:
                                body_bytes += str(chunk).encode('utf-8')
                            except:
                                pass
                except Exception:
                    body_bytes = b""
                # 重建响应以保持行为一致
                try:
                    from starlette.responses import Response as StarletteResponse
                    headers = dict(response.headers) if hasattr(response, "headers") else {}
                    media_type = getattr(response, "media_type", None)
                    response = StarletteResponse(content=body_bytes, status_code=status_code, headers=headers, media_type=media_type)
                except Exception:
                    pass
                try:
                    response_body = body_bytes.decode('utf-8', errors='ignore')
                except:
                    response_body = str(body_bytes)
        except Exception:
            response_body = None

        # 截断过长内容以防数据库过大，保留前 2000 字符
        if response_body:
            if len(response_body) > 2000:
                response_body = response_body[:2000] + '...'
        else:
            response_body = f"Status: {status_code}"
        # 兼容：某些 Response 实例（如 JSONResponse）可能把内容放在 .media 属性
        if response_body and response_body.startswith("Status:") and hasattr(response, "media"):
            try:
                import json as _json
                media = getattr(response, "media")
                if media is not None:
                    try:
                        response_text_from_media = _json.dumps(media, ensure_ascii=False)
                    except Exception:
                        response_text_from_media = str(media)
                    if response_text_from_media:
                        response_body = response_text_from_media if len(response_text_from_media) <= 2000 else response_text_from_media[:2000] + '...'
            except Exception:
                pass
        
        # 异步记录日志（避免阻塞）
        try:
            db = SessionLocal()
            try:
                api_log = APILog(
                    user_id=user_id,
                    admin_id=scope.get("verified_admin_id") if scope.get("verified_admin_id") else None,
                    endpoint=url,
                    method=method,
                    request_data=request_body if request_body else None,
                    # store the captured response body (or fallback status)
                    response_data=response_body if response_body else None,
                    ip_address=client_ip,
                    user_agent=user_agent,
                    response_code=status_code,
                    execution_time=execution_time,
                    created_at=beijing_now_naive()
                )
                db.add(api_log)
                db.commit()
            except Exception as e:
                db.rollback()
                logger.exception("记录API日志失败")
            finally:
                db.close()
        except:
            pass
        
        return response


class APILoggingASGIMiddleware:
    """更可靠的 ASGI 层日志中间件，捕获请求/响应流（用于记录完整返回报文）
    
    增强功能：
    - 记录入参密文（原始请求体）
    - 记录入参明文（解密后的请求参数，需要业务层配合设置）
    - 记录出参明文（业务响应数据，需要业务层配合设置）
    - 记录出参密文（加密后的响应数据）
    - 日志文件按天滚动，每天一个日志文件
    """

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return

        start_time = time.time()
        method = scope.get("method")
        server = scope.get("server") or ("localhost", 80)
        path = scope.get("raw_path", b"").decode("utf-8", errors="ignore")
        query_string = scope.get("query_string", b"").decode("utf-8", errors="ignore")
        scheme = scope.get("scheme", "http")
        url = f"{scheme}://{server[0]}:{server[1]}{path}" + (f"?{query_string}" if query_string else "")
        client = scope.get("client")
        client_ip = client[0] if client else None
        headers = {k.decode("utf-8"): v.decode("utf-8") for k, v in scope.get("headers", [])}
        user_agent = headers.get("user-agent", "")

        # 直接从请求头获取user_id和user_name，进行组合校验
        user_id = None
        admin_id = None
        try:
            # 从header获取user_id和user_name
            header_user_id = headers.get("x-user-id") or headers.get("X-User-Id")
            header_user_name = headers.get("x-user-name") or headers.get("X-User-Name")

            if header_user_id and header_user_name:
                # 验证header中的user_id和user_name是否匹配
                db_temp = SessionLocal()
                try:
                    # 首先尝试从User表查询
                    user = db_temp.query(User).filter(
                        User.id == int(header_user_id),
                        User.username == header_user_name,
                        User.is_deleted == 0
                    ).first()
                    if user:
                        user_id = user.id
                    else:
                        # 如果User表中没有，尝试从Admin表查询
                        from app.models.admin import Admin
                        admin = db_temp.query(Admin).filter(
                            Admin.id == int(header_user_id),
                            Admin.username == header_user_name
                        ).first()
                        if admin:
                            admin_id = admin.id
                finally:
                    db_temp.close()
        except Exception:
            # 如果header验证失败，记录为None
            user_id = None
            admin_id = None

        # 包装 receive，收集请求体（仅在有 body 时）
        body_bytes = b""
        async def receive_wrapper():
            nonlocal body_bytes
            message = await receive()
            if message.get("type") == "http.request":
                body_bytes += message.get("body", b"")
            return message

        # 包装 send，收集响应状态与 body 分片
        status_code = 500
        resp_chunks = []
        async def send_wrapper(message):
            nonlocal status_code, resp_chunks
            if message.get("type") == "http.response.start":
                status_code = message.get("status", 0)
            if message.get("type") == "http.response.body":
                chunk = message.get("body", b"")
                if chunk:
                    resp_chunks.append(chunk)
            await send(message)

        # 请求体要等下游（如 secure_endpoint）调用 request.body() 时才会通过 receive_wrapper 填入 body_bytes，
        # 因此此处仅初始化上下文；encrypted_request 在 finally 中根据收集到的 body_bytes 再赋值。
        set_request_context({
            "endpoint": url,
            "method": method,
            "user_id": user_id,
            "client_ip": client_ip,
            "encrypted_request": None,
            "encrypted_response": None,
            "decrypted_request": None,
            "decrypted_response": None,
        })

        error_message = None
        try:
            # 调用下游应用
            await self.app(scope, receive_wrapper, send_wrapper)
        except Exception as e:
            error_message = str(e)
            raise

        finally:
            try:
                # 此时 body_bytes 已由下游调用 request.body() 时通过 receive_wrapper 收集完毕
                encrypted_request = None
                if body_bytes:
                    try:
                        encrypted_request = body_bytes.decode("utf-8", errors="ignore")
                    except Exception:
                        encrypted_request = str(body_bytes)

                # 收集响应数据（出参密文）
                response_body_bytes = b"".join(resp_chunks) if resp_chunks else b""
                encrypted_response = None
                if response_body_bytes:
                    try:
                        encrypted_response = response_body_bytes.decode("utf-8", errors="ignore")
                    except Exception:
                        encrypted_response = str(response_body_bytes)

                # 计算执行时间
                processing_time_ms = int((time.time() - start_time) * 1000)

                # 获取解密后的请求和响应数据
                decrypted_request = get_decrypted_request()
                decrypted_response = get_decrypted_response()

                # 记录增强的API日志（包含明文和密文）
                api_logger.log_api_call(
                    endpoint=url,
                    method=method,
                    encrypted_request=encrypted_request,
                    decrypted_request=decrypted_request,
                    encrypted_response=encrypted_response,
                    decrypted_response=decrypted_response,
                    processing_time_ms=processing_time_ms,
                    status_code=status_code,
                    user_id=user_id,
                    client_ip=client_ip,
                    error_message=error_message,
                )

                # 根据路径选择对应的logger用于控制台输出
                console_logger = get_logger_by_path(path)

                # 在控制台打印API入参出参
                try:
                    request_display = encrypted_request[:500] + ('...[truncated]' if len(encrypted_request or "") > 500 else '') if encrypted_request else 'None'
                    response_display = encrypted_response[:500] + ('...[truncated]' if len(encrypted_response or "") > 500 else '') if encrypted_response else 'None'
                    console_logger.info(f"[API请求] {method} {url} - 请求数据(密文): {request_display}")
                    console_logger.info(f"[API响应] {method} {url} - 响应数据(密文): {response_display} - 执行时间: {processing_time_ms}ms")
                    
                    # 如果有明文数据，也打印明文（脱敏后）
                    if decrypted_request:
                        try:
                            import json
                            req_json = json.dumps(decrypted_request, ensure_ascii=False, default=str)[:500]
                            console_logger.info(f"[API请求] {method} {url} - 请求数据(明文): {req_json}...")
                        except Exception:
                            pass
                    if decrypted_response:
                        try:
                            import json
                            res_json = json.dumps(decrypted_response, ensure_ascii=False, default=str)[:500]
                            console_logger.info(f"[API响应] {method} {url} - 响应数据(明文): {res_json}...")
                        except Exception:
                            pass
                except Exception as e:
                    console_logger.warning(f"[API日志] 控制台日志输出失败: {e}")

                # 写入数据库：按文档 4.5 记录明文（脱敏），无明文时回退为密文截断
                try:
                    db = SessionLocal()
                    try:
                        request_data_for_db = (
                            redact_for_db(decrypted_request, 2000)
                            if decrypted_request is not None
                            else (encrypted_request[:2000] + "..." if encrypted_request and len(encrypted_request) > 2000 else encrypted_request)
                        )
                        response_data_for_db = (
                            redact_for_db(decrypted_response, 2000)
                            if decrypted_response is not None
                            else (encrypted_response[:2000] + "..." if encrypted_response and len(encrypted_response) > 2000 else encrypted_response)
                        )
                        api_log = APILog(
                            user_id=user_id,
                            admin_id=scope.get("verified_admin_id") if scope.get("verified_admin_id") else None,
                            endpoint=url,
                            method=method,
                            request_data=request_data_for_db,
                            response_data=response_data_for_db,
                            ip_address=client_ip,
                            user_agent=user_agent,
                            response_code=status_code,
                            execution_time=processing_time_ms,
                            created_at=beijing_now_naive()
                        )
                        db.add(api_log)
                        db.commit()
                        console_logger.debug(f"[APILoggingASGI] 成功记录API日志: {url} {method} {status_code} {processing_time_ms}ms")
                    except Exception as e:
                        db.rollback()
                        console_logger.error(f"[APILoggingASGI] 记录API日志失败: {url} {method} {status_code}", exc_info=True)
                    finally:
                        db.close()
                except Exception as e:
                    console_logger.error(f"[APILoggingASGI] 创建数据库连接失败: {e}")
            except Exception:
                pass