forked from erp-dev/erp
117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
"""
|
||
SSE认证工具模块
|
||
|
||
提供统一的认证和商户验证功能，用于SSE连接端点和其他需要认证的地方
|
||
"""
|
||
|
||
from django.http import HttpResponseForbidden
|
||
from rest_framework.authentication import get_authorization_header
|
||
from rest_framework_simplejwt.authentication import JWTAuthentication
|
||
from rest_framework_simplejwt.exceptions import InvalidToken
|
||
from rest_framework.exceptions import AuthenticationFailed
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def get_user_merchant_id(request):
|
||
"""获取用户所属商户ID"""
|
||
try:
|
||
# 检查用户是否已认证
|
||
if not request.user.is_authenticated:
|
||
return None
|
||
# 获取商户ID
|
||
return request.user.employee.merchant_id
|
||
except (AttributeError, AttributeError):
|
||
return None
|
||
|
||
|
||
def authenticate_sse_request(request):
|
||
"""
|
||
对SSE连接请求进行认证
|
||
|
||
参数:
|
||
- request: Django HTTP请求对象
|
||
|
||
返回:
|
||
- tuple: (user, merchant_id) 认证成功返回用户和商户ID
|
||
- None: 认证失败返回None
|
||
"""
|
||
jwt_authenticator = JWTAuthentication()
|
||
|
||
try:
|
||
# 获取Authorization头
|
||
auth_header = get_authorization_header(request).split()
|
||
if not auth_header or auth_header[0].lower() != b'bearer':
|
||
logger.warning("SSE认证失败: 缺少Bearer token")
|
||
return None
|
||
|
||
if len(auth_header) < 2:
|
||
logger.warning("SSE认证失败: Bearer token格式错误")
|
||
return None
|
||
|
||
token = auth_header[1].decode('utf-8')
|
||
|
||
# 验证token
|
||
validated_token = jwt_authenticator.get_validated_token(token)
|
||
user = jwt_authenticator.get_user(validated_token)
|
||
|
||
# 检查用户是否有关联的员工和商户
|
||
try:
|
||
merchant_id = user.employee.merchant_id
|
||
if not merchant_id:
|
||
logger.warning(f"用户 {user.username} 无关联商户")
|
||
return None
|
||
return user, merchant_id
|
||
except AttributeError:
|
||
logger.warning(f"用户 {user.username} 无关联员工")
|
||
return None
|
||
|
||
except (IndexError, InvalidToken, AuthenticationFailed) as e:
|
||
logger.warning(f"SSE认证失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
def require_sse_authentication(view_func):
|
||
"""
|
||
装饰器：要求SSE连接认证
|
||
|
||
用于SSE连接端点的认证，如果认证失败则返回403响应
|
||
OPTIONS请求跳过认证
|
||
"""
|
||
def wrapper(request, *args, **kwargs):
|
||
# OPTIONS请求跳过认证
|
||
if request.method == 'OPTIONS':
|
||
return view_func(request, *args, **kwargs)
|
||
|
||
auth_result = authenticate_sse_request(request)
|
||
if not auth_result:
|
||
return HttpResponseForbidden("Authentication failed or user has no associated merchant")
|
||
|
||
user, merchant_id = auth_result
|
||
request.user = user
|
||
request.merchant_id = merchant_id
|
||
|
||
return view_func(request, *args, **kwargs)
|
||
|
||
return wrapper
|
||
|
||
|
||
def validate_sse_request(request):
|
||
"""
|
||
验证SSE请求的认证状态
|
||
|
||
返回:
|
||
- tuple: (is_valid, error_response)
|
||
- is_valid: bool 表示请求是否有效
|
||
- error_response: HttpResponse 如果无效则返回错误响应，否则为None
|
||
"""
|
||
auth_result = authenticate_sse_request(request)
|
||
if not auth_result:
|
||
return False, HttpResponseForbidden("认证失败或用户无关联商户")
|
||
|
||
user, merchant_id = auth_result
|
||
request.user = user
|
||
request.merchant_id = merchant_id
|
||
|
||
return True, None |