95 lines
3.4 KiB
Python
95 lines
3.4 KiB
Python
# -*- coding = utf-8 -*-
|
||
# @Time : 2025/2/17 下午8:14
|
||
# @File : services.py
|
||
# @Software : PyCharm
|
||
# @Author : xingc
|
||
# @Desc :
|
||
from datetime import datetime, timedelta
|
||
from typing import Annotated, Optional
|
||
|
||
import jwt
|
||
import ujson
|
||
from fastapi import HTTPException, status, Depends, Request, WebSocket
|
||
from fastapi.security import OAuth2PasswordBearer
|
||
import redis.asyncio as aioredis
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from config import settings
|
||
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 . import schema
|
||
|
||
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
||
|
||
|
||
def create_access_token(data: dict):
|
||
to_encode = data.copy()
|
||
expire = datetime.utcnow() + timedelta(
|
||
minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES
|
||
)
|
||
to_encode.update({"exp": expire})
|
||
encoded_jwt = jwt.encode(to_encode, settings.ACCESS_SECRET_KEY, algorithm=settings.ALGORITHM)
|
||
return {'token_type': 'bearer', 'access_token': encoded_jwt}
|
||
|
||
|
||
def verify_access_token(token: str):
|
||
try:
|
||
# payload = jwt.decode(token, settings.ACCESS_SECRET_KEY, algorithms=[settings.ALGORITHM])
|
||
payload = jwt.decode(token, settings.ACCESS_SECRET_KEY, algorithms=[settings.ALGORITHM])
|
||
return payload
|
||
except jwt.ExpiredSignatureError:
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Token has expired")
|
||
except jwt.InvalidTokenError:
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
|
||
|
||
|
||
async def get_current_user(
|
||
token: str | None, db_session: AsyncSession, redis_conn: aioredis
|
||
) -> schema.LoginUser:
|
||
"""验证登录状态、获取的用户基础信息"""
|
||
|
||
|
||
result = verify_access_token(token)
|
||
user_id = result.get("sub")
|
||
if user_id is None:
|
||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
|
||
|
||
keyname = generate_keyname(f'login_cache', use_today=True)
|
||
user_str = await redis_conn.hget(keyname, user_id)
|
||
if user_str is None:
|
||
user_model = await CRUD(db_session, User).get(filters={'user_id': user_id})
|
||
user = schema.LoginUser(**user_model.dict())
|
||
await redis_conn.hset(keyname, user_id, user.json())
|
||
else:
|
||
user = schema.LoginUser(**ujson.decode(user_str))
|
||
|
||
return user
|
||
|
||
|
||
async def get_current_user_form_http(
|
||
token: Annotated[str, Depends(oauth2_scheme)],
|
||
db_session: Annotated[AsyncSession, Depends(get_async_session)],
|
||
redis_conn: aioredis = Depends(get_cache),
|
||
) -> schema.LoginUser:
|
||
"""在http headers中提取token，并验证再获取当前用户"""
|
||
user = await get_current_user(token, db_session, redis_conn)
|
||
return user
|
||
|
||
|
||
async def get_current_user_form_ws(
|
||
websocket: WebSocket,
|
||
db_session: Annotated[AsyncSession, Depends(get_async_session)] = None,
|
||
redis_conn: aioredis = Depends(get_cache),
|
||
) -> schema.LoginUser:
|
||
"""在ws headers中提取token，并验证再获取当前用户"""
|
||
auth_header = websocket.headers.get("authorization") or websocket.headers.get("Authorization")
|
||
if auth_header and auth_header.lower().startswith("bearer "):
|
||
token = auth_header.removeprefix("Bearer ").strip()
|
||
else:
|
||
token = ''
|
||
user = await get_current_user(token, db_session, redis_conn)
|
||
return user
|