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

import json
import secrets
from typing import Annotated

from fastapi import APIRouter, Depends, Path, Query, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi_cache import FastAPICache
from fastapi_cache.decorator import cache
from redis.asyncio.client import Redis
from sqlalchemy.ext.asyncio import AsyncSession

from app.common.response import ErrorResponse, ResponseSchema, SuccessResponse
from app.config.setting import settings
from app.core.base_schema import (
    AuthSchema,
    JWTOutSchema,
    LogoutPayloadSchema,
    RefreshTokenPayloadSchema,
)
from app.core.dependencies import db_getter, get_current_user, redis_getter
from app.core.exceptions import CustomException
from app.core.logger import logger
from app.core.redis_crud import RedisCURD
from app.core.router_class import OperationLogRoute
from app.core.security import CustomOAuth2PasswordRequestForm

from .oauth_service import (
    STATE_PREFIX,
    _callback_url,
    build_authorize_url,
    complete_oauth_login,
    oauth_service_error_redirect,
    oauth_service_frontend_redirect_from_token,
    save_oauth_state,
)
from .schema import (
    AutoLoginTokenSchema,
    AutoLoginUserSchema,
    CaptchaOutSchema,
    ForgotPasswordSchema,
    LoginWithTenantsSchema,
    ResetPasswordWithTokenSchema,
    SelectTenantOutSchema,
    SelectTenantSchema,
    TenantOptionSchema,
    TenantRegisterOutSchema,
    TenantRegisterSchema,
)
from .service import (
    AutoLoginService,
    CaptchaService,
    LoginService,
    PasswordResetService,
    TenantRegisterService,
)

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

_AUTH_TENANTS_NS = "auth_tenants"


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

    参数:
    - request (Request): FastAPI请求对象
    - redis (Redis): Redis 客户端对象
    - login_form (CustomOAuth2PasswordRequestForm): 登录表单数据
    - db (AsyncSession): 数据库会话对象

    返回:
    - LoginWithTenantsSchema: 包含令牌、租户列表和用户信息的响应模型

    异常:
    - CustomException: 认证失败时抛出异常。
    """
    login_result = await LoginService.authenticate_user_service(
        request=request, redis=redis, login_form=login_form, db=db
    )

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

    # 如果是文档请求，则不记录日志
    if settings.DOCS_URL in request.headers.get("referer", ""):
        return login_result
    return SuccessResponse(data=login_result, msg="登录成功")


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

    参数:
    - payload (RefreshTokenPayloadSchema): 刷新令牌负载模型
    - db (AsyncSession): 数据库会话对象
    - redis (Redis): Redis 客户端对象

    返回:
    - JWTOutSchema: 包含新的访问令牌和刷新令牌的响应模型

    异常:
    - CustomException: 刷新令牌失败时抛出异常。
    """
    new_token = await LoginService.refresh_token_service(
        db=db, redis=redis, refresh_token=payload
    )
    return SuccessResponse(data=new_token, msg="刷新成功")


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

    参数:
    - redis (Redis): Redis客户端对象

    返回:
    - CaptchaOutSchema: 包含验证码图片和key的响应模型

    异常:
    - CustomException: 获取验证码失败时抛出异常。
    """
    # 获取验证码
    captcha = await CaptchaService.get_captcha_service(redis=redis)
    return SuccessResponse(data=captcha, msg="获取验证码成功")


@AuthRouter.post(
    "/logout",
    summary="退出登录",
    dependencies=[Depends(get_current_user)],
    response_model=ResponseSchema[None],
)
async def logout_controller(
    payload: LogoutPayloadSchema,
    redis: Annotated[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="退出失败")


@AuthRouter.get(
    "/auto-login/users",
    summary="获取免登录用户列表",
    response_model=ResponseSchema[list[AutoLoginUserSchema]],
)
async def get_auto_login_users_controller(
    auth: Annotated[AuthSchema, Depends(get_current_user)],
    db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse:
    """
    获取免登录用户列表

    参数:
    - auth (AuthSchema): 认证信息
    - db (AsyncSession): 数据库会话对象

    返回:
    - list[AutoLoginUserSchema]: 免登录用户列表
    """
    tenant_id = None if auth.user.is_superuser else auth.user.tenant_id
    users = await AutoLoginService.get_auto_login_users_service(db=db, tenant_id=tenant_id)
    return SuccessResponse(data=users, msg="获取成功")


@AuthRouter.post(
    "/auto-login/token",
    summary="获取免登录Token",
    response_model=ResponseSchema[AutoLoginTokenSchema],
)
async def get_auto_login_token_controller(
    auth: Annotated[AuthSchema, Depends(get_current_user)],
    redis: Annotated[Redis, Depends(redis_getter)],
    db: Annotated[AsyncSession, Depends(db_getter)],
    user_id: int,
) -> JSONResponse:
    """
    获取免登录Token

    参数:
    - auth (AuthSchema): 认证信息
    - redis (Redis): Redis客户端对象
    - db (AsyncSession): 数据库会话对象
    - user_id (int): 用户ID

    返回:
    - AutoLoginTokenSchema: 免登录Token和用户信息
    """
    tenant_id = None if auth.user.is_superuser else auth.user.tenant_id
    result = await AutoLoginService.create_auto_login_token_service(
        redis=redis, db=db, user_id=user_id, tenant_id=tenant_id
    )
    return SuccessResponse(data=result, msg="获取成功")


@AuthRouter.post(
    "/auto-login",
    summary="免登录",
    response_model=ResponseSchema[JWTOutSchema],
)
async def auto_login_controller(
    request: Request,
    redis: Annotated[Redis, Depends(redis_getter)],
    db: Annotated[AsyncSession, Depends(db_getter)],
    token: str,
) -> JSONResponse:
    """
    免登录

    参数:
    - request (Request): FastAPI请求对象
    - redis (Redis): Redis客户端对象
    - db (AsyncSession): 数据库会话对象
    - token (str): 免登录Token

    返回:
    - JWTOutSchema: JWT令牌信息
    """
    login_token = await AutoLoginService.auto_login_service(
        request=request, redis=redis, db=db, token=token
    )
    logger.info("用户免登录成功")
    return SuccessResponse(data=login_token, msg="登录成功")


@AuthRouter.post(
    "/select-tenant",
    summary="选择租户",
    response_model=ResponseSchema[SelectTenantOutSchema],
    dependencies=[Depends(get_current_user)],
)
async def select_tenant_controller(
    request: Request,
    data: SelectTenantSchema,
    auth: Annotated[AuthSchema, Depends(get_current_user)],
    redis: Annotated[Redis, Depends(redis_getter)],
) -> JSONResponse:
    """
    选择租户

    验证用户是否属于该租户，签发包含 tenant_id 的新 JWT Token。

    参数:
    - request (Request): FastAPI请求对象
    - data (SelectTenantSchema): 租户选择请求
    - auth (AuthSchema): 当前认证信息
    - redis (Redis): Redis客户端对象

    返回:
    - SelectTenantOutSchema: 包含新令牌的响应
    """
    result = await LoginService.select_tenant_service(
        request=request, redis=redis, auth=auth, tenant_id=data.tenant_id
    )
    await FastAPICache.clear(namespace=_AUTH_TENANTS_NS)
    return SuccessResponse(data=result, msg="租户切换成功")


@AuthRouter.get(
    "/tenants",
    summary="获取可选租户列表",
    response_model=ResponseSchema[list[TenantOptionSchema]],
    dependencies=[Depends(get_current_user)],
)
@cache(expire=120, namespace=_AUTH_TENANTS_NS)
async def get_user_tenants_controller(
    auth: Annotated[AuthSchema, Depends(get_current_user)],
    db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse:
    """
    获取当前用户的租户列表

    参数:
    - auth (AuthSchema): 当前认证信息
    - db (AsyncSession): 数据库会话对象

    返回:
    - list[TenantOptionSchema]: 租户选项列表
    """
    tenants = await LoginService.get_user_tenants_service(auth=auth, db=db)
    return SuccessResponse(data=tenants, msg="获取租户列表成功")


@AuthRouter.get(
    "/oauth/{provider}/login",
    summary="第三方OAuth跳转",
)
async def oauth_login_redirect_controller(
    request: Request,
    redis: Annotated[Redis, Depends(redis_getter)],
    provider: Annotated[str, Path(description="wechat | qq | github | gitee")],
    redirect_uri: Annotated[
        str | None,
        Query(description="OAuth 完成后浏览器回到的前端登录页完整 URL"),
    ] = None,
) -> RedirectResponse:
    allowed = {"wechat", "qq", "github", "gitee"}
    fe = redirect_uri or settings.OAUTH_FRONTEND_FALLBACK
    if provider not in allowed:
        return RedirectResponse(
            url=oauth_service_error_redirect(fe, "不支持的 OAuth 渠道"),
            status_code=302,
        )
    if not redirect_uri:
        return RedirectResponse(
            url=oauth_service_error_redirect(fe, "缺少 redirect_uri 参数"),
            status_code=302,
        )
    try:
        state = secrets.token_urlsafe(32)
        await save_oauth_state(
            redis=redis,
            state=state,
            provider=provider,
            frontend_redirect=redirect_uri,
        )
        cb = _callback_url(request, provider)
        url = build_authorize_url(provider=provider, callback_url=cb, state=state)
        return RedirectResponse(url=url, status_code=302)
    except CustomException as e:
        return RedirectResponse(
            url=oauth_service_error_redirect(redirect_uri, e.msg),
            status_code=302,
        )


@AuthRouter.get(
    "/oauth/{provider}/callback",
    summary="第三方OAuth回调",
    include_in_schema=False,
)
async def oauth_callback_controller(
    request: Request,
    redis: Annotated[Redis, Depends(redis_getter)],
    db: Annotated[AsyncSession, Depends(db_getter)],
    provider: Annotated[str, Path()],
    code: Annotated[str | None, Query()] = None,
    state: Annotated[str | None, Query()] = None,
) -> RedirectResponse:
    fe_fallback = settings.OAUTH_FRONTEND_FALLBACK

    async def resolve_frontend() -> str:
        if not state:
            return fe_fallback
        raw = await RedisCURD(redis).get(f"{STATE_PREFIX}{state}")
        if not raw:
            return fe_fallback
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")
        try:
            payload = json.loads(raw)
            return str(payload.get("frontend_redirect") or fe_fallback).strip() or fe_fallback
        except json.JSONDecodeError:
            return fe_fallback

    if provider not in {"wechat", "qq", "github", "gitee"}:
        url = oauth_service_error_redirect(await resolve_frontend(), "不支持的 OAuth 渠道")
        return RedirectResponse(url=url, status_code=302)
    if not code or not state:
        url = oauth_service_error_redirect(await resolve_frontend(), "授权被取消或参数不完整")
        return RedirectResponse(url=url, status_code=302)
    try:
        token, fe = await complete_oauth_login(
            request=request,
            redis=redis,
            db=db,
            provider=provider,
            code=code,
            state=state,
        )
        success_url = oauth_service_frontend_redirect_from_token(fe, token)
        return RedirectResponse(url=success_url, status_code=302)
    except CustomException as e:
        fe = await resolve_frontend()
        return RedirectResponse(url=oauth_service_error_redirect(fe, e.msg), status_code=302)


@AuthRouter.post(
    "/tenant/register",
    summary="租户自助注册",
    response_model=ResponseSchema[TenantRegisterOutSchema],
)
async def tenant_register_controller(
    data: TenantRegisterSchema,
    db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse:
    """
    租户自助注册（PRD §4.5）

    参数:
    - data (TenantRegisterSchema): 注册信息

    返回:
    - TenantRegisterOutSchema: 注册结果，含 tenant_id/user_id/试用到期日
    """
    result = await TenantRegisterService.register(
        db=db,
        username=data.username,
        password=data.password,
        email=data.email,
        tenant_name=data.tenant_name,
    )
    logger.info(f"新租户注册: username={data.username} tenant={result.tenant_name}")
    return SuccessResponse(data=result, msg=result.message)


# ─── 忘记密码（自助重置）────────────────────────────────────

@AuthRouter.post(
    "/forgot-password",
    summary="忘记密码",
    response_model=ResponseSchema[str],
)
async def forgot_password_controller(
    data: ForgotPasswordSchema,
    redis: Annotated[Redis, Depends(redis_getter)],
    db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse:
    """
    忘记密码：提交邮箱，接收重置邮件

    参数:
    - data (ForgotPasswordSchema): 邮箱

    返回:
    - str: 提示信息（无论成功与否均返回相同文案）
    """
    msg = await PasswordResetService.forgot_password_service(
        redis=redis, db=db, email=data.email
    )
    return SuccessResponse(data=msg, msg="邮件已发送")


@AuthRouter.post(
    "/reset-password",
    summary="重置密码",
    response_model=ResponseSchema[str],
)
async def reset_password_controller(
    data: ResetPasswordWithTokenSchema,
    redis: Annotated[Redis, Depends(redis_getter)],
    db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse:
    """
    重置密码：使用令牌 + 新密码

    参数:
    - data (ResetPasswordWithTokenSchema): 令牌 + 新密码

    返回:
    - str: 重置结果
    """
    msg = await PasswordResetService.reset_password_with_token_service(
        redis=redis, db=db, token=data.token, new_password=data.new_password
    )
    return SuccessResponse(data=msg, msg="密码已重置")