fastapi-best-architecture/backend/plugin/oauth2/service/oauth2_service.py

import json

from typing import Any

from fast_captcha import text_captcha
from fastapi import BackgroundTasks, Response
from sqlalchemy.ext.asyncio import AsyncSession

from backend.app.admin.crud.crud_user import user_dao
from backend.app.admin.schema.token import GetLoginToken
from backend.app.admin.schema.user import AddOAuth2UserParam
from backend.app.admin.service.login_log_service import login_log_service
from backend.common.context import ctx
from backend.common.enums import LoginLogStatusType
from backend.common.exception import errors
from backend.common.i18n import t
from backend.common.security import jwt
from backend.core.conf import settings
from backend.database.redis import redis_client
from backend.plugin.oauth2.crud.crud_user_social import user_social_dao
from backend.plugin.oauth2.enums import UserSocialAuthType, UserSocialType
from backend.plugin.oauth2.schema.user_social import CreateUserSocialParam
from backend.plugin.oauth2.service.user_social_service import user_social_service
from backend.utils.timezone import timezone


class OAuth2Service:
    """OAuth2 认证服务类"""

    @staticmethod
    async def login(
        *,
        db: AsyncSession,
        response: Response,
        background_tasks: BackgroundTasks,
        sid: str,
        source: UserSocialType,
        username: str | None = None,
        nickname: str | None = None,
        email: str | None = None,
        avatar: str | None = None,
    ) -> GetLoginToken:
        """
        OAuth2 用户登录

        :param db: 数据库会话
        :param response: FastAPI 响应对象
        :param background_tasks: FastAPI 后台任务
        :param sid: 社交账号唯一编码
        :param source: 社交平台
        :param username: 用户名
        :param nickname: 昵称
        :param email: 邮箱
        :param avatar: 头像地址
        :return:
        """
        user_social = await user_social_dao.get_by_sid(db, sid, source.value)
        if user_social:
            sys_user = await user_dao.get(db, user_social.user_id)
            # 更新用户头像
            if not sys_user.avatar and avatar is not None:
                await user_dao.update_avatar(db, sys_user.id, avatar)
        else:
            sys_user = None
            # 检测系统用户是否已存在
            if email:
                sys_user = await user_dao.check_email(db, email)

            # 创建系统用户
            if not sys_user:
                while await user_dao.get_by_username(db, username):
                    username = f'{username}_{text_captcha(5)}'
                new_sys_user = AddOAuth2UserParam(
                    username=username,
                    password=None,
                    nickname=nickname,
                    email=email,
                    avatar=avatar,
                )
                await user_dao.add_by_oauth2(db, new_sys_user)
                await db.flush()
                sys_user = await user_dao.get_by_username(db, username)

            # 绑定社交账号
            new_user_social = CreateUserSocialParam(sid=sid, source=source.value, user_id=sys_user.id)
            await user_social_dao.create(db, new_user_social)

        # 创建 token
        access_token_data = await jwt.create_access_token(
            sys_user.id,
            multi_login=sys_user.is_multi_login,
            # extra info
            username=sys_user.username,
            nickname=sys_user.nickname,
            last_login_time=timezone.to_str(timezone.now()),
            ip=ctx.ip,
            os=ctx.os,
            browser=ctx.browser,
            device=ctx.device,
        )
        refresh_token_data = await jwt.create_refresh_token(
            access_token_data.session_uuid,
            sys_user.id,
            multi_login=sys_user.is_multi_login,
        )
        await user_dao.update_login_time(db, sys_user.username)
        await db.refresh(sys_user)
        background_tasks.add_task(
            login_log_service.create,
            user_uuid=sys_user.uuid,
            username=sys_user.username,
            login_time=timezone.now(),
            status=LoginLogStatusType.success.value,
            msg=t('success.login.oauth2_success'),
        )
        await redis_client.delete(f'{settings.LOGIN_CAPTCHA_REDIS_PREFIX}:{ctx.ip}')
        response.set_cookie(
            key=settings.COOKIE_REFRESH_TOKEN_KEY,
            value=refresh_token_data.refresh_token,
            max_age=settings.COOKIE_REFRESH_TOKEN_EXPIRE_SECONDS,
            expires=timezone.to_utc(refresh_token_data.refresh_token_expire_time),
            httponly=True,
        )
        data = GetLoginToken(
            access_token=access_token_data.access_token,
            access_token_expire_time=access_token_data.access_token_expire_time,
            session_uuid=access_token_data.session_uuid,
            user=sys_user,  # type: ignore
        )
        return data

    async def login_or_binding(
        self,
        *,
        db: AsyncSession,
        response: Response,
        background_tasks: BackgroundTasks,
        user: dict[str, Any],
        social: UserSocialType,
        state: str | None = None,
    ) -> GetLoginToken | None:
        """
        OAuth2 登录或绑定

        :param db: 数据库会话
        :param response: FastAPI 响应对象
        :param background_tasks: FastAPI 后台任务
        :param user: OAuth2 用户信息
        :param social: 社交平台类型
        :param state: OAuth2 state 参数
        :return:
        """

        sid = user.get('uuid')
        username = user.get('username')
        nickname = user.get('nickname')
        email = user.get('email')
        avatar = user.get('avatar_url')

        match social:
            case UserSocialType.github:
                sid = user.get('id')
                username = user.get('login')
                nickname = user.get('name')
            case UserSocialType.google:
                sid = user.get('id')
                username = user.get('name')
                nickname = user.get('given_name')
                avatar = user.get('picture')
            case _:
                raise errors.ForbiddenError(msg=f'暂不支持 {social} OAuth2 登录')

        if not state:
            raise errors.ForbiddenError(msg='OAuth2 状态信息缺失')

        state_data = await redis_client.get(f'{settings.OAUTH2_STATE_REDIS_PREFIX}:{state}')
        if not state_data:
            raise errors.ForbiddenError(msg='OAuth2 状态信息无效或缺失')

        state_info = json.loads(state_data)
        await redis_client.delete(f'{settings.OAUTH2_STATE_REDIS_PREFIX}:{state}')

        # 绑定流程
        if state_info.get('type') == UserSocialAuthType.binding.value:
            user_id = state_info.get('user_id')
            if not user_id:
                raise errors.ForbiddenError(msg='非法操作，OAuth2 状态信息无效')
            await user_social_service.binding_with_oauth2(
                db=db,
                user_id=user_id,
                sid=str(sid),
                source=social,
            )
            return None

        # 登录流程
        if state_info.get('type') != UserSocialAuthType.login.value:
            raise errors.ForbiddenError(msg='OAuth2 状态信息无效')

        return await self.login(
            db=db,
            response=response,
            background_tasks=background_tasks,
            sid=str(sid),
            source=social,
            username=username,
            nickname=nickname,
            email=email,
            avatar=avatar,
        )


oauth2_service: OAuth2Service = OAuth2Service()