79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
|
"""
|
|||
|
|
JWT 工具函数
|
|||
|
|
"""
|
|||
|
|
import jwt
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
from typing import Optional
|
|||
|
|
from fastapi import HTTPException, Depends, Header
|
|||
|
|
from app.config import get_settings
|
|||
|
|
|
|||
|
|
|
|||
|
|
def create_token(user_id: int, openid: str) -> str:
|
|||
|
|
"""
|
|||
|
|
创建 JWT Token
|
|||
|
|
"""
|
|||
|
|
settings = get_settings()
|
|||
|
|
expire = datetime.utcnow() + timedelta(hours=settings.jwt_expire_hours)
|
|||
|
|
|
|||
|
|
payload = {
|
|||
|
|
"user_id": user_id,
|
|||
|
|
"openid": openid,
|
|||
|
|
"exp": expire,
|
|||
|
|
"iat": datetime.utcnow()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
token = jwt.encode(payload, settings.jwt_secret_key, algorithm="HS256")
|
|||
|
|
return token
|
|||
|
|
|
|||
|
|
|
|||
|
|
def verify_token(token: str) -> dict:
|
|||
|
|
"""
|
|||
|
|
验证 JWT Token
|
|||
|
|
返回 payload 或抛出异常
|
|||
|
|
"""
|
|||
|
|
settings = get_settings()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
payload = jwt.decode(token, settings.jwt_secret_key, algorithms=["HS256"])
|
|||
|
|
return payload
|
|||
|
|
except jwt.ExpiredSignatureError:
|
|||
|
|
raise HTTPException(status_code=401, detail="Token已过期，请重新登录")
|
|||
|
|
except jwt.InvalidTokenError:
|
|||
|
|
raise HTTPException(status_code=401, detail="无效的Token")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_current_user_id(authorization: Optional[str] = Header(None, alias="Authorization")) -> int:
|
|||
|
|
"""
|
|||
|
|
从 Header 中获取并验证 Token，返回 user_id
|
|||
|
|
用作 FastAPI 依赖注入
|
|||
|
|
"""
|
|||
|
|
if not authorization:
|
|||
|
|
raise HTTPException(status_code=401, detail="未提供身份令牌")
|
|||
|
|
|
|||
|
|
# 支持 "Bearer xxx" 格式
|
|||
|
|
token = authorization
|
|||
|
|
if authorization.startswith("Bearer "):
|
|||
|
|
token = authorization[7:]
|
|||
|
|
|
|||
|
|
payload = verify_token(token)
|
|||
|
|
return payload.get("user_id")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_optional_user_id(authorization: Optional[str] = Header(None, alias="Authorization")) -> Optional[int]:
|
|||
|
|
"""
|
|||
|
|
可选的用户验证，未提供 Token 时返回 None
|
|||
|
|
用于不强制要求登录的接口
|
|||
|
|
"""
|
|||
|
|
if not authorization:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
token = authorization
|
|||
|
|
if authorization.startswith("Bearer "):
|
|||
|
|
token = authorization[7:]
|
|||
|
|
|
|||
|
|
payload = verify_token(token)
|
|||
|
|
return payload.get("user_id")
|
|||
|
|
except HTTPException:
|
|||
|
|
return None
|