FastAPI与Swagger UI集成OAuth2认证：提升API测试效率

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from fastapi import HTTPException
from typing import Callable
import time
# from firebase_admin import auth # 假设已初始化Firebase Admin SDK

class AuthenticationMiddleware(BaseHTTPMiddleware):
    """用于使用Firebase Auth验证请求的中间件。"""

    async def dispatch(self, request: Request, call_next: Callable):
        path = request.url.path
        # 排除特定路径，例如健康检查、登录、文档页面
        if path in ["/health", "/auth/login", "/docs", "/openapi.json"]:
            return await call_next(request)

        headers = request.headers
        token = headers.get("Authorization")

        if not token:
            raise HTTPException(status_code=401, detail="Authorization token is missing")

        try:
            # 假设令牌格式为 "Bearer <token>"
            scheme, credentials = token.split()
            if scheme.lower() != 'bearer':
                raise HTTPException(status_code=401, detail="Invalid authorization scheme")

            # 实际应用中，此处应调用Firebase Admin SDK验证ID Token
            # user_info = auth.verify_id_token(credentials)
            # if user_info.get("exp") < time.time():
            #     raise HTTPException(status_code=401, detail="Token expired")

            # 简化示例：假设令牌验证成功，并获取到用户信息
            user_info = {"uid": "test_user_id", "email": "test@example.com", "exp": time.time() + 3600} # 模拟有效令牌
            request.state.user = user_info
            return await call_next(request)

        except Exception as err:
            # print(f"Authentication error: {err}") # 生产环境应记录日志
            raise HTTPException(status_code=401, detail="Invalid token") from err

from fastapi import FastAPI, Depends, HTTPException, Security
from fastapi.security import OAuth2AuthorizationCodeBearer, HTTPAuthorizationCredentials
import time

app = FastAPI()

# 配置OAuth2授权码流
# 注意：这里以Google OAuth2端点作为示例，因为Firebase认证通常会与Google账户关联。
# 您需要根据实际的OAuth2提供商和您的Firebase项目配置进行调整。
oauth2_scheme = OAuth2AuthorizationCodeBearer(
    authorizationUrl="https://accounts.google.com/o/oauth2/v2/auth", # OAuth2提供商的授权端点
    tokenUrl="https://oauth2.googleapis.com/token", # OAuth2提供商的令牌端点
    scopes={"openid": "获取用户OpenID", "email": "获取用户邮箱地址"}, # 定义请求的作用域
    # refreshUrl="https://oauth2.googleapis.com/token" # 可选，用于令牌刷新
)

# ... (您的AuthenticationMiddleware定义，如果需要继续使用) ...
# app.add_middleware(AuthenticationMiddleware)

# 认证依赖函数
async def get_current_user_token(
    credentials: HTTPAuthorizationCredentials = Security(oauth2_scheme),
) -> str:
    """
    此依赖项将从请求头中提取Bearer令牌。
    在实际应用中，您应在此处验证此令牌的有效性，
    例如，如果令牌是Firebase ID Token，则使用Firebase Admin SDK进行验证。
    """
    token = credentials.credentials # 这是实际的Bearer令牌

    try:
        # 示例：验证令牌（如果它是Firebase ID Token）
        # from firebase_admin import auth
        # decoded_token = auth.verify_id_token(token)
        # if decoded_token.get("exp") < time.time():
        #     raise HTTPException(status_code=401, detail="Token expired")
        # request.state.user = decoded_token # 将解码后的用户信息存储在请求状态中

        # 为简化教程，此处仅返回令牌，实际应用中需进行严格验证
        print(f"Received token: {token}")
        return token
    except Exception as e:
        raise HTTPException(status_code=401, detail=f"Invalid authentication credentials: {e}")

@app.get("/protected-resource", summary="获取受保护的资源")
async def read_protected_resource(current_token: str = Depends(get_current_user_token)):
    """
    这是一个受OAuth2保护的API端点。
    只有通过认证的用户才能访问。
    """
    return {"message": "您已成功访问受保护的资源！", "your_token": current_token}

@app.get("/public-resource", summary="获取公开资源")
async def read_public_resource():
    """
    这是一个公开的API端点，无需认证即可访问。
    """
    return {"message": "这是一个公开的资源。"}

from starlette.middleware.authentication import AuthenticationMiddleware, SimpleUser, AuthCredentials
from starlette.authentication import BaseUser, UnauthenticatedUser
from starlette.requests import Request
from starlette.responses import PlainTextResponse
from starlette.types import ASGIApp
# from firebase_admin import auth # 假设已初始化Firebase Admin SDK

class AuthenticatedUser(SimpleUser):
    """自定义用户类，可以存储更多用户信息，如Firebase UID、email等。"""
    def __init__(self, uid: str, email: str, display_name: str = None):
        super().__init__(display_name or email) # 使用email作为名称，如果display_name不存在
        self.uid = uid
        self.email = email

class FirebaseAuthenticationBackend:
    async def authenticate(self, conn: Request):
        if "Authorization" not in conn.headers:
            return

        auth_header = conn.headers["Authorization"]
        try:
            scheme, credentials = auth_header.split()
            if scheme.lower() == 'bearer':
                # 实际应用中，这里应调用Firebase Admin SDK验证credentials
                # decoded_token = auth.verify_id_token(credentials)
                # return AuthCredentials(["authenticated"]), AuthenticatedUser(
                #     uid=decoded_token['uid'],
                #     email=decoded_token['email'],
                #     display_name=decoded_token.get('name')
                # )

                # 简化示例：假设令牌验证成功，并模拟用户信息
                mock_uid = f"user_{credentials[:5]}" # 从令牌前缀生成模拟UID
                mock_email = f"{mock_uid}@example.com"
                return AuthCredentials(["authenticated"]), AuthenticatedUser(uid=mock_uid, email=mock_email)
        except ValueError:
            pass # 令牌格式不正确
        except Exception as e:
            # print(f"Firebase authentication failed: {e}") # 记录日志
            pass # 令牌验证失败

        return # 认证失败

# app = FastAPI() # 假设已创建FastAPI实例
# app.add_middleware(AuthenticationMiddleware, backend=FirebaseAuthenticationBackend())

# 之后，在任何路由中都可以通过 request.user 访问认证用户
# @app.get("/whoami")
# async def who_am_i(request: Request):
#     if request.user.is_authenticated:
#         return {"uid": request.user.uid, "email": request.user.email}
#     return {"message": "Not authenticated"}