LongjoeAgent/middlewares/registration_middleware.py

from fastapi import Request, HTTPException
from fastapi.responses import JSONResponse
from utils.registration_manager import registration_manager
from utils.logger import chat_logger


async def registration_check_middleware(request: Request, call_next):
    """
    注册检查中间件
    对需要注册验证的接口进行检查
    """
    # 定义需要检查注册状态的接口列表
    protected_paths = ["/chat", "/message_create_bill", "/ocr_create_bill", "/ocr"]
    # 只拦截POST /chat请求
    if request.url.path in protected_paths and request.method == "POST":
        if not await registration_manager.check_registration():
            chat_logger.warning(f"拒绝未注册访问: {request.client.host}")

            if request.url.path == "/chat":
                response_data = {
                    "final_answer": "服务未注册或注册已过期，请联系管理员。"
                }
            else:
                response_data = "服务未注册或注册已过期，请联系管理员。"

            return JSONResponse(status_code=403, content=response_data)

    # 其他请求直接放行
    response = await call_next(request)
    return response