|
| 1 | +import hashlib |
| 2 | +import hmac |
| 3 | +import time |
| 4 | +from functools import wraps |
| 5 | + |
| 6 | +from django.http import JsonResponse |
| 7 | + |
| 8 | +from maxkb.const import CONFIG |
| 9 | + |
| 10 | + |
| 11 | +def mcp_token_required(view_func): |
| 12 | + """MCP内部令牌验证装饰器""" |
| 13 | + |
| 14 | + @wraps(view_func) |
| 15 | + def wrapper(self, request, *args, **kwargs): |
| 16 | + # 1. 验证IP白名单 |
| 17 | + client_ip = request.META.get('REMOTE_ADDR') |
| 18 | + if client_ip not in ['127.0.0.1', '::1']: |
| 19 | + return JsonResponse({'code': 403, 'message': 'Access denied'}, status=403) |
| 20 | + |
| 21 | + # 2. 验证MCP令牌 |
| 22 | + mcp_token = request.headers.get('X-MCP-Token') |
| 23 | + timestamp = request.headers.get('X-MCP-Timestamp') |
| 24 | + |
| 25 | + # 允许无令牌请求通过 |
| 26 | + if not mcp_token or not timestamp: |
| 27 | + return view_func(self, request, *args, **kwargs) |
| 28 | + |
| 29 | + # 3. 验证时间戳(防止重放攻击,5分钟有效期) |
| 30 | + try: |
| 31 | + ts = int(timestamp) |
| 32 | + if abs(time.time() - ts) > 300: # 5分钟 |
| 33 | + return JsonResponse({'code': 401, 'message': 'Token expired'}, status=401) |
| 34 | + except ValueError: |
| 35 | + return JsonResponse({'code': 401, 'message': 'Invalid timestamp'}, status=401) |
| 36 | + |
| 37 | + # 4. 验证令牌签名 |
| 38 | + secret = CONFIG.get('MCP_INTERNAL_SECRET', 'your-secret-key') |
| 39 | + |
| 40 | + # 从Authorization获取API Key |
| 41 | + auth_header = request.headers.get('Authorization', '') |
| 42 | + api_key = auth_header.replace('Bearer ', '') if auth_header.startswith('Bearer ') else '' |
| 43 | + |
| 44 | + # 重新计算签名 |
| 45 | + token_data = f"{api_key}:{timestamp}" |
| 46 | + expected_token = hmac.new( |
| 47 | + secret.encode(), |
| 48 | + token_data.encode(), |
| 49 | + hashlib.sha256 |
| 50 | + ).hexdigest() |
| 51 | + |
| 52 | + # print(expected_token, mcp_token) |
| 53 | + |
| 54 | + if not hmac.compare_digest(mcp_token, expected_token): |
| 55 | + return JsonResponse({'code': 401, 'message': 'Invalid MCP token'}, status=401) |
| 56 | + |
| 57 | + return view_func(self, request, *args, **kwargs) |
| 58 | + |
| 59 | + return wrapper |
0 commit comments