You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

318 lines
10 KiB
Python

import json
import secrets
from typing import Annotated
from fastapi import APIRouter, BackgroundTasks, Depends, Path, Query, Request
from fastapi.responses import JSONResponse, RedirectResponse
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 import cache_util
from app.core.base_schema import (
AuthSchema,
JWTOutSchema,
LogoutPayloadSchema,
RefreshTokenPayloadSchema,
)
from app.core.cache_util import cache
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,
LoginWithTenantsSchema,
SelectTenantOutSchema,
SelectTenantSchema,
TenantOptionSchema,
TenantRegisterOutSchema,
TenantRegisterSchema,
)
from .service import (
AutoLoginService,
CaptchaService,
LoginService,
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,
background_tasks: BackgroundTasks,
redis: Annotated[Redis, Depends(redis_getter)],
login_form: Annotated[CustomOAuth2PasswordRequestForm, Depends()],
db: Annotated[AsyncSession, Depends(db_getter)],
) -> JSONResponse | dict:
login_result = await LoginService.authenticate_user(
request=request, redis=redis, login_form=login_form, db=db, background_tasks=background_tasks
)
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:
new_token = await LoginService.refresh_token(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:
captcha = await CaptchaService.get_captcha(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:
if await LoginService.logout(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:
tenant_id = None if auth.user.is_superuser else auth.user.tenant_id
users = await AutoLoginService.get_auto_login_users(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:
tenant_id = None if auth.user.is_superuser else auth.user.tenant_id
result = await AutoLoginService.create_auto_login_token(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:
login_token = await AutoLoginService.auto_login(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:
result = await LoginService(auth).select_tenant(request=request, redis=redis, tenant_id=data.tenant_id)
await cache_util.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:
service = LoginService(auth)
tenants = await service.get_user_tenants()
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:
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)