130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
# -*- coding = utf-8 -*-
|
||
# @Time : 2025/2/13 下午5:02
|
||
# @File : routes.py
|
||
# @Software : PyCharm
|
||
# @Author : xingc
|
||
# @Desc :
|
||
import logging
|
||
from typing import Annotated
|
||
|
||
import ujson
|
||
from fastapi import APIRouter, HTTPException, Depends, Query, status
|
||
import redis.asyncio as aioredis
|
||
from fastapi.security import OAuth2PasswordRequestForm
|
||
from simple_spider_tool import md5, jsonpath
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from config import settings
|
||
from core.base import schema as base_schema
|
||
from core.base.exceptions import WechatApiError
|
||
from core.db.crud import CRUD
|
||
from core.db.engine import get_async_session, get_cache
|
||
from core.db.models import User
|
||
from core.utils.helper import generate_keyname
|
||
from core.wechat.api import wechat_api, decrypt_sensitive_data
|
||
|
||
from . import schema
|
||
from .services import create_access_token
|
||
|
||
logger = logging.getLogger(__name__)
|
||
auth_router = APIRouter(prefix='/auth', tags=['认证模块'])
|
||
test_auth_router = APIRouter(tags=['认证模块'])
|
||
|
||
|
||
@auth_router.get(
|
||
'/phone',
|
||
summary='获取绑定的手机号',
|
||
response_model=base_schema.BaseResponse
|
||
)
|
||
async def get_wechat_phone(
|
||
code: str = Query(pattern=r'^[a-zA-Z0-9]+'),
|
||
):
|
||
try:
|
||
result = await wechat_api.get_phone_number(code)
|
||
except WechatApiError:
|
||
raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail='Invalid code')
|
||
return {'data': {'phone': jsonpath(result, '$.phone_info.phoneNumber', first=True)}}
|
||
|
||
|
||
@auth_router.post(
|
||
'/wechat_login',
|
||
summary='微信小程序授权登录',
|
||
response_model=base_schema.BaseResponse[schema.Token]
|
||
)
|
||
async def wechat_login(
|
||
payload: schema.WechatLoginIn,
|
||
db_session: Annotated[AsyncSession, Depends(get_async_session)],
|
||
redis_conn: Annotated[aioredis.Redis, Depends(get_cache)]
|
||
):
|
||
db_crud = CRUD(db_session, User)
|
||
keyname = generate_keyname(f'wechat_session:{md5(payload.code)}')
|
||
# 验证code的有效性
|
||
wechat_session_str = await redis_conn.get(keyname)
|
||
if wechat_session_str is None:
|
||
try:
|
||
wechat_session = await wechat_api.code_to_session(payload.code)
|
||
except WechatApiError:
|
||
raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail='Invalid code')
|
||
logger.debug(f"WeChat session: {wechat_session}")
|
||
else:
|
||
wechat_session = ujson.decode(wechat_session_str)
|
||
|
||
openid = wechat_session.get('openid')
|
||
if openid is None:
|
||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='Invalid code')
|
||
|
||
# 验证用户的存在/是否位首次登录
|
||
user: User | None = await db_crud.get(filters={'wechat_openid': openid})
|
||
if user is None:
|
||
checked = all([
|
||
payload.phone,
|
||
payload.nickname
|
||
])
|
||
if checked is False and wechat_session_str is None:
|
||
await redis_conn.set(keyname, ujson.encode(wechat_session), ex=60 * 5)
|
||
return {'status': 4001, 'message': '首次登录需要用户授权信息'}
|
||
|
||
if checked is False:
|
||
raise HTTPException(status.HTTP_400_BAD_REQUEST, detail='请重试登录操作')
|
||
|
||
# 注册用户
|
||
union_id = wechat_session.get('union_id')
|
||
user: User | bool = await db_crud.create(
|
||
data={
|
||
'nickname': payload.nickname,
|
||
'wechat_openid': openid,
|
||
'wechat_union_id': union_id,
|
||
'phone': payload.phone,
|
||
'avatar_url': payload.avatar_url,
|
||
},
|
||
filters={'wechat_openid': openid}
|
||
)
|
||
await redis_conn.delete(keyname)
|
||
|
||
result = create_access_token({
|
||
'sub': str(user.user_id)
|
||
})
|
||
return {'data': result}
|
||
|
||
|
||
if settings.ENVIRONMENT == 'dev':
|
||
@test_auth_router.post(
|
||
'/token',
|
||
summary='开发',
|
||
description='点击“Authorize”在username输入“user_id”, password任意输入字符，点再点击“Authorize”即可'
|
||
)
|
||
async def login_for_access_token(
|
||
form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
|
||
db_session: Annotated[AsyncSession, Depends(get_async_session)]
|
||
) -> schema.Token:
|
||
logger.info(form_data.username, form_data.password)
|
||
db_crud = CRUD(db_session, User)
|
||
user = await db_crud.get(filters={'user_id': form_data.username})
|
||
if not user:
|
||
raise HTTPException(status_code=400, detail='Incorrect user_id')
|
||
|
||
result = create_access_token({
|
||
'sub': str(user.user_id)
|
||
})
|
||
return result
|