FastapiAdmin/backend/app/api/v1/module_system/auth/controller.py

# -*- coding: utf-8 -*-

import time
from typing import Union, Dict
from fastapi import APIRouter, Depends, Request
from fastapi.responses import JSONResponse
from sqlalchemy.ext.asyncio import AsyncSession
from redis.asyncio.client import Redis

from app.common.response import ErrorResponse, SuccessResponse
from app.core.router_class import OperationLogRoute
from app.core.security import CustomOAuth2PasswordRequestForm
from app.core.logger import logger
from app.config.setting import settings
from app.core.dependencies import (
    db_getter,
    get_current_user,
    redis_getter
)
from .service import (
    LoginService,
    CaptchaService
)
from .schema import (
    CaptchaOutSchema,
    JWTOutSchema,
    RefreshTokenPayloadSchema,
    LogoutPayloadSchema
)


AuthRouter = APIRouter(route_class=OperationLogRoute, prefix="/auth", tags=["认证授权"])


@AuthRouter.post("/login", summary="登录", description="登录", response_model=JWTOutSchema)
async def login_for_access_token_controller(
    request: Request,
    redis: Redis = Depends(redis_getter), 
    login_form: CustomOAuth2PasswordRequestForm = Depends(),
    db: AsyncSession = Depends(db_getter),
) -> Union[JSONResponse, Dict]:
    """
    用户登录

    参数:
    - request (Request): FastAPI请求对象
    - login_form (CustomOAuth2PasswordRequestForm): 登录表单数据
    - db (AsyncSession): 数据库会话对象
        
    返回:
    - JWTOutSchema: 包含访问令牌和刷新令牌的响应模型
        
    异常:
    - CustomException: 认证失败时抛出异常。
    """
    login_token = await LoginService.authenticate_user_service(request=request, redis=redis, login_form=login_form, db=db)

    logger.info(f"用户{login_form.username}登录成功")

    # 如果是文档请求，则不记录日志:http://localhost:8000/api/v1/docs
    if settings.DOCS_URL in request.headers.get("referer", ""):
        return login_token.model_dump()
    return SuccessResponse(data=login_token.model_dump(), msg="登录成功")


@AuthRouter.post("/token/refresh", summary="刷新token", description="刷新token", response_model=JWTOutSchema, dependencies=[Depends(get_current_user)])
async def get_new_token_controller(
    request: Request,
    payload: RefreshTokenPayloadSchema,
    db: AsyncSession = Depends(db_getter),
    redis: Redis = Depends(redis_getter) 
) -> JSONResponse:
    """
    刷新token

    参数:
    - request (Request): FastAPI请求对象
    - payload (RefreshTokenPayloadSchema): 刷新令牌负载模型
        
    返回:
    - JWTOutSchema: 包含新的访问令牌和刷新令牌的响应模型
        
    异常:
    - CustomException: 刷新令牌失败时抛出异常。
    """
    # 解析当前的访问Token以获取用户名
    new_token = await LoginService.refresh_token_service(db=db, request=request, redis=redis, refresh_token=payload)
    token_dict = new_token.model_dump()
    logger.info(f"刷新token成功: {token_dict}")
    return SuccessResponse(data=token_dict, msg="刷新成功")


@AuthRouter.get("/captcha/get", summary="获取验证码", description="获取登录验证码", response_model=CaptchaOutSchema)
async def get_captcha_for_login_controller(
    redis: Redis = Depends(redis_getter)
) -> JSONResponse:
    """
    获取登录验证码

    参数:
    - redis (Redis): Redis客户端对象
        
    返回:
    - CaptchaOutSchema: 包含验证码图片和key的响应模型
        
    异常:
    - CustomException: 获取验证码失败时抛出异常。
    """
    # 获取验证码
    captcha = await CaptchaService.get_captcha_service(redis=redis)
    logger.info(f"获取验证码成功")
    return SuccessResponse(data=captcha, msg="获取验证码成功")


@AuthRouter.post('/logout', summary="退出登录", description="退出登录", dependencies=[Depends(get_current_user)])
async def logout_controller(
    payload: LogoutPayloadSchema,
    redis: Redis = Depends(redis_getter)
) -> JSONResponse:
    """
    退出登录

    参数:
    - payload (LogoutPayloadSchema): 退出登录负载模型
    - redis (Redis): Redis客户端对象
        
    返回:
    - JSONResponse: 包含退出登录结果的响应模型
        
    异常:
    - CustomException: 退出登录失败时抛出异常。
    """
    if await LoginService.logout_service(redis=redis, token=payload):
        logger.info('退出成功')
        return SuccessResponse(msg='退出成功')
    return ErrorResponse(msg='退出失败')