Upload proxy_handler.py

proxy_handler.py

@@ -6,6 +6,10 @@ from typing import AsyncGenerator, Dict, Any, Tuple, List
 import httpx
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
+import hashlib
+import hmac
+import urllib.parse
+from datetime import datetime, timezone
 
 from config import settings
 from cookie_manager import cookie_manager
@@ -21,11 +25,43 @@ class ProxyHandler:
             limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
             http2=True,
         )
+        # 从JavaScript代码中获取的固定密钥
+        self.primary_secret = "junjie".encode('utf-8')
 
     async def aclose(self):
         if not self.client.is_closed:
             await self.client.aclose()
-    
+            
+    # --- 新增方法：获取毫秒级时间戳 ---
+    def _get_timestamp_millis(self) -> int:
+        return int(time.time() * 1000)
+
+    # --- 新增方法：解析JWT以获取用户ID ---
+    def _parse_jwt_token(self, token: str) -> Dict[str, str]:
+        try:
+            parts = token.split('.')
+            if len(parts) != 3:
+                return {"userId": ""}
+            payload = json.loads(base64.urlsafe_b64decode(parts[1] + '==').decode('utf-8'))
+            return {"userId": payload.get("sub", "")}
+        except Exception:
+            return {"userId": ""}
+
+    # --- 新增方法：生成签名 ---
+    def _generate_signature(self, token: str, payload_str: str, mt: str) -> Tuple[str, int]:
+        timestamp_ms = self._get_timestamp_millis()
+        minute_bucket = str(timestamp_ms // 60000)
+
+        level1_data = f"{token}|{minute_bucket}".encode('utf-8')
+        mac1 = hmac.new(self.primary_secret, level1_data, hashlib.sha256)
+        derived_key_hex = mac1.hexdigest()
+
+        level2_data = f"{payload_str}|{mt}|{timestamp_ms}".encode('utf-8')
+        mac2 = hmac.new(derived_key_hex.encode('utf-8'), level2_data, hashlib.sha256)
+        signature = mac2.hexdigest()
+        
+        return signature, timestamp_ms
+
     def _clean_thinking_content(self, text: str) -> str:
         """
         Aggressively cleans raw thinking content strings based on observed patterns
@@ -84,19 +120,82 @@ class ProxyHandler:
             else: out.append({"role": getattr(m, "role", "user"), "content": getattr(m, "content", str(m))})
         return out
 
-    async def _prep_upstream(self, req: ChatCompletionRequest) -> Tuple[Dict[str, Any], Dict[str, str], str]:
-        """Prepares the request body, headers, and cookie for the upstream API."""
+    # --- 重构 _prep_upstream 方法以加入签名逻辑 ---
+    async def _prep_upstream(self, req: ChatCompletionRequest) -> Tuple[Dict[str, Any], Dict[str, str], str, str]:
+        """Prepares the request body, headers, cookie, and URL for the upstream API."""
         ck = await cookie_manager.get_next_cookie()
         if not ck: raise HTTPException(503, "No available cookies")
+
         model = settings.UPSTREAM_MODEL if req.model == settings.MODEL_NAME else req.model
-        body = { "stream": True, "model": model, "messages": self._serialize_msgs(req.messages), "background_tasks": {"title_generation": True, "tags_generation": True}, "chat_id": str(uuid.uuid4()), "features": {"image_generation": False, "code_interpreter": False, "web_search": False, "auto_web_search": False, "enable_thinking": True,}, "id": str(uuid.uuid4()), "mcp_servers": ["deep-web-search"], "model_item": {"id": model, "name": "GLM-4.6", "owned_by": "openai"}, "params": {}, "tool_servers": [], "variables": {"{{USER_NAME}}": "User", "{{USER_LOCATION}}": "Unknown", "{{CURRENT_DATETIME}}": time.strftime("%Y-%m-%d %H:%M:%S"),},}
-        headers = { "Content-Type": "application/json", "Authorization": f"Bearer {ck}", "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"), "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN", "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.79","X-Signature": "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=", "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",}
-        return body, headers, ck
+        chat_id = str(uuid.uuid4())
+        request_id = str(uuid.uuid4())
+        user_info = self._parse_jwt_token(ck)
+        user_id = user_info.get("userId", "")
+        
+        body = { "stream": True, "model": model, "messages": self._serialize_msgs(req.messages), "background_tasks": {"title_generation": True, "tags_generation": True}, "chat_id": chat_id, "features": {"image_generation": False, "code_interpreter": False, "web_search": False, "auto_web_search": False, "enable_thinking": True,}, "id": request_id, "mcp_servers": ["deep-web-search"], "model_item": {"id": model, "name": "GLM-4.6", "owned_by": "openai"}, "params": {}, "tool_servers": [], "variables": {"{{USER_NAME}}": "User", "{{USER_LOCATION}}": "Unknown", "{{CURRENT_DATETIME}}": time.strftime("%Y-%m-%d %H:%M:%S"),},}
+        
+        # 构造用于签名的负载
+        timestamp = self._get_timestamp_millis()
+        now = datetime.now(timezone.utc)
+        
+        payload_data = {
+            'timestamp': str(timestamp),
+            'requestId': request_id,
+            'user_id': user_id,
+            'token': ck,
+            'user_agent': "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36",
+            'current_url': f"https://chat.z.ai/c/{chat_id}",
+            'pathname': f"/c/{chat_id}",
+            'timezone': 'UTC', # 简化时区处理
+            'timezone_offset': '0',
+            'local_time': now.isoformat(),
+            'utc_time': now.strftime('%a, %d %b %Y %H:%M:%S GMT'),
+            'version': '0.0.1',
+            'platform': 'web',
+            'language': 'zh-CN',
+            'languages': 'zh-CN,en',
+            'cookie_enabled': 'true',
+            'screen_width': '2560',
+            'screen_height': '1440',
+            'screen_resolution': '2560x1440',
+            'viewport_height': '1328',
+            'viewport_width': '1342',
+            'viewport_size': '1342x1328',
+            'color_depth': '24',
+            'pixel_ratio': '2',
+            'search': '',
+            'hash': '',
+            'host': 'chat.z.ai',
+            'hostname': 'chat.z.ai',
+            'protocol': 'https:',
+            'referrer': '',
+            'title': 'Chat with Z.ai - Free AI Chatbot powered by GLM-4.5',
+            'is_mobile': 'false',
+            'is_touch': 'false',
+            'max_touch_points': '0',
+            'browser_name': 'Chrome',
+            'os_name': 'Mac OS'
+        }
+        
+        keys = sorted(payload_data.keys())
+        sorted_payload = ",".join([f"{k},{payload_data[k]}" for k in keys])
+        url_params = urllib.parse.urlencode(payload_data)
+
+        # 获取最后一条消息作为 mt
+        last_message = req.messages[-1].content if req.messages else ""
+        
+        signature, sig_timestamp = self._generate_signature(ck, sorted_payload, last_message)
+        
+        final_url = f"{settings.UPSTREAM_URL}?{url_params}&signature_timestamp={sig_timestamp}"
+
+        headers = { "Content-Type": "application/json", "Authorization": f"Bearer {ck}", "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"), "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN", "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.79", "X-Signature": signature, "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",}
+        
+        return body, headers, ck, final_url
         
     async def stream_proxy_response(self, req: ChatCompletionRequest) -> AsyncGenerator[str, None]:
         ck = None
         try:
-            body, headers, ck = await self._prep_upstream(req)
+            body, headers, ck, url = await self._prep_upstream(req)
             comp_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
             think_open = False
             yielded_think_buffer = "" 
@@ -127,7 +226,7 @@ class ProxyHandler:
                     if cleaned_text:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': cleaned_text}, 'finish_reason': None}]})}\n\n"
 
-            async with self.client.stream("POST", settings.UPSTREAM_URL, json=body, headers=headers) as resp:
+            async with self.client.stream("POST", url, json=body, headers=headers) as resp:
                 if resp.status_code != 200:
                     await cookie_manager.mark_cookie_failed(ck); err_body = await resp.aread()
                     err_msg = f"Error: {resp.status_code} - {err_body.decode(errors='ignore')}"
@@ -186,10 +285,10 @@ class ProxyHandler:
     async def non_stream_proxy_response(self, req: ChatCompletionRequest) -> ChatCompletionResponse:
         ck = None
         try:
-            body, headers, ck = await self._prep_upstream(req)
+            body, headers, ck, url = await self._prep_upstream(req)
             last_thinking_content = ""
             raw_answer_parts = []
-            async with self.client.stream("POST", settings.UPSTREAM_URL, json=body, headers=headers) as resp:
+            async with self.client.stream("POST", url, json=body, headers=headers) as resp:
                 if resp.status_code != 200:
                     await cookie_manager.mark_cookie_failed(ck); error_detail = await resp.text()
                     raise HTTPException(resp.status_code, f"Upstream error: {error_detail}")