fastapi_best_architecture/backend/app/common/jwt.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from datetime import datetime, timedelta

from fastapi import Depends, Request
from fastapi.security import OAuth2PasswordBearer
from fastapi.security.utils import get_authorization_scheme_param
from jose import jwt
from passlib.context import CryptContext
from pydantic import ValidationError
from typing_extensions import Annotated

from backend.app.common.exception.errors import AuthorizationError, TokenError
from backend.app.common.redis import redis_client
from backend.app.core.conf import settings
from backend.app.crud.crud_user import UserDao
from backend.app.database.db_mysql import CurrentSession
from backend.app.models import User
from backend.app.schemas.token import RefreshTokenTime

pwd_context = CryptContext(schemes=['bcrypt'], deprecated='auto')

oauth2_schema = OAuth2PasswordBearer(tokenUrl=settings.TOKEN_URL_SWAGGER)


def get_hash_password(password: str) -> str:
    """
    Encrypt passwords using the hash algorithm

    :param password:
    :return:
    """
    return pwd_context.hash(password)


def password_verify(plain_password: str, hashed_password: str) -> bool:
    """
    Password verification

    :param plain_password: The password to verify
    :param hashed_password: The hash ciphers to compare
    :return:
    """
    return pwd_context.verify(plain_password, hashed_password)


async def create_access_token(sub: str, expires_delta: timedelta | None = None, **kwargs) -> tuple[str, datetime]:
    """
    Generate encryption token

    :param sub: The subject/userid of the JWT
    :param expires_delta: Increased expiry time
    :return:
    """
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
        expire_seconds = int(expires_delta.total_seconds())
    else:
        expire = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
        expire_seconds = settings.TOKEN_EXPIRE_SECONDS
    to_encode = {'exp': expire, 'sub': sub, **kwargs}
    token = jwt.encode(to_encode, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
    if sub not in settings.TOKEN_WHITE_LIST:
        await redis_client.delete_prefix(f'{settings.TOKEN_REDIS_PREFIX}:{sub}:')
    key = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:{token}'
    await redis_client.setex(key, expire_seconds, token)
    return token, expire


async def create_refresh_token(
    sub: str, expire_time: datetime | None = None, custom_expire_time: RefreshTokenTime | None = None, **kwargs
) -> tuple[str, datetime]:
    """
    Generate encryption refresh token

    :param sub: The subject/userid of the JWT
    :param expire_time: expiry time
    :param custom_expire_time: custom expiry time
    :return:
    """
    if expire_time:
        expire = expire_time + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
        expire_seconds = int((expire - datetime.utcnow()).total_seconds())
    elif custom_expire_time:
        expire = custom_expire_time.expire_time
        expire_seconds = int((expire - datetime.utcnow()).total_seconds())
    else:
        expire = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
        expire_seconds = settings.TOKEN_EXPIRE_SECONDS
    to_encode = {'exp': expire, 'sub': sub, **kwargs}
    token = jwt.encode(to_encode, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
    # 刷新 token 时，保持旧 token 有效，不执行删除操作
    key = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:{token}'
    await redis_client.setex(key, expire_seconds, token)
    return token, expire


def get_token(request: Request) -> str:
    """
    Get token for request header

    :return:
    """
    authorization = request.headers.get('Authorization')
    scheme, param = get_authorization_scheme_param(authorization)
    if not authorization or scheme.lower() != 'bearer':
        raise TokenError
    return param


def jwt_decode(token: str) -> tuple[int, list[int]]:
    """
    Decode token

    :param token:
    :return:
    """
    try:
        payload = jwt.decode(token, settings.TOKEN_SECRET_KEY, algorithms=[settings.TOKEN_ALGORITHM])
        user_id = int(payload.get('sub'))
        user_roles = list(payload.get('role_ids'))
        if not user_id or not user_roles:
            raise TokenError
    except (jwt.JWTError, ValidationError, Exception):
        raise TokenError
    return user_id, user_roles


async def jwt_authentication(token: str = Depends(oauth2_schema)) -> dict[str, int]:
    """
    JWT authentication

    :param token:
    :return:
    """
    user_id, _ = jwt_decode(token)
    key = f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{token}'
    token_verify = await redis_client.get(key)
    if not token_verify:
        raise TokenError
    return {'sub': user_id}


async def get_current_user(db: CurrentSession, data: dict = Depends(jwt_authentication)) -> User:
    """
    Get the current user through tokens

    :param db:
    :param data:
    :return:
    """
    user_id = data.get('sub')
    user = await UserDao.get_with_relation(db, user_id=user_id)
    if not user:
        raise TokenError
    return user


async def get_current_is_superuser(user: User = Depends(get_current_user)):
    """
    Verify the current user permissions through token

    :param user:
    :return:
    """
    is_superuser = user.is_superuser
    if not is_superuser:
        raise AuthorizationError
    return is_superuser


# User Annotated
CurrentUser = Annotated[User, Depends(get_current_user)]
CurrentSuperUser = Annotated[bool, Depends(get_current_is_superuser)]
# Token dependency injection
CurrentJwtAuth = Annotated[dict, Depends(jwt_authentication)]
# Permission dependency injection
DependsUser = Depends(get_current_user)
DependsSuperUser = Depends(get_current_is_superuser)