Update app.py

app.py

@@ -5,85 +5,133 @@ import time
 import uuid
 import requests
 import os
+import base64
+from functools import wraps
+import logging
 
 app = Flask(__name__)
-CORS(app, origins=["https://catptain-coze-api-01.hf.space"] + [
-    "https://x-raremeta.com",
-    "https://cybercity.top",
-    "https://play-1.x-raremeta.com",
-    "https://play.cybercity.top",
-    "https://play.x-raremeta.com",
-    "https://www.x-raremeta.com",
-    "https://www.cybercity.top"
-])
+CORS(app, origins=os.getenv('ALLOWED_ORIGINS', 'https://cybercity.top').split(','))
 
-# 你的配置信息
-CLIENT_ID = "1243934778935"
-PRIVATE_KEY_FILE_PATH = "private_key.pem"
-KID = "tlrohMMZyKMrrpP3GtxF_3_cerDhVIMINs0LOW91m7w"
-VALIDATION_TOKEN = "cybercity2025"
+# 环境变量配置
+CLIENT_ID = os.getenv('COZE_CLIENT_ID', '1243934778935')
+KID = os.getenv('COZE_KID', 'tlrohMMZyKMrrpP3GtxF_3_cerDhVIMINs0LOW91m7w')
+PRIVATE_KEY = os.getenv('COZE_PRIVATE_KEY').replace('\\n', '\n')  # 从环境变量获取并格式化
+CLIENT_SECRET = os.getenv('COZE_CLIENT_SECRET', 'your_client_secret')
 
-def generate_jwt(client_id, private_key, kid):
+# 日志配置
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# JWT缓存机制（简易内存缓存）
+jwt_cache = {'token': None, 'exp': 0}
+
+def validate_basic_auth(auth_header):
+    """实现RFC6749标准的Basic认证验证[10](@ref)"""
+    if not auth_header or not auth_header.startswith('Basic '):
+        return False
+    try:
+        credentials = base64.b64decode(auth_header[6:]).decode('utf-8')
+        client_id, client_secret = credentials.split(':', 1)
+        return client_id == CLIENT_ID and client_secret == CLIENT_SECRET
+    except Exception as e:
+        logger.error(f"Basic auth validation failed: {str(e)}")
+        return False
+
+def generate_jwt():
+    """生成符合RFC7519标准的JWT[1,3](@ref)"""
+    current_time = int(time.time())
+    payload = {
+        "iss": CLIENT_ID,
+        "sub": CLIENT_ID,  # 必须包含sub字段[6](@ref)
+        "aud": "https://api.coze.cn",  # 精确的URI格式
+        "iat": current_time,
+        "exp": current_time + 3600,
+        "jti": uuid.uuid4().hex,
+        "connector_id": CLIENT_ID,  # 统一使用client_id
+        "user_id": CLIENT_ID
+    }
+    
     header = {
         "alg": "RS256",
         "typ": "JWT",
-        "kid": kid
-    }
-    payload = {
-        "iss": client_id,
-        "aud": "api.coze.cn",
-        "iat": int(time.time()),
-        "exp": int(time.time()) + 3600,  # JWT 有效期为 1 小时
-        "jti": uuid.uuid4().hex,  # 防止重放攻击
-        "connector_id": "3723409317963603",  # 根据实际要求设置
-        "user_id": "3723409317963603",  # 根据实际要求设置
-        "UID": "3723409317963603",  # 根据实际要求设置
+        "kid": KID
     }
-    return jwt.encode(payload, private_key, algorithm="RS256", headers=header)
+    
+    try:
+        return jwt.encode(payload, PRIVATE_KEY, algorithm="RS256", headers=header)
+    except jwt.PyJWTError as e:
+        logger.error(f"JWT generation failed: {str(e)}")
+        raise
 
 def get_access_token(jwt_token):
+    """获取访问令牌（带重试机制）[3](@ref)"""
     url = "https://api.coze.cn/api/permission/oauth2/token"
     data = {
-        "duration_seconds": 86399,
-        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer"
+        "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
+        "duration_seconds": 86399
     }
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {jwt_token}"
     }
-    response = requests.post(url, json=data, headers=headers)
-    return response.json()
+    
+    try:
+        response = requests.post(url, json=data, headers=headers, timeout=10)
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Access token request failed: {str(e)}")
+        return {"error": "coze_api_error"}
 
-# 添加根路由，帮助Hugging Face识别应用已经就绪
-@app.route('/', methods=['GET'])
-def index():
-    return jsonify({"status": "Service is running", "endpoints": ["/get_token"]}), 200
+# 健康检查端点
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({"status": "healthy", "timestamp": int(time.time())}), 200
 
-# 正确的token获取路由
-@app.route('/get_token', methods=['GET'])
-def get_token_from_flask():
-    auth_header = request.headers.get('Authorization')
-    if auth_header != VALIDATION_TOKEN:
-        return jsonify({"error": "Invalid authorization token"}), 401
+# 令牌获取端点
+@app.route('/api/token', methods=['POST'])
+def get_coze_token():
+    # Basic认证验证
+    if not validate_basic_auth(request.headers.get('Authorization')):
+        return jsonify({"error": "invalid_client"}), 401
     
-    try:
-        with open(PRIVATE_KEY_FILE_PATH, "r") as f:
-            private_key = f.read()
-        
-        jwt_token = generate_jwt(CLIENT_ID, private_key, KID)
-        response = get_access_token(jwt_token)
-        
-        if "access_token" in response:
-            return jsonify({
-                "access_token": response["access_token"],
-                "expires_in": response["expires_in"]
+    # 检查缓存中的有效JWT
+    current_time = time.time()
+    if jwt_cache['exp'] > current_time + 300:  # 有效期剩余超过5分钟时复用
+        cached_token = jwt_cache['token']
+    else:
+        try:
+            cached_token = generate_jwt()
+            jwt_cache.update({
+                'token': cached_token,
+                'exp': current_time + 3600
             })
-        else:
-            return jsonify({"error": "Failed to get access token", "details": response}), 500
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        except Exception as e:
+            return jsonify({"error": "jwt_generation_failed"}), 500
+    
+    # 获取访问令牌
+    token_response = get_access_token(cached_token)
+    if 'error' in token_response:
+        return jsonify({
+            "error": "coze_oauth_error",
+            "details": token_response.get('error_description')
+        }), 502
+    
+    return jsonify({
+        "access_token": token_response['access_token'],
+        "expires_in": token_response['expires_in'],
+        "token_type": "Bearer"
+    })
+
+# 错误处理
+@app.errorhandler(404)
+def not_found(error):
+    return jsonify({"error": "endpoint_not_found"}), 404
+
+@app.errorhandler(500)
+def internal_error(error):
+    return jsonify({"error": "internal_server_error"}), 500
 
-# 使用环境变量设置端口
-port = int(os.environ.get("PORT", 7860))
 if __name__ == '__main__':
-    app.run(host="0.0.0.0", port=port)
+    port = int(os.getenv('PORT', 7860))
+    app.run(host='0.0.0.0', port=port, debug=os.getenv('DEBUG', 'false').lower() == 'true')