102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
|
import jwt
|
|||
|
|
import redis
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
from typing import Optional, Dict, Any
|
|||
|
|
from app.core.config import REDIS_CONFIG
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
# JWT配置
|
|||
|
|
JWT_SECRET_KEY = os.getenv('JWT_SECRET_KEY', 'your-super-secret-key-change-in-production')
|
|||
|
|
JWT_ALGORITHM = "HS256"
|
|||
|
|
ACCESS_TOKEN_EXPIRE_MINUTES = 60 * 24 * 7 # 7天
|
|||
|
|
|
|||
|
|
class JWTService:
|
|||
|
|
def __init__(self):
|
|||
|
|
self.redis_client = redis.Redis(**REDIS_CONFIG)
|
|||
|
|
|
|||
|
|
def create_access_token(self, data: Dict[str, Any]) -> str:
|
|||
|
|
"""创建JWT访问令牌"""
|
|||
|
|
to_encode = data.copy()
|
|||
|
|
expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
|
|||
|
|
to_encode.update({"exp": expire, "type": "access"})
|
|||
|
|
|
|||
|
|
encoded_jwt = jwt.encode(to_encode, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
|||
|
|
|
|||
|
|
# 将token存储到Redis，用于管理和撤销
|
|||
|
|
user_id = data.get("user_id")
|
|||
|
|
if user_id:
|
|||
|
|
self.redis_client.setex(
|
|||
|
|
f"token:{user_id}:{encoded_jwt}",
|
|||
|
|
ACCESS_TOKEN_EXPIRE_MINUTES * 60, # Redis需要秒
|
|||
|
|
"active"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return encoded_jwt
|
|||
|
|
|
|||
|
|
def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
|
|||
|
|
"""验证JWT令牌"""
|
|||
|
|
try:
|
|||
|
|
# 解码JWT
|
|||
|
|
payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
|
|||
|
|
|
|||
|
|
# 检查token类型
|
|||
|
|
if payload.get("type") != "access":
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
user_id = payload.get("user_id")
|
|||
|
|
if not user_id:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 检查token是否在Redis中且未被撤销
|
|||
|
|
redis_key = f"token:{user_id}:{token}"
|
|||
|
|
if not self.redis_client.exists(redis_key):
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
return payload
|
|||
|
|
|
|||
|
|
except jwt.ExpiredSignatureError:
|
|||
|
|
return None
|
|||
|
|
except jwt.InvalidTokenError:
|
|||
|
|
return None
|
|||
|
|
except Exception:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def revoke_token(self, token: str, user_id: int) -> bool:
|
|||
|
|
"""撤销token"""
|
|||
|
|
try:
|
|||
|
|
redis_key = f"token:{user_id}:{token}"
|
|||
|
|
return self.redis_client.delete(redis_key) > 0
|
|||
|
|
except:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def revoke_all_user_tokens(self, user_id: int) -> int:
|
|||
|
|
"""撤销用户的所有token"""
|
|||
|
|
try:
|
|||
|
|
pattern = f"token:{user_id}:*"
|
|||
|
|
keys = self.redis_client.keys(pattern)
|
|||
|
|
if keys:
|
|||
|
|
return self.redis_client.delete(*keys)
|
|||
|
|
return 0
|
|||
|
|
except:
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
def refresh_token(self, token: str) -> Optional[str]:
|
|||
|
|
"""刷新token（可选功能）"""
|
|||
|
|
payload = self.verify_token(token)
|
|||
|
|
if not payload:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 撤销旧token
|
|||
|
|
user_id = payload.get("user_id")
|
|||
|
|
self.revoke_token(token, user_id)
|
|||
|
|
|
|||
|
|
# 创建新token
|
|||
|
|
new_data = {
|
|||
|
|
"user_id": user_id,
|
|||
|
|
"username": payload.get("username"),
|
|||
|
|
"caption": payload.get("caption")
|
|||
|
|
}
|
|||
|
|
return self.create_access_token(new_data)
|
|||
|
|
|
|||
|
|
# 全局实例
|
|||
|
|
jwt_service = JWTService()
|