86 lines
2.2 KiB
Python
86 lines
2.2 KiB
Python
from contextvars import ContextVar
|
|
from datetime import datetime, timedelta
|
|
from typing import Optional
|
|
|
|
from fastapi import Depends
|
|
from fastapi.security import OAuth2PasswordBearer
|
|
import jwt
|
|
from passlib.context import CryptContext
|
|
|
|
from middleware.exceptions import get_credentials_exception
|
|
|
|
# 密钥和算法配置
|
|
SECRET_KEY = "sjdi@!#3ksj2780se1283"
|
|
ALGORITHM = "HS256"
|
|
ACCESS_TOKEN_EXPIRE_MINUTES = 30
|
|
|
|
# 密码哈希上下文
|
|
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
|
|
|
# OAuth2 方案
|
|
# 设置默认登录接口
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="session")
|
|
|
|
# 请求上下文
|
|
context_sub: ContextVar[str] = ContextVar('sub', default='')
|
|
|
|
|
|
# 获取数据库 加密密码
|
|
def get_password_hash(password: str):
|
|
return pwd_context.hash(password)
|
|
|
|
|
|
# 验证数据库密码
|
|
def verify_password(plain_password: str, hashed_password: str):
|
|
return pwd_context.verify(plain_password, hashed_password)
|
|
|
|
|
|
# 生成payload
|
|
def create_payload(sub: str) -> dict:
|
|
return {
|
|
"sub": sub
|
|
}
|
|
|
|
|
|
# 验证payload
|
|
def verify_payload(payload: dict) -> bool:
|
|
return "sub" in payload
|
|
|
|
|
|
# 从payload中获取用户
|
|
def get_payload_sub(payload: dict) -> str:
|
|
return payload["sub"]
|
|
|
|
|
|
# 创建token
|
|
def create_token(payload: dict, expires_delta: Optional[timedelta] = None):
|
|
# 复制一份
|
|
payload_copy = payload.copy()
|
|
|
|
# 加上有效时间
|
|
if expires_delta:
|
|
expire = datetime.now() + expires_delta
|
|
else:
|
|
expire = datetime.now() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
|
|
payload_copy.update({"exp": expire})
|
|
|
|
# 生成jwt Token
|
|
return jwt.encode(payload_copy, SECRET_KEY, algorithm=ALGORITHM)
|
|
|
|
|
|
# 验证token
|
|
async def verify_token(token: str = Depends(oauth2_scheme)):
|
|
try:
|
|
# 1. 检验token信息
|
|
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
|
# 2. 校验payload信息
|
|
if not verify_payload(payload):
|
|
raise get_credentials_exception()
|
|
# 3. 校验数据库中是否存在 payload中的账户信息
|
|
sub = get_payload_sub(payload)
|
|
|
|
# 4. 存储账户信息
|
|
context_sub.set(sub)
|
|
except jwt.exceptions.InvalidTokenError:
|
|
raise get_credentials_exception()
|