383 lines
12 KiB
Python
383 lines
12 KiB
Python
import json
|
||
from datetime import datetime, timedelta
|
||
from typing import Optional, Dict, Any, List
|
||
from pydantic import BaseModel
|
||
from dataclasses import dataclass, asdict
|
||
import hashlib
|
||
import threading
|
||
|
||
import jwt
|
||
|
||
|
||
@dataclass
|
||
class TokenCacheData:
|
||
"""Token 缓存数据结构"""
|
||
|
||
token_hash: str
|
||
user_id: str
|
||
device_id: Optional[str] = None
|
||
user_agent: Optional[str] = None
|
||
ip_address: Optional[str] = None
|
||
created_at: Optional[str] = None
|
||
expires_at: Optional[str] = None
|
||
last_used: Optional[str] = None
|
||
is_active: bool = True
|
||
payload_data: Optional[Dict[str, Any]] = None
|
||
token_type: str = "access" # access 或 refresh
|
||
|
||
def __post_init__(self):
|
||
if self.created_at is None:
|
||
self.created_at = datetime.utcnow().isoformat()
|
||
if self.last_used is None:
|
||
self.last_used = datetime.utcnow().isoformat()
|
||
|
||
|
||
class JWTCache:
|
||
"""JWT Token 缓存管理器"""
|
||
|
||
def __init__(self, config, cache):
|
||
self.config = config
|
||
|
||
self.lock = threading.RLock()
|
||
|
||
self.cache = cache
|
||
|
||
self.secret_key = self.config["security"]["secret_key"]
|
||
self.algorithm = self.config["security"]["algorithm"]
|
||
|
||
def _json_encode(self, obj):
|
||
"""自定义 JSON 编码器，处理 datetime 对象"""
|
||
|
||
def convert_datetime(o):
|
||
if isinstance(o, datetime):
|
||
return o.isoformat()
|
||
raise TypeError(
|
||
f"Object of type {o.__class__.__name__} is not JSON serializable"
|
||
)
|
||
|
||
return json.dumps(obj, default=convert_datetime)
|
||
|
||
def _json_decode(self, s):
|
||
"""自定义 JSON 解码器"""
|
||
return json.loads(s)
|
||
|
||
def _hash_token(self, token: str) -> str:
|
||
"""Token 哈希化（避免存储明文 Token）"""
|
||
return hashlib.sha256(token.encode()).hexdigest()
|
||
|
||
def create_token(
|
||
self,
|
||
user_id: str,
|
||
user_data: Dict[str, Any],
|
||
token_type: str = "access",
|
||
device_id: Optional[str] = None,
|
||
user_agent: Optional[str] = None,
|
||
ip_address: Optional[str] = None,
|
||
) -> str:
|
||
"""
|
||
创建并缓存 JWT Token
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
user_data: 用户数据
|
||
token_type: Token 类型（access 或 refresh）
|
||
device_id: 设备ID
|
||
user_agent: 用户代理
|
||
ip_address: IP地址
|
||
|
||
Returns:
|
||
JWT Token 字符串
|
||
"""
|
||
with self.lock:
|
||
# 计算过期时间
|
||
if token_type == "access":
|
||
expire_minutes = self.config["security"]["access_token_expire_minutes"]
|
||
expires_delta = timedelta(minutes=expire_minutes)
|
||
else:
|
||
expire_days = self.config["security"]["refresh_token_expire_days"]
|
||
expires_delta = timedelta(days=expire_days)
|
||
|
||
expire = datetime.utcnow() + expires_delta
|
||
|
||
# 构建 payload
|
||
payload = {
|
||
"sub": user_id,
|
||
"exp": expire,
|
||
"iat": datetime.utcnow(),
|
||
"type": token_type,
|
||
"data": user_data,
|
||
"device_id": device_id,
|
||
"jti": self._generate_jti(), # JWT ID，唯一标识
|
||
}
|
||
|
||
# 生成 JWT Token
|
||
token = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
|
||
if isinstance(token, bytes):
|
||
token = token.decode("utf-8")
|
||
|
||
# 创建缓存数据
|
||
token_hash = self._hash_token(token)
|
||
cache_data = TokenCacheData(
|
||
token_hash=token_hash,
|
||
user_id=user_id,
|
||
device_id=device_id,
|
||
user_agent=user_agent,
|
||
ip_address=ip_address,
|
||
expires_at=expire.isoformat(), # 转换为字符串
|
||
payload_data=payload,
|
||
token_type=token_type,
|
||
)
|
||
|
||
# 存储到缓存
|
||
self.cache.set(f"token:{token_hash}", asdict(cache_data))
|
||
|
||
# 更新用户会话索引
|
||
self._update_user_session_index(user_id, token_hash, token_type)
|
||
|
||
# 检查用户会话数限制
|
||
self._enforce_session_limit(user_id)
|
||
|
||
return token
|
||
|
||
def _generate_jti(self) -> str:
|
||
"""生成唯一的 JWT ID"""
|
||
import uuid
|
||
|
||
return str(uuid.uuid4())
|
||
|
||
def _update_user_session_index(
|
||
self, user_id: str, token_hash: str, token_type: str
|
||
):
|
||
"""更新用户会话索引"""
|
||
with self.lock:
|
||
key = f"user:{user_id}:{token_type}"
|
||
user_tokens = self.cache.get(key, [])
|
||
|
||
# 添加新 token
|
||
user_tokens.append(token_hash)
|
||
|
||
# 只保留最近的一定数量的 token
|
||
max_tokens = 100 # 每个用户每种类型最多保留100个 token
|
||
if len(user_tokens) > max_tokens:
|
||
user_tokens = user_tokens[-max_tokens:]
|
||
|
||
self.cache.set(key, user_tokens)
|
||
|
||
def _enforce_session_limit(self, user_id: str):
|
||
"""强制用户会话限制"""
|
||
max_sessions = self.config["security"]["max_sessions_per_user"]
|
||
|
||
access_key = f"user:{user_id}:access"
|
||
refresh_key = f"user:{user_id}:refresh"
|
||
|
||
# 检查 access token 数量
|
||
access_tokens = self.cache.get(access_key, [])
|
||
if len(access_tokens) > max_sessions:
|
||
# 删除最早的 token
|
||
tokens_to_remove = access_tokens[:-max_sessions]
|
||
for token_hash in tokens_to_remove:
|
||
if self.cache.exists(f"token:{token_hash}"):
|
||
self.cache.delete(f"token:{token_hash}")
|
||
|
||
# 更新索引
|
||
self.cache.set(access_key, access_tokens[-max_sessions:])
|
||
|
||
# 检查 refresh token 数量
|
||
refresh_tokens = self.cache.get(refresh_key, [])
|
||
if len(refresh_tokens) > max_sessions:
|
||
tokens_to_remove = refresh_tokens[:-max_sessions]
|
||
for token_hash in tokens_to_remove:
|
||
if self.cache.exists(f"token:{token_hash}"):
|
||
self.cache.delete(f"token:{token_hash}")
|
||
|
||
self.cache.set(refresh_key, refresh_tokens[-max_sessions:])
|
||
|
||
def validate_token(
|
||
self, token: str, token_type: str = "access"
|
||
) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
验证 JWT Token
|
||
|
||
Args:
|
||
token: JWT Token
|
||
token_type: 期望的 Token 类型
|
||
|
||
Returns:
|
||
验证成功返回用户数据，失败返回 None
|
||
"""
|
||
try:
|
||
# 解码验证 Token
|
||
payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
|
||
|
||
# 检查 Token 类型
|
||
print(f"payload: {payload}")
|
||
if payload.get("type") != token_type:
|
||
return None
|
||
|
||
user_id = payload.get("sub")
|
||
jti = payload.get("jti")
|
||
|
||
if not user_id or not jti:
|
||
return None
|
||
|
||
token_hash = self._hash_token(token)
|
||
|
||
# 检查缓存
|
||
with self.lock:
|
||
cache_data = self.cache.get(f"token:{token_hash}")
|
||
if not cache_data:
|
||
return None
|
||
|
||
# 检查 Token 是否激活
|
||
if not cache_data.get("is_active", True):
|
||
return None
|
||
|
||
# 检查是否过期
|
||
expires_at = datetime.fromisoformat(cache_data["expires_at"])
|
||
if datetime.utcnow() > expires_at:
|
||
# 标记为不活跃
|
||
cache_data["is_active"] = False
|
||
self.cache.set(f"token:{token_hash}", cache_data)
|
||
return None
|
||
|
||
# 更新最后使用时间
|
||
cache_data["last_used"] = datetime.utcnow().isoformat()
|
||
self.cache.set(f"token:{token_hash}", cache_data)
|
||
|
||
return {
|
||
"user_id": user_id,
|
||
"device_id": cache_data.get("device_id"),
|
||
"payload": payload,
|
||
"cache_data": cache_data,
|
||
}
|
||
|
||
except jwt.ExpiredSignatureError:
|
||
return None
|
||
except jwt.InvalidTokenError:
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
def invalidate_token(self, token: str):
|
||
"""使单个 Token 失效"""
|
||
with self.lock:
|
||
token_hash = self._hash_token(token)
|
||
|
||
# 标记为不活跃
|
||
if self.cache.exists(f"token:{token_hash}"):
|
||
cache_data = self.cache.get(f"token:{token_hash}")
|
||
cache_data["is_active"] = False
|
||
self.cache.set(f"token:{token_hash}", cache_data)
|
||
|
||
def invalidate_user_tokens(self, user_id: str, token_type: Optional[str] = None):
|
||
"""使用户的所有 Token 失效"""
|
||
with self.lock:
|
||
if token_type:
|
||
keys = [f"user:{user_id}:{token_type}"]
|
||
else:
|
||
keys = [f"user:{user_id}:access", f"user:{user_id}:refresh"]
|
||
|
||
for key in keys:
|
||
token_hashes = self.cache.get(key, [])
|
||
for token_hash in token_hashes:
|
||
# 添加到黑名单
|
||
# 注意：这里我们只有 token_hash，没有原始 token
|
||
# 所以只能标记为不活跃
|
||
if self.cache.exists(f"token:{token_hash}"):
|
||
cache_data = self.cache.get(f"token:{token_hash}")
|
||
cache_data["is_active"] = False
|
||
self.cache.set(f"token:{token_hash}", cache_data)
|
||
|
||
# 清空用户会话索引
|
||
self.cache.set(key, [])
|
||
|
||
def get_user_sessions(self, user_id: str) -> List[Dict[str, Any]]:
|
||
"""获取用户的所有活跃会话"""
|
||
with self.lock:
|
||
sessions = []
|
||
|
||
for token_type in ["access", "refresh"]:
|
||
key = f"user:{user_id}:{token_type}"
|
||
token_hashes = self.cache.get(key, [])
|
||
|
||
for token_hash in token_hashes:
|
||
if self.cache.exists(f"token:{token_hash}"):
|
||
cache_data = self.cache.get(f"token:{token_hash}")
|
||
|
||
# 检查是否活跃且未过期
|
||
if cache_data.get("is_active", True):
|
||
expires_at = datetime.fromisoformat(
|
||
cache_data["expires_at"]
|
||
)
|
||
if datetime.utcnow() <= expires_at:
|
||
sessions.append(
|
||
{
|
||
"token_type": token_type,
|
||
"device_id": cache_data.get("device_id"),
|
||
"user_agent": cache_data.get("user_agent"),
|
||
"ip_address": cache_data.get("ip_address"),
|
||
"created_at": cache_data["created_at"],
|
||
"last_used": cache_data["last_used"],
|
||
"expires_at": cache_data["expires_at"],
|
||
}
|
||
)
|
||
|
||
return sessions
|
||
|
||
def cleanup_expired(self):
|
||
"""清理过期数据"""
|
||
with self.lock:
|
||
cleaned_count = 0
|
||
current_time = datetime.utcnow()
|
||
|
||
# 清理过期 Token
|
||
for token_hash in self.cache.keys("token:*"):
|
||
expires_at = datetime.fromisoformat(self.cache.get(token_hash))
|
||
if current_time > expires_at:
|
||
self.cache.delete(token_hash)
|
||
cleaned_count += 1
|
||
|
||
return cleaned_count
|
||
|
||
|
||
# 全局缓存实例
|
||
_jwt_cache = None
|
||
|
||
|
||
def get_jwt_cache(config=None, cache=None) -> JWTCache:
|
||
"""获取 JWT 缓存实例（单例模式）"""
|
||
global _jwt_cache
|
||
if _jwt_cache is None:
|
||
if config is None or cache is None:
|
||
raise ValueError("参数 config 和 cache 不能为空")
|
||
_jwt_cache = JWTCache(config, cache)
|
||
return _jwt_cache
|
||
|
||
|
||
# Token 响应模型
|
||
class TokenResponse(BaseModel):
|
||
"""Token 响应模型"""
|
||
|
||
access_token: str
|
||
refresh_token: str
|
||
token_type: str = "bearer"
|
||
expires_in: int
|
||
refresh_expires_in: int
|
||
|
||
|
||
class RefreshTokenResponse(BaseModel):
|
||
"""刷新 Token 响应模型"""
|
||
|
||
access_token: str
|
||
token_type: str = "bearer"
|
||
|
||
|
||
# Token 数据模型
|
||
class TokenData(BaseModel):
|
||
"""Token 数据模型"""
|
||
|
||
user_id: str
|
||
username: Optional[str] = None
|
||
email: Optional[str] = None
|
||
scopes: List[str] = []
|