Update app.py

app.py

@@ -26,6 +26,7 @@ CONTENT_TYPE_EVENT_STREAM = 'text/event-stream'
 _BASE_URL = "https://chat.notdiamond.ai"
 _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"
 
 app = Flask(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -34,103 +35,23 @@ CORS(app, resources={r"/*": {"origins": "*"}})
 executor = ThreadPoolExecutor(max_workers=10)
 
 proxy_url = os.getenv('PROXY_URL')
+NOTDIAMOND_IP = os.getenv('NOTDIAMOND_IP')
 NOTDIAMOND_DOMAIN = os.getenv('NOTDIAMOND_DOMAIN')
-if not NOTDIAMOND_DOMAIN:
-    logger.error("NOTDIAMOND_DOMAIN environment variable is not set!")
-    raise ValueError("NOTDIAMOND_DOMAIN must be set")
-
-# IP缓存配置
-IP_CACHE_TTL = 3600  # IP缓存1小时过期
-MAX_IP_RETRIES = 3   # 最大重试次数
-RETRY_DELAY = 5      # 重试延迟（秒）
-
-# 动态IP缓存
-ip_cache = TTLCache(maxsize=1, ttl=IP_CACHE_TTL)
-ip_refresh_lock = threading.Lock()
-
-def resolve_domain_ip(domain: str, retry_count: int = 0) -> str:
-    """
-    解析域名对应的IP地址，带重试机制
-    
-    Args:
-        domain: 要解析的域名
-        retry_count: 当前重试次数
-        
-    Returns:
-        str: 解析得到的IP地址
-        
-    Raises:
-        Exception: 当所有重试都失败时抛出异常
-    """
-    try:
-        logger.info(f"Attempting to resolve IP for domain: {domain} (attempt {retry_count + 1})")
-        ip = socket.gethostbyname(domain)
-        logger.info(f"Successfully resolved IP: {ip} for domain: {domain}")
-        return ip
-    except socket.gaierror as e:
-        if retry_count < MAX_IP_RETRIES - 1:
-            logger.warning(f"DNS resolution failed for {domain} (attempt {retry_count + 1}): {e}")
-            time.sleep(RETRY_DELAY)
-            return resolve_domain_ip(domain, retry_count + 1)
-        else:
-            logger.error(f"DNS resolution failed after {MAX_IP_RETRIES} attempts for {domain}: {e}")
-            raise Exception(f"Failed to resolve IP for {domain} after {MAX_IP_RETRIES} attempts")
 
-def get_current_ip() -> str:
-    """
-    获取当前域名对应的IP地址，使用缓存机制
-    
-    Returns:
-        str: 当前域名对应的IP地址
-        
-    Raises:
-        Exception: 当IP解析失败时抛出异常
-    """
-    try:
-        with ip_refresh_lock:
-            if 'current_ip' not in ip_cache:
-                logger.info(f"IP cache miss for domain: {NOTDIAMOND_DOMAIN}")
-                ip = resolve_domain_ip(NOTDIAMOND_DOMAIN)
-                ip_cache['current_ip'] = ip
-                logger.info(f"Updated IP cache with: {ip}")
-            return ip_cache['current_ip']
-    except Exception as e:
-        logger.error(f"Failed to get current IP: {e}")
-        raise
+if not NOTDIAMOND_IP:
+    logger.error("NOTDIAMOND_IP environment variable is not set!")
+    raise ValueError("NOTDIAMOND_IP must be set")
 
-def refresh_ip_cache() -> None:
-    """
-    强制刷新IP缓存
-    """
-    try:
-        with ip_refresh_lock:
-            ip_cache.clear()
-            logger.info("IP cache cleared for refresh")
-            _ = get_current_ip()  # 触发新的IP解析
-            logger.info("IP cache refreshed successfully")
-    except Exception as e:
-        logger.error(f"Error refreshing IP cache: {e}")
-        raise
+refresh_token_cache = TTLCache(maxsize=1000, ttl=3600)
+headers_cache = TTLCache(maxsize=1, ttl=3600)  # 1小时过期
+token_refresh_lock = threading.Lock()
 
 # 自定义连接函数
 def patched_create_connection(address, *args, **kwargs):
-    """
-    自定义连接函数，用于替换默认的连接函数
-    支持动态IP解析和自动重试
-    """
     host, port = address
     if host == NOTDIAMOND_DOMAIN:
-        try:
-            current_ip = get_current_ip()
-            logger.info(f"Using resolved IP: {current_ip} for connection to {NOTDIAMOND_DOMAIN}")
-            return create_connection((current_ip, port), *args, **kwargs)
-        except Exception as e:
-            logger.error(f"Connection failed using resolved IP: {e}")
-            # 如果连接失败，尝试刷新IP缓存并重试
-            refresh_ip_cache()
-            current_ip = get_current_ip()
-            logger.info(f"Retrying connection with fresh IP: {current_ip}")
-            return create_connection((current_ip, port), *args, **kwargs)
+        logger.info(f"Connecting to {NOTDIAMOND_DOMAIN} using IP: {NOTDIAMOND_IP}")
+        return create_connection((NOTDIAMOND_IP, port), *args, **kwargs)
     return create_connection(address, *args, **kwargs)
 
 # 替换 urllib3 的默认连接函数
@@ -509,23 +430,25 @@ def generate_stream_response(response, model, prompt_tokens):
     yield "data: [DONE]\n\n"
 
 def get_auth_credentials():
-    """从请求头中获取多个认证凭据"""
-    auth_header = request.headers.get('Authorization')
-    if not auth_header or not auth_header.startswith('Bearer '):
-        logger.error("Authorization header is missing or invalid")
-        return []
-    
+    """从API获取认证凭据"""
     try:
-        credentials_string = auth_header.split('Bearer ')[1]
-        credentials_list = credentials_string.split(';')
-        parsed_credentials = []
-        for cred in credentials_list:
-            email, password = cred.split('|')
-            parsed_credentials.append((email.strip(), password.strip()))
-        logger.info(f"Extracted {len(parsed_credentials)} sets of credentials")
-        return parsed_credentials
+        headers = {
+            'accept': '*/*',
+            'accept-language': 'zh-CN,zh;q=0.9',
+            'user-agent': _USER_AGENT,
+            'x-password': '321'
+        }
+        response = requests.get(_PASTE_API_URL, headers=headers)
+        if response.status_code == 200:
+            data = response.json()
+            if data.get('status') == 'success' and data.get('content'):
+                credentials_string = data['content']
+                email, password = credentials_string.split('|')
+                return [(email.strip(), password.strip())]
+        logger.error(f"Failed to get credentials from API: {response.status_code}")
+        return []
     except Exception as e:
-        logger.error(f"Error parsing Authorization header: {e}")
+        logger.error(f"Error getting credentials from API: {e}")
         return []
 
 @app.before_request
@@ -544,10 +467,6 @@ def root():
         "usage": {
             "endpoint": "/ai/v1/chat/completions",
             "method": "POST",
-            "headers": {
-                "Content-Type": "application/json",
-                "Authorization": "Bearer YOUR_EMAIL1|YOUR_PASSWORD1;YOUR_EMAIL2|YOUR_PASSWORD2"
-            },
             "body": {
                 "model": "One of: " + ", ".join(MODEL_INFO.keys()),
                 "messages": [
@@ -559,7 +478,7 @@ def root():
             }
         },
         "availableModels": list(MODEL_INFO.keys()),
-        "note": "Replace YOUR_EMAIL and YOUR_PASSWORD with your actual Not Diamond credentials."
+        "note": "Credentials are automatically fetched from the API."
     })
 
 @app.route('/ai/v1/models', methods=['GET'])
@@ -744,15 +663,6 @@ def health_check():
                 if current_time.tm_hour == 0 and current_time.tm_min == 0:
                     multi_auth_manager.reset_all_model_status()
                     logger.info("Reset model status for all accounts")
-                
-                # 每小时刷新一次IP缓存
-                if current_time.tm_min == 0:
-                    try:
-                        ip_cache.clear()
-                        logger.info("Cleared IP cache for periodic refresh")
-                    except Exception as e:
-                        logger.error(f"Error clearing IP cache: {e}")
-                
         except Exception as e:
             logger.error(f"Health check error: {e}")
         time.sleep(60)  # 每分钟检查一次