Update app.py

app.py

@@ -28,6 +28,11 @@ _API_BASE_URL = "https://spuckhogycrxcbomznwo.supabase.co"
 _USER_AGENT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36'
 _PASTE_API_URL = "https://page.zhoudan.icu/api/paste/b40v96oX"
 
+# 从环境变量获取API密钥
+API_KEY = os.getenv('API_KEY')
+if not API_KEY:
+    raise ValueError("API_KEY environment variable must be set")
+
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -42,6 +47,25 @@ if not NOTDIAMOND_IP:
     logger.error("NOTDIAMOND_IP environment variable is not set!")
     raise ValueError("NOTDIAMOND_IP must be set")
 
+# API密钥验证装饰器
+def require_api_key(f):
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        auth_header = request.headers.get('Authorization')
+        if not auth_header:
+            return jsonify({'error': 'No API key provided'}), 401
+        
+        try:
+            # 从 Bearer token 中提取API密钥
+            provided_key = auth_header.split('Bearer ')[-1].strip()
+            if provided_key != API_KEY:
+                return jsonify({'error': 'Invalid API key'}), 401
+        except Exception:
+            return jsonify({'error': 'Invalid Authorization header format'}), 401
+            
+        return f(*args, **kwargs)
+    return decorated_function
+
 refresh_token_cache = TTLCache(maxsize=1000, ttl=3600)
 headers_cache = TTLCache(maxsize=1, ttl=3600)  # 1小时过期
 token_refresh_lock = threading.Lock()
@@ -470,12 +494,16 @@ def before_request():
         multi_auth_manager = None
 
 @app.route('/', methods=['GET'])
+@require_api_key
 def root():
     return jsonify({
         "service": "AI Chat Completion Proxy",
         "usage": {
             "endpoint": "/ai/v1/chat/completions",
             "method": "POST",
+            "headers": {
+                "Authorization": "Bearer YOUR_API_KEY"
+            },
             "body": {
                 "model": "One of: " + ", ".join(MODEL_INFO.keys()),
                 "messages": [
@@ -487,10 +515,11 @@ def root():
             }
         },
         "availableModels": list(MODEL_INFO.keys()),
-        "note": "Credentials are automatically fetched from the API."
+        "note": "API key authentication is required for all endpoints."
     })
 
 @app.route('/ai/v1/models', methods=['GET'])
+@require_api_key
 def proxy_models():
     """返回可用模型列表。"""
     models = [
@@ -510,6 +539,7 @@ def proxy_models():
     })
 
 @app.route('/ai/v1/chat/completions', methods=['POST'])
+@require_api_key
 def handle_request():
     global multi_auth_manager
     if not multi_auth_manager: