python-fastapi-mm-zhcs-yj-api/common/security.py

# -*- coding: utf-8 -*-
from fastapi import Header
from datetime import datetime, timedelta
import jwt
from passlib.context import CryptContext
from sqlalchemy.orm import Session
from models.base import AppInfo
from exceptions import TokenException
from config import settings
from extensions import logger

# https://fastapi.tiangolo.com/tutorial/request-forms/
# https://fastapi.tiangolo.com/tutorial/security/oauth2-jwt/

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

def valid_access_token(Authorization: str = Header(..., alias="Authorization")) -> int:
    access_token = Authorization.removeprefix("Bearer ") 
    
    token_exception = TokenException()

    try:
        payload = jwt.decode(access_token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        user_id: str = payload.get("sub")
        logger.info('sub user_id: {}', user_id)
    except Exception:
        raise token_exception
    
    return int(user_id)

def verify_secret(plain_secret, hashed_secret):
    return pwd_context.verify(plain_secret, hashed_secret)

def get_secret_hash(secret):
    return pwd_context.hash(secret)

def create_access_token(*, data: dict, expires_delta: timedelta = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=10)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
    return encoded_jwt