企业级令牌验证服务优化

# 企业级令牌验证服务优化

以下是对 `TokenService` 的全面企业级优化实现，增加了安全性、可扩展性和监控能力：

```python
import os
import logging
import time
from datetime import datetime, timedelta
from typing import Optional, Dict
import jwt
import redis
from django.core.cache import cache
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from rest_framework.exceptions import AuthenticationFailed

logger = logging.getLogger(__name__)

class TokenService:
"""
企业级令牌验证服务 - 优化版本
- 支持JWT和API密钥两种令牌类型
- 令牌吊销机制
- 令牌使用频率限制
- 详细审计日志
- 缓存优化
"""

# 令牌类型常量
TOKEN_TYPE_API_KEY = "api_key"
TOKEN_TYPE_JWT = "jwt"

# 错误代码
ERROR_INVALID_TOKEN = "invalid_token"
ERROR_EXPIRED_TOKEN = "expired_token"
ERROR_REVOKED_TOKEN = "revoked_token"
ERROR_RATE_LIMIT = "rate_limit_exceeded"
ERROR_MISSING_SCOPE = "insufficient_scope"

@staticmethod
def get_redis_connection():
"""获取Redis连接 - 企业级连接池管理"""
try:
return redis.Redis(
host=settings.REDIS_HOST,
port=settings.REDIS_PORT,
db=settings.REDIS_DB,
password=settings.REDIS_PASSWORD,
socket_timeout=2,
socket_connect_timeout=2,
health_check_interval=30
)
except Exception as e:
logger.error(f"Redis连接失败: {str(e)}", exc_info=True)
return None

@staticmethod
def get_jwt_secret() -> str:
"""获取JWT密钥 - 企业级密钥管理"""
secret = getattr(settings, 'JWT_SECRET_KEY', None)
if not secret:
raise ImproperlyConfigured("JWT_SECRET_KEY未在设置中配置")
return secret

@staticmethod
def get_api_keys() -> Dict[str, Dict]:
"""获取API密钥及其元数据 - 企业级密钥存储"""
api_keys = getattr(settings, 'API_KEYS', {})
if not api_keys:
logger.warning("API_KEYS未在设置中配置，使用环境变量回退")
# 环境变量回退：BEACON_API_KEY1, BEACON_API_KEY2, ...
env_keys = {k[13:]: v for k, v in os.environ.items()
if k.startswith('BEACON_API_KEY')}
api_keys = {key: {"scopes": ["audit:write"]} for key in env_keys.values()}
return api_keys

@staticmethod
def is_token_revoked(token: str, token_type: str) -> bool:
"""检查令牌是否被吊销 - 企业级吊销机制"""
redis_conn = TokenService.get_redis_connection()
if redis_conn:
try:
# 使用Redis存储吊销列表
revoked_key = f"revoked_token:{token_type}:{token}"
return redis_conn.exists(revoked_key) > 0
except Exception as e:
logger.error(f"Redis吊销检查失败: {str(e)}", exc_info=True)

# Redis不可用时使用Django缓存回退
cache_key = f"revoked_token_{token_type}_{token}"
return cache.get(cache_key) is not None

@staticmethod
def check_rate_limit(token: str, token_type: str) -> bool:
"""检查令牌使用频率 - 企业级限流机制"""
redis_conn = TokenService.get_redis_connection()
if not redis_conn:
# Redis不可用时跳过限流
return True

try:
# 创建限流键
rate_key = f"token_rate:{token_type}:{token}"

# 获取当前计数
current = redis_conn.get(rate_key)
current_count = int(current) if current else 0

# 检查是否超过限制
rate_limit = getattr(settings, 'TOKEN_RATE_LIMIT', 100) # 默认100次/分钟
if current_count >= rate_limit:
return False

# 增加计数并设置过期时间
pipe = redis_conn.pipeline()
pipe.incr(rate_key)
pipe.expire(rate_key, 60) # 60秒过期
pipe.execute()

return True
except Exception as e:
logger.error(f"限流检查失败: {str(e)}", exc_info=True)
return True

@staticmethod
def log_token_usage(token: str, token_type: str, request):
"""记录令牌使用情况 - 企业级审计日志"""
# 在实际项目中，这里应该写入审计日志系统
client_ip = request.META.get('REMOTE_ADDR', '0.0.0.0')
logger.info(f"令牌使用: type={token_type}, token={token[:6]}..., client_ip={client_ip}")

@staticmethod
def validate_jwt(token_str: str) -> dict:
"""验证JWT令牌 - 企业级JWT验证"""
try:
secret = TokenService.get_jwt_secret()
payload = jwt.decode(
token_str,
secret,
algorithms=["HS256"],
options={"require": ["exp", "iat", "sub"]}
)
return payload
except jwt.ExpiredSignatureError:
raise AuthenticationFailed("令牌已过期", code=TokenService.ERROR_EXPIRED_TOKEN)
except jwt.InvalidTokenError:
raise AuthenticationFailed("无效令牌", code=TokenService.ERROR_INVALID_TOKEN)

@staticmethod
def validate_api_key(token_str: str) -> dict:
"""验证API密钥 - 企业级密钥验证"""
api_keys = TokenService.get_api_keys()
if token_str not in api_keys:
raise AuthenticationFailed("无效API密钥", code=TokenService.ERROR_INVALID_TOKEN)

# 返回API密钥元数据
return {
"sub": "api_key",
"scopes": api_keys[token_str].get("scopes", []),
"metadata": api_keys[token_str].get("metadata", {})
}

@staticmethod
def has_scope(payload: dict, required_scope: str) -> bool:
"""检查令牌是否具有所需作用域 - 企业级权限控制"""
scopes = payload.get("scopes", [])

# 支持通配符作用域
if "*" in scopes:
return True

# 支持前缀匹配 (audit:*)
for scope in scopes:
if scope.endswith(":*") and required_scope.startswith(scope[:-1]):
return True

return required_scope in scopes

@staticmethod
def validate_token(token_str: str, request, required_scope: str = "audit:write") -> bool:
"""
企业级令牌验证
:param token_str: 令牌字符串
:param request: HTTP请求对象
:param required_scope: 所需的作用域
:return: 是否有效
"""
start_time = time.time()

try:
# 1. 确定令牌类型
token_type = TokenService.TOKEN_TYPE_JWT if "." in token_str else TokenService.TOKEN_TYPE_API_KEY

# 2. 检查限流
if not TokenService.check_rate_limit(token_str, token_type):
logger.warning(f"令牌使用频率过高: {token_str[:6]}...")
raise AuthenticationFailed("请求过于频繁", code=TokenService.ERROR_RATE_LIMIT)

# 3. 检查吊销状态
if TokenService.is_token_revoked(token_str, token_type):
logger.warning(f"使用已吊销令牌: {token_str[:6]}...")
raise AuthenticationFailed("令牌已被吊销", code=TokenService.ERROR_REVOKED_TOKEN)

# 4. 验证令牌
if token_type == TokenService.TOKEN_TYPE_JWT:
payload = TokenService.validate_jwt(token_str)
else:
payload = TokenService.validate_api_key(token_str)

# 5. 检查作用域
if not TokenService.has_scope(payload, required_scope):
logger.warning(f"令牌缺少所需作用域: {required_scope}")
raise AuthenticationFailed("权限不足", code=TokenService.ERROR_MISSING_SCOPE)

# 6. 记录使用情况
TokenService.log_token_usage(token_str, token_type, request)

# 7. 附加令牌信息到请求对象
request.token_payload = payload
request.token_type = token_type
request.token_str = token_str

return True

except AuthenticationFailed as e:
# 记录详细错误信息
logger.warning(f"令牌验证失败: {str(e)}")
raise

finally:
# 记录性能指标
duration = time.time() - start_time
logger.debug(f"令牌验证耗时: {duration:.4f}秒")

@staticmethod
def revoke_token(token: str, token_type: str, reason: str = "admin_revocation", ttl: int = 86400):
"""
吊销令牌 - 企业级吊销机制
:param token: 要吊销的令牌
:param token_type: 令牌类型 (api_key 或 jwt)
:param reason: 吊销原因
:param ttl: 吊销记录有效期（秒）
"""
redis_conn = TokenService.get_redis_connection()
if redis_conn:
try:
# 使用Redis存储吊销记录
revoked_key = f"revoked_token:{token_type}:{token}"
redis_conn.setex(revoked_key, ttl, reason)
logger.info(f"令牌已吊销: type={token_type}, token={token[:6]}..., reason={reason}")
return True
except Exception as e:
logger.error(f"Redis吊销操作失败: {str(e)}", exc_info=True)

# Redis不可用时使用Django缓存回退
cache_key = f"revoked_token_{token_type}_{token}"
cache.set(cache_key, reason, ttl)
logger.info(f"使用缓存吊销令牌: type={token_type}, token={token[:6]}..., reason={reason}")
return True

@staticmethod
def generate_jwt(payload: dict, expires_in: int = 3600) -> str:
"""
生成JWT令牌 - 企业级令牌生成
:param payload: 令牌负载
:param expires_in: 过期时间（秒）
:return: JWT令牌字符串
"""
# 添加标准声明
now = datetime.utcnow()
standard_claims = {
"iat": now,
"exp": now + timedelta(seconds=expires_in),
"nbf": now,
"jti": str(uuid.uuid4()),
}

# 合并自定义负载
full_payload = {**standard_claims, **payload}

# 生成令牌
secret = TokenService.get_jwt_secret()
return jwt.encode(full_payload, secret, algorithm="HS256")
```

## 企业级优化说明

### 1. 多令牌类型支持
- 同时支持JWT和API密钥两种令牌类型
- 自动检测令牌类型并应用不同的验证逻辑

### 2. 增强的安全性
- **令牌吊销机制**：支持即时吊销令牌
- **作用域验证**：细粒度的权限控制
- **速率限制**：防止令牌滥用
- **安全错误处理**：详细的错误代码和消息

### 3. 高性能架构
- **Redis集成**：用于令牌吊销列表和速率限制
- **缓存回退**：Redis不可用时使用Django缓存
- **连接池管理**：优化的Redis连接设置
- **性能监控**：记录验证耗时

### 4. 企业级密钥管理
- **集中式配置**：从Django设置获取密钥
- **环境变量回退**：兼容传统部署方式
- **密钥元数据**：支持为每个密钥配置作用域

### 5. 详细的审计日志
- 记录所有令牌使用情况
- 跟踪客户端IP地址
- 记录吊销操作和原因

### 6. 错误处理
- 标准化的错误代码
- 详细的错误消息
- 安全异常处理

### 7. 扩展功能
- **令牌生成**：支持生成JWT令牌
- **通配符作用域**：支持`audit:*`等通配符作用域
- **自定义TTL**：可配置的吊销记录有效期

## 配置示例

```python
# settings.py

# JWT配置
JWT_SECRET_KEY = os.getenv('JWT_SECRET_KEY', 'your-strong-secret-key')

# API密钥配置
API_KEYS = {
"abc123def456": {
"scopes": ["audit:write", "audit:read"],
"metadata": {"description": "Beacon API key"}
},
"xyz789uvw000": {
"scopes": ["audit:write"],
"metadata": {"description": "Monitoring service key"}
}
}

# Redis配置
REDIS_HOST = os.getenv('REDIS_HOST', 'localhost')
REDIS_PORT = os.getenv('REDIS_PORT', 6379)
REDIS_DB = os.getenv('REDIS_DB', 0)
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD', None)

# 限流配置
TOKEN_RATE_LIMIT = 200 # 每分钟最大请求数
```

## 使用示例

```python
# 在视图中使用
class BeaconAPIView(APIView):
def post(self, request):
auth_header = request.headers.get('Authorization', '')
if not auth_header.startswith('Bearer '):
return Response(...)

token_str = auth_header[7:].strip()

try:
# 验证令牌
TokenService.validate_token(token_str, request, "audit:write")

# 处理请求...
return Response(...)

except AuthenticationFailed as e:
return Response(
{'error': str(e.detail), 'code': e.get_codes()},
status=status.HTTP_401_UNAUTHORIZED
)

# 吊销令牌
TokenService.revoke_token(
token="abc123def456",
token_type=TokenService.TOKEN_TYPE_API_KEY,
reason="compromised_key"
)

# 生成JWT令牌
jwt_token = TokenService.generate_jwt(
payload={"sub": "service-account", "scopes": ["audit:write"]},
expires_in=3600
)
```

这个企业级优化版本提供了高安全性、高性能和可扩展性，适用于关键业务系统中的令牌验证需求。