buff/backend/app/middleware/logging.py

"""
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
from typing import Any

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），其它方法不记录请求体
        request_body = None
        if method in ["POST", "PUT", "PATCH"]:
            try:
                body = await request.body()
                if body:  # 只在有内容时记录
                    request_body = body.decode('utf-8')[:1000]  # 限制长度
            except:
                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()
                )
                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)

        # 调用下游应用并在返回后记录日志
        try:
            await self.app(scope, receive_wrapper, send_wrapper)
        finally:
            try:
                response_body_bytes = b"".join(resp_chunks) if resp_chunks else b""
                try:
                    response_text = response_body_bytes.decode("utf-8", errors="ignore") if response_body_bytes else None
                except Exception:
                    response_text = str(response_body_bytes)
                if not response_text:
                    response_text = f"Status: {status_code}"
                if len(response_text) > 2000:
                    response_text = response_text[:2000] + "..."

                request_text = None
                if body_bytes:
                    try:
                        request_text = body_bytes.decode("utf-8", errors="ignore")
                        if len(request_text) > 2000:
                            request_text = request_text[:2000] + "..."
                    except Exception:
                        request_text = str(body_bytes)

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

                # 根据路径选择对应的logger
                api_logger = get_logger_by_path(path)

                # 在控制台打印API入参出参
                try:
                    request_display = request_text[:1000] + ('...' if request_text and len(request_text) > 1000 else '') if request_text else 'None'
                    response_display = response_text[:1000] + ('...' if response_text and len(response_text) > 1000 else '') if response_text else 'None'
                    api_logger.info(f"[API请求] {method} {url} - 请求数据: {request_display}")
                    api_logger.info(f"[API响应] {method} {url} - 响应数据: {response_display} - 执行时间: {execution_time}ms")
                except Exception as e:
                    api_logger.warning(f"[API日志] 控制台日志输出失败: {e}")

                # 写入数据库（异步性与错误容忍）
                try:
                    db = SessionLocal()
                    try:
                        # 临时调试日志：打印请求头中的 X-User-Id/X-User-Name 与 scope 注入的 verified ids
                        try:
                            headers_lower = {k.lower(): v for k, v in headers.items()} if 'headers' in locals() else {}
                            xu = headers_lower.get('x-user-id') or headers_lower.get('x-user-name') or None
                        except Exception:
                            xu = None
                        api_logger.debug(f"[APILoggingASGI][DEBUG] endpoint={url} verified_user_id={scope.get('verified_user_id')} verified_admin_id={scope.get('verified_admin_id')} header_X_User_Id={xu}")
                        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_text if request_text else None,
                            response_data=response_text if response_text else None,
                            ip_address=client_ip,
                            user_agent=user_agent,
                            response_code=status_code,
                            execution_time=execution_time,
                            created_at=beijing_now()
                        )
                        db.add(api_log)
                        db.commit()
                        api_logger.debug(f"[APILoggingASGI] 成功记录API日志: {url} {method} {status_code} {execution_time}ms")
                    except Exception as e:
                        db.rollback()
                        api_logger.error(f"[APILoggingASGI] 记录API日志失败: {url} {method} {status_code}", exc_info=True)
                    finally:
                        db.close()
                except Exception as e:
                    api_logger.error(f"[APILoggingASGI] 创建数据库连接失败: {e}")
            except Exception:
                pass