Skip to content

Commit 4d28001

Browse files
committed
feat: add MCP token validation decorator for API access control
1 parent 8a785b2 commit 4d28001

File tree

1 file changed

+59
-0
lines changed

1 file changed

+59
-0
lines changed

apps/common/auth/mcp_auth_token.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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

Comments
 (0)