Upload 2 files

Dockerfile

@@ -2,7 +2,7 @@ FROM python:3.10-slim
 
 WORKDIR /app
 
-RUN pip install --no-cache-dir flask requests curl_cffi werkzeug loguru 
+RUN pip install --no-cache-dir flask requests curl_cffi werkzeug loguru dotenv
 
 COPY . .
 

app.py

@@ -5,13 +5,19 @@ import time
 import base64
 import sys
 import inspect
+import secrets
 from loguru import logger
+from dotenv import load_dotenv
 
 import requests
-from flask import Flask, request, Response, jsonify, stream_with_context
+from flask import Flask, request, Response, jsonify, stream_with_context, render_template, redirect, session
 from curl_cffi import requests as curl_requests
 from werkzeug.middleware.proxy_fix import ProxyFix
 
+current_dir = os.path.dirname(os.path.abspath(__file__))
+env_path = os.path.join(current_dir, '.env')
+
+load_dotenv(env_path)
 
 class Logger:
     def __init__(self, level="INFO", colorize=True, format=None):
@@ -93,27 +99,34 @@ CONFIG = {
         "grok-3-reasoning": "grok-3"
     },
     "API": {
-        "IS_TEMP_CONVERSATION": os.environ.get("IS_TEMP_CONVERSATION", "true").lower() == "true",
-        "IS_CUSTOM_SSO": os.environ.get("IS_CUSTOM_SSO", "false").lower() == "true",
-        "BASE_URL": "https://shadowfetch.guaq.workers.dev/image/https://grok.com",
-        "API_KEY": os.environ.get("API_KEY", "sk-123456"),
+        "IS_TEMP_CONVERSATION": os.getenv("IS_TEMP_CONVERSATION", "true").lower() == "true",
+        "IS_CUSTOM_SSO": os.getenv("IS_CUSTOM_SSO", "false").lower() == "true",
+        "BASE_URL": "https://grok.com",
+        "API_KEY": os.getenv("API_KEY", "sk-123456"),
         "SIGNATURE_COOKIE": None,
-        "PICGO_KEY": os.environ.get("PICGO_KEY") or None,
-        "TUMY_KEY": os.environ.get("TUMY_KEY") or None,
+        "PICGO_KEY": os.getenv("PICGO_KEY") or None,
+        "TUMY_KEY": os.getenv("TUMY_KEY") or None,
         "RETRY_TIME": 1000,
-        "PROXY": os.environ.get("PROXY") or None
+        "PROXY": os.getenv("PROXY") or None
+    },
+    "ADMIN": {
+        "MANAGER_SWITCH": os.getenv("MANAGER_SWITCH") or None,
+        "PASSWORD": os.getenv("ADMINPASSWORD") or None 
     },
     "SERVER": {
-        "PORT": int(os.environ.get("PORT", 5200))
+        "COOKIE": None,
+        "CF_CLEARANCE":os.getenv("CF_CLEARANCE") or None,
+        "PORT": int(os.getenv("PORT", 5200))
     },
     "RETRY": {
+        "RETRYSWITCH": False,
         "MAX_ATTEMPTS": 2
     },
-    "SHOW_THINKING": os.environ.get("SHOW_THINKING") == "true",
+    "SHOW_THINKING": os.getenv("SHOW_THINKING") == "true",
     "IS_THINKING": False,
     "IS_IMG_GEN": False,
     "IS_IMG_GEN2": False,
-    "ISSHOW_SEARCH_RESULTS": os.environ.get("ISSHOW_SEARCH_RESULTS", "true").lower() == "true"
+    "ISSHOW_SEARCH_RESULTS": os.getenv("ISSHOW_SEARCH_RESULTS", "true").lower() == "true"
 }
 
 
@@ -217,14 +230,48 @@ class AuthTokenManager:
         except Exception as error:
             logger.error(f"令牌删除失败: {str(error)}")
             return False
-
-    def get_next_token_for_model(self, model_id):
+    def reduce_token_request_count(self, model_id, count):
+        try:
+            normalized_model = self.normalize_model_name(model_id)
+            
+            if normalized_model not in self.token_model_map:
+                logger.error(f"模型 {normalized_model} 不存在", "TokenManager")
+                return False
+                
+            if not self.token_model_map[normalized_model]:
+                logger.error(f"模型 {normalized_model} 没有可用的token", "TokenManager")
+                return False
+                
+            token_entry = self.token_model_map[normalized_model][0]
+            
+            # 确保RequestCount不会小于0
+            new_count = max(0, token_entry["RequestCount"] - count)
+            reduction = token_entry["RequestCount"] - new_count
+            
+            token_entry["RequestCount"] = new_count
+            
+            # 更新token状态
+            if token_entry["token"]:
+                sso = token_entry["token"].split("sso=")[1].split(";")[0]
+                if sso in self.token_status_map and normalized_model in self.token_status_map[sso]:
+                    self.token_status_map[sso][normalized_model]["totalRequestCount"] = max(
+                        0, 
+                        self.token_status_map[sso][normalized_model]["totalRequestCount"] - reduction
+                    )
+            return True
+            
+        except Exception as error:
+            logger.error(f"重置校对token请求次数时发生错误: {str(error)}", "TokenManager")
+            return False
+    def get_next_token_for_model(self, model_id, is_return=False):
         normalized_model = self.normalize_model_name(model_id)
 
         if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:
             return None
 
         token_entry = self.token_model_map[normalized_model][0]
+        if is_return:
+            return token_entry["token"]
 
         if token_entry:
             if token_entry["StartCallTime"] is None:
@@ -377,6 +424,14 @@ class AuthTokenManager:
             for entry in model_tokens:
                 all_tokens.add(entry["token"])
         return list(all_tokens)
+    def get_current_token(self, model_id):
+        normalized_model = self.normalize_model_name(model_id)
+
+        if normalized_model not in self.token_model_map or not self.token_model_map[normalized_model]:
+            return None
+
+        token_entry = self.token_model_map[normalized_model][0]
+        return token_entry["token"]
 
     def get_token_status_map(self):
         return self.token_status_map
@@ -401,8 +456,8 @@ class Utils:
         return '\n\n'.join(formatted_results)
 
     @staticmethod
-    def create_auth_headers(model):
-        return token_manager.get_next_token_for_model(model)
+    def create_auth_headers(model, is_return=False):
+        return token_manager.get_next_token_for_model(model, is_return)
 
     @staticmethod
     def get_proxy_options():
@@ -411,12 +466,17 @@ class Utils:
 
         if proxy:
             logger.info(f"使用代理: {proxy}", "Server")
-            proxy_options["proxies"] = {"https": proxy, "http": proxy}
-
+            
             if proxy.startswith("socks5://"):
-                proxy_options["proxies"] = {"https": proxy, "http": proxy}
-                proxy_options["proxy_type"] = "socks5"
-
+                proxy_options["proxy"] = proxy
+            
+                if '@' in proxy:
+                    auth_part = proxy.split('@')[0].split('://')[1]
+                    if ':' in auth_part:
+                        username, password = auth_part.split(':')
+                        proxy_options["proxy_auth"] = (username, password)
+            else:
+                proxy_options["proxies"] = {"https": proxy, "http": proxy}     
         return proxy_options
 
 class GrokApiClient:
@@ -445,7 +505,41 @@ class GrokApiClient:
             "mimeType": mime_type,
             "fileName": file_name
         }
+    def upload_base64_file(self, message, model):
+        try:
+            message_base64 = base64.b64encode(message.encode('utf-8')).decode('utf-8')
+            upload_data = {
+                "fileName": "message.txt",
+                "fileMimeType": "text/plain",
+                "content": message_base64
+            }
+
+            logger.info("发送文字文件请求", "Server")
+            cookie = f"{Utils.create_auth_headers(model, True)};{CONFIG['SERVER']['CF_CLEARANCE']}" 
+            proxy_options = Utils.get_proxy_options()
+            response = curl_requests.post(
+                "https://grok.com/rest/app-chat/upload-file",
+                headers={
+                    **DEFAULT_HEADERS,
+                    "Cookie":cookie
+                },
+                json=upload_data,
+                impersonate="chrome133a",
+                verify=False,
+                **proxy_options
+            )
+
+            if response.status_code != 200:
+                logger.error(f"上传文件失败,状态码:{response.status_code}", "Server")
+                raise Exception(f"上传文件失败,状态码:{response.status_code}")
+
+            result = response.json()
+            logger.info(f"上传文件成功: {result}", "Server")
+            return result.get("fileMetadataId", "")
 
+        except Exception as error:
+            logger.error(str(error), "Server")
+            raise Exception(f"上传文件失败,状态码:{response.status_code}")
     def upload_base64_image(self, base64_data, url):
         try:
             if 'data:image' in base64_data:
@@ -473,10 +567,11 @@ class GrokApiClient:
                 url,
                 headers={
                     **DEFAULT_HEADERS,
-                    "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
+                    "Cookie":CONFIG["SERVER"]['COOKIE']
                 },
                 json=upload_data,
                 impersonate="chrome133a",
+                verify=False,
                 **proxy_options
             )
 
@@ -504,11 +599,13 @@ class GrokApiClient:
             if last_message["role"] != 'user':
                 raise ValueError('此模型最后一条消息必须是用户消息!')
             todo_messages = [last_message]
-
         file_attachments = []
         messages = ''
         last_role = None
         last_content = ''
+        message_length = 0
+        convert_to_file = False
+        last_message_content = ''
         search = request["model"] in ['grok-2-search', 'grok-3-search']
 
         # 移除<think>标签及其内容和base64图片
@@ -558,7 +655,9 @@ class GrokApiClient:
 
 
             text_content = process_content(current.get("content", ""))
-
+            if is_last_message and convert_to_file:
+                last_message_content = f"{role.upper()}: {text_content or '[图片]'}\n"
+                continue
             if text_content or (is_last_message and file_attachments):
                 if role == last_role and text_content:
                     last_content += '\n' + text_content
@@ -567,9 +666,22 @@ class GrokApiClient:
                     messages += f"{role.upper()}: {text_content or '[图片]'}\n"
                     last_content = text_content
                     last_role = role
-
+            message_length += len(messages)
+            if message_length >= 40000:
+                convert_to_file = True
+               
+        if convert_to_file:
+            file_id = self.upload_base64_file(messages, request["model"])
+            if file_id:
+                file_attachments.insert(0, file_id)
+            messages = last_message_content.strip()
+        if messages.strip() == '':
+            if convert_to_file:
+                messages = '基于txt文件内容进行回复：'
+            else:
+                raise ValueError('消息内容为空!')
         return {
-            "temporary": CONFIG["API"]["IS_TEMP_CONVERSATION"],
+            "temporary": CONFIG["API"].get("IS_TEMP_CONVERSATION", False),
             "modelName": self.model_id,
             "message": messages.strip(),
             "fileAttachments": file_attachments[:4],
@@ -684,12 +796,12 @@ def handle_image_response(image_url):
         try:
             proxy_options = Utils.get_proxy_options()
             image_base64_response = curl_requests.get(
-                f"https://shadowfetch.guaq.workers.dev/image/https://assets.grok.com/{image_url}",
+                f"https://assets.grok.com/{image_url}",
                 headers={
                     **DEFAULT_HEADERS,
-                    "Cookie": CONFIG["API"]["SIGNATURE_COOKIE"]
+                    "Cookie":CONFIG["SERVER"]['COOKIE']
                 },
-                impersonate="chrome120",
+                impersonate="chrome133a",
                 **proxy_options
             )
 
@@ -855,7 +967,7 @@ def handle_stream_response(response, model):
     return generate()
 
 def initialization():
-    sso_array = os.environ.get("SSO", "").split(',')
+    sso_array = os.getenv("SSO", "").split(',')
     logger.info("开始加载令牌", "Server")
     for sso in sso_array:
         if sso:
@@ -872,11 +984,76 @@ logger.info("初始化完成", "Server")
 
 app = Flask(__name__)
 app.wsgi_app = ProxyFix(app.wsgi_app)
+app.secret_key = os.getenv('FLASK_SECRET_KEY') or secrets.token_hex(16)
+app.json.sort_keys = False
+
+@app.route('/manager/login', methods=['GET', 'POST'])
+def manager_login():
+    if CONFIG["ADMIN"]["MANAGER_SWITCH"]:
+        if request.method == 'POST':
+            password = request.form.get('password')
+            if password == CONFIG["ADMIN"]["PASSWORD"]:
+                session['is_logged_in'] = True
+                return redirect('/manager')
+            return render_template('login.html', error=True)
+        return render_template('login.html', error=False)
+    else:
+        return redirect('/')
+
+def check_auth():
+    return session.get('is_logged_in', False)
+
+@app.route('/manager')
+def manager():
+    if not check_auth():
+        return redirect('/manager/login')
+    return render_template('manager.html')
+
+@app.route('/manager/api/get')
+def get_manager_tokens():
+    if not check_auth():
+        return jsonify({"error": "Unauthorized"}), 401
+    return jsonify(token_manager.get_token_status_map())
 
+@app.route('/manager/api/add', methods=['POST'])
+def add_manager_token():
+    if not check_auth():
+        return jsonify({"error": "Unauthorized"}), 401
+    try:
+        sso = request.json.get('sso')
+        if not sso:
+            return jsonify({"error": "SSO token is required"}), 400
+        token_manager.add_token(f"sso-rw={sso};sso={sso}")
+        return jsonify({"success": True})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+@app.route('/manager/api/delete', methods=['POST'])
+def delete_manager_token():
+    if not check_auth():
+        return jsonify({"error": "Unauthorized"}), 401
+    try:
+        sso = request.json.get('sso')
+        if not sso:
+            return jsonify({"error": "SSO token is required"}), 400
+        token_manager.delete_token(f"sso-rw={sso};sso={sso}")
+        return jsonify({"success": True})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+    
+@app.route('/manager/api/cf_clearance', methods=['POST'])   
+def setCf_Manager_clearance():
+    if not check_auth():
+        return jsonify({"error": "Unauthorized"}), 401
+    try:
+        cf_clearance = request.json.get('cf_clearance')
+        if not cf_clearance:
+            return jsonify({"error": "cf_clearance is required"}), 400
+        CONFIG["SERVER"]['CF_CLEARANCE'] = cf_clearance
+        return jsonify({"success": True})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
-@app.before_request
-def log_request_info():
-    logger.info(f"{request.method} {request.path}", "Request")
 
 @app.route('/get/tokens', methods=['GET'])
 def get_tokens():
@@ -885,7 +1062,6 @@ def get_tokens():
         return jsonify({"error": '自定义的SSO令牌模式无法获取轮询sso令牌状态'}), 403
     elif auth_token != CONFIG["API"]["API_KEY"]:
         return jsonify({"error": 'Unauthorized'}), 401
-
     return jsonify(token_manager.get_token_status_map())
 
 @app.route('/add/token', methods=['POST'])
@@ -903,7 +1079,20 @@ def add_token():
     except Exception as error:
         logger.error(str(error), "Server")
         return jsonify({"error": '添加sso令牌失败'}), 500
-
+    
+@app.route('/set/cf_clearance', methods=['POST'])
+def setCf_clearance():
+    auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
+    if auth_token != CONFIG["API"]["API_KEY"]:
+        return jsonify({"error": 'Unauthorized'}), 401
+    try:
+        cf_clearance = request.json.get('cf_clearance')
+        CONFIG["SERVER"]['CF_CLEARANCE'] = cf_clearance
+        return jsonify({"message": '设置cf_clearance成功'}), 200
+    except Exception as error:
+        logger.error(str(error), "Server")
+        return jsonify({"error": '设置cf_clearance失败'}), 500
+    
 @app.route('/delete/token', methods=['POST'])
 def delete_token():
     auth_token = request.headers.get('Authorization', '').replace('Bearer ', '')
@@ -937,6 +1126,7 @@ def get_models():
 
 @app.route('/v1/chat/completions', methods=['POST'])
 def chat_completions():
+    response_status_code = 500
     try:
         auth_token = request.headers.get('Authorization',
                                          '').replace('Bearer ', '')
@@ -956,63 +1146,69 @@ def chat_completions():
         retry_count = 0
         grok_client = GrokApiClient(model)
         request_payload = grok_client.prepare_chat_request(data)
+        
 
         while retry_count < CONFIG["RETRY"]["MAX_ATTEMPTS"]:
             retry_count += 1
-            CONFIG["API"]["SIGNATURE_COOKIE"] = Utils.create_auth_headers(
-                model)
+            CONFIG["API"]["SIGNATURE_COOKIE"] = Utils.create_auth_headers(model)
 
             if not CONFIG["API"]["SIGNATURE_COOKIE"]:
                 raise ValueError('该模型无可用令牌')
 
             logger.info(
-                f"当前令牌: {json.dumps(CONFIG['API']['SIGNATURE_COOKIE'], indent=2)}",
-                "Server")
+                f"当前令牌: {json.dumps(CONFIG['API']['SIGNATURE_COOKIE'], indent=2)}","Server")
             logger.info(
-                f"当前可用模型的全部可用数量: {json.dumps(token_manager.get_remaining_token_request_capacity(), indent=2)}",
-                "Server")
-
+                f"当前可用模型的全部可用数量: {json.dumps(token_manager.get_remaining_token_request_capacity(), indent=2)}","Server")
+            
+            if CONFIG['SERVER']['CF_CLEARANCE']:
+                CONFIG["SERVER"]['COOKIE'] = f"{CONFIG['API']['SIGNATURE_COOKIE']};{CONFIG['SERVER']['CF_CLEARANCE']}" 
+            else:
+                CONFIG["SERVER"]['COOKIE'] = CONFIG['API']['SIGNATURE_COOKIE']
+            logger.info(json.dumps(request_payload,indent=2),"Server")
             try:
                 proxy_options = Utils.get_proxy_options()
                 response = curl_requests.post(
                     f"{CONFIG['API']['BASE_URL']}/rest/app-chat/conversations/new",
                     headers={
-                        **DEFAULT_HEADERS, "Cookie":
-                        CONFIG["API"]["SIGNATURE_COOKIE"]
+                        **DEFAULT_HEADERS, 
+                        "Cookie":CONFIG["SERVER"]['COOKIE']
                     },
-                    data=json.dumps(request_payload),
+                    json=request_payload,
                     impersonate="chrome133a",
+                    verify=False,
                     stream=True,
                     **proxy_options)
+                logger.info(CONFIG["SERVER"]['COOKIE'],"Server")
                 if response.status_code == 200:
+                    response_status_code = 200
                     logger.info("请求成功", "Server")
-                    logger.info(
-                        f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",
-                        "Server")
+                    logger.info(f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}","Server")
 
                     try:
                         if stream:
                             return Response(stream_with_context(
-                                handle_stream_response(response, model)),
-                                            content_type='text/event-stream')
+                                handle_stream_response(response, model)),content_type='text/event-stream')
                         else:
-                            content = handle_non_stream_response(
-                                response, model)
+                            content = handle_non_stream_response(response, model)
                             return jsonify(
-                                MessageProcessor.create_chat_response(
-                                    content, model))
+                                MessageProcessor.create_chat_response(content, model))
 
                     except Exception as error:
                         logger.error(str(error), "Server")
                         if CONFIG["API"]["IS_CUSTOM_SSO"]:
                             raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
-
-                        token_manager.remove_token_from_model(
-                            model, CONFIG["API"]["SIGNATURE_COOKIE"])
+                        token_manager.remove_token_from_model(model, CONFIG["API"]["SIGNATURE_COOKIE"])
                         if token_manager.get_token_count_for_model(model) == 0:
                             raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")
-
+                elif response.status_code == 403:
+                    response_status_code = 403
+                    token_manager.reduce_token_request_count(model,1)#重置去除当前因为错误未成功请求的次数，确保不会因为错误未成功请求的次数导致次数上限
+                    if token_manager.get_token_count_for_model(model) == 0:
+                        raise ValueError(f"{model} 次数已达上限，请切换其他模型或者重新对话")
+                    raise ValueError(f"IP暂时被封无法破盾，请稍后重试或者更换ip")
                 elif response.status_code == 429:
+                    response_status_code = 429
+                    token_manager.reduce_token_request_count(model,1)
                     if CONFIG["API"]["IS_CUSTOM_SSO"]:
                         raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
 
@@ -1025,10 +1221,8 @@ def chat_completions():
                     if CONFIG["API"]["IS_CUSTOM_SSO"]:
                         raise ValueError(f"自定义SSO令牌当前模型{model}的请求次数已失效")
 
-                    logger.error(f"令牌异常错误状态!status: {response.status_code}",
-                                 "Server")
-                    token_manager.remove_token_from_model(
-                        model, CONFIG["API"]["SIGNATURE_COOKIE"])
+                    logger.error(f"令牌异常错误状态!status: {response.status_code}","Server")
+                    token_manager.remove_token_from_model(model, CONFIG["API"]["SIGNATURE_COOKIE"])
                     logger.info(
                         f"当前{model}剩余可用令牌数: {token_manager.get_token_count_for_model(model)}",
                         "Server")
@@ -1038,8 +1232,10 @@ def chat_completions():
                 if CONFIG["API"]["IS_CUSTOM_SSO"]:
                     raise
                 continue
-
-        raise ValueError('当前模型所有令牌都已耗尽')
+        if response_status_code == 403:
+            raise ValueError('IP暂时被封无法破盾，请稍后重试或者更换ip')
+        elif response_status_code == 500:
+            raise ValueError('当前模型所有令牌暂无可用，请稍后重试')    
 
     except Exception as error:
         logger.error(str(error), "ChatAPI")
@@ -1047,7 +1243,7 @@ def chat_completions():
             {"error": {
                 "message": str(error),
                 "type": "server_error"
-            }}), 500
+            }}), response_status_code
 
 @app.route('/', defaults={'path': ''})
 @app.route('/<path:path>')
@@ -1062,4 +1258,4 @@ if __name__ == '__main__':
         host='0.0.0.0',
         port=CONFIG["SERVER"]["PORT"],
         debug=False
-    )
+    )