Upload utils.py

src/utils.py

@@ -0,0 +1,208 @@
+import os
+import uuid
+import aiohttp
+import aiofiles
+import logging
+import ssl
+from urllib.parse import urlparse
+from .config import Config
+
+# 初始化日志
+logger = logging.getLogger("sora-api.utils")
+
+# 图片本地化调试开关
+IMAGE_DEBUG = os.getenv("IMAGE_DEBUG", "").lower() in ("true", "1", "yes")
+
+# 修复Python 3.11之前版本中HTTPS代理处理HTTPS请求的问题
+# 参考: https://docs.aiohttp.org/en/stable/client_advanced.html#proxy-support
+try:
+    import aiohttp.connector
+    orig_create_connection = aiohttp.connector.TCPConnector._create_connection
+
+    async def patched_create_connection(self, req, traces, timeout):
+        if req.ssl and req.proxy and req.proxy.scheme == 'https':
+            # 为代理连接创建SSL上下文
+            proxy_ssl = ssl.create_default_context()
+            req.proxy_ssl = proxy_ssl
+            
+            if IMAGE_DEBUG:
+                logger.debug("已应用HTTPS代理补丁")
+        
+        return await orig_create_connection(self, req, traces, timeout)
+    
+    # 应用猴子补丁
+    aiohttp.connector.TCPConnector._create_connection = patched_create_connection
+    
+    if IMAGE_DEBUG:
+        logger.debug("已启用aiohttp HTTPS代理支持补丁")
+except Exception as e:
+    logger.warning(f"应用HTTPS代理补丁失败: {e}")
+
+async def download_and_save_image(image_url: str) -> str:
+    """
+    下载图片并保存到本地
+    
+    Args:
+        image_url: 图片URL
+        
+    Returns:
+        本地化后的图片URL
+    """
+    # 如果未启用本地化或URL已经是本地路径，直接返回
+    if not Config.IMAGE_LOCALIZATION:
+        if IMAGE_DEBUG:
+            logger.debug(f"图片本地化未启用，返回原始URL: {image_url}")
+        return image_url
+    
+    # 检查是否已经是本地URL
+    # 准备URL检测所需的信息
+    parsed_base_url = urlparse(Config.BASE_URL)
+    base_path = parsed_base_url.path.rstrip('/')
+    
+    # 直接检查常见的图片URL模式
+    local_url_patterns = [
+        "/images/",
+        "/static/images/",
+        f"{base_path}/images/",
+        f"{base_path}/static/images/"
+    ]
+    
+    # 检查是否有自定义前缀
+    if Config.STATIC_PATH_PREFIX:
+        prefix = Config.STATIC_PATH_PREFIX
+        if not prefix.startswith('/'):
+            prefix = f"/{prefix}"
+        local_url_patterns.append(f"{prefix}/images/")
+        local_url_patterns.append(f"{base_path}{prefix}/images/")
+    
+    # 检查是否符合任何本地URL模式
+    is_local_url = any(pattern in image_url for pattern in local_url_patterns)
+    
+    if is_local_url:
+        if IMAGE_DEBUG:
+            logger.debug(f"URL已是本地图片路径: {image_url}")
+        
+        # 如果是相对路径，补充完整的URL
+        if image_url.startswith("/"):
+            return f"{parsed_base_url.scheme}://{parsed_base_url.netloc}{image_url}"
+        return image_url
+    
+    try:
+        # 生成文件名和保存路径
+        parsed_url = urlparse(image_url)
+        file_extension = os.path.splitext(parsed_url.path)[1] or ".png"
+        filename = f"{uuid.uuid4()}{file_extension}"
+        save_path = os.path.join(Config.IMAGE_SAVE_DIR, filename)
+        
+        # 确保保存目录存在
+        os.makedirs(Config.IMAGE_SAVE_DIR, exist_ok=True)
+        
+        if IMAGE_DEBUG:
+            logger.debug(f"下载图片: {image_url} -> {save_path}")
+        
+        # 配置代理
+        proxy = None
+        if Config.PROXY_HOST and Config.PROXY_PORT:
+            proxy_auth = None
+            if Config.PROXY_USER and Config.PROXY_PASS:
+                proxy_auth = aiohttp.BasicAuth(Config.PROXY_USER, Config.PROXY_PASS)
+            
+            proxy_url = f"http://{Config.PROXY_HOST}:{Config.PROXY_PORT}"
+            if IMAGE_DEBUG:
+                auth_info = f" (使用认证)" if proxy_auth else ""
+                logger.debug(f"使用代理: {proxy_url}{auth_info}")
+            proxy = proxy_url
+        
+        # 下载图片
+        timeout = aiohttp.ClientTimeout(total=60)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            # 创建请求参数
+            request_kwargs = {"timeout": 30}
+            if proxy:
+                request_kwargs["proxy"] = proxy
+                if Config.PROXY_USER and Config.PROXY_PASS:
+                    request_kwargs["proxy_auth"] = aiohttp.BasicAuth(Config.PROXY_USER, Config.PROXY_PASS)
+            
+            async with session.get(image_url, **request_kwargs) as response:
+                if response.status != 200:
+                    logger.warning(f"下载失败，状态码: {response.status}, URL: {image_url}")
+                    return image_url
+                
+                content = await response.read()
+                if not content:
+                    logger.warning("下载内容为空")
+                    return image_url
+                
+                # 保存图片
+                async with aiofiles.open(save_path, "wb") as f:
+                    await f.write(content)
+        
+        # 检查文件是否成功保存
+        if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:
+            logger.warning(f"图片保存失败: {save_path}")
+            if os.path.exists(save_path):
+                os.remove(save_path)
+            return image_url
+        
+        # 返回本地URL
+        # 获取文件名
+        filename = os.path.basename(save_path)
+        
+        # 解析基础URL
+        parsed_base_url = urlparse(Config.BASE_URL)
+        base_path = parsed_base_url.path.rstrip('/')
+        
+        # 使用固定的图片URL格式
+        relative_url = f"/images/{filename}"
+        
+        # 如果设置了STATIC_PATH_PREFIX，优先使用该前缀
+        if Config.STATIC_PATH_PREFIX:
+            prefix = Config.STATIC_PATH_PREFIX
+            if not prefix.startswith('/'):
+                prefix = f"/{prefix}"
+            relative_url = f"{prefix}/images/{filename}"
+        
+        # 如果BASE_URL有子路径，添加到相对路径前
+        if base_path:
+            relative_url = f"{base_path}{relative_url}"
+        
+        # 处理重复斜杠
+        while "//" in relative_url:
+            relative_url = relative_url.replace("//", "/")
+        
+        # 生成完整URL
+        full_url = f"{parsed_base_url.scheme}://{parsed_base_url.netloc}{relative_url}"
+        
+        if IMAGE_DEBUG:
+            logger.debug(f"图片保存成功: {full_url}")
+            logger.debug(f"图片保存路径: {save_path}")
+        
+        return full_url
+    except Exception as e:
+        logger.error(f"图片下载失败: {str(e)}", exc_info=IMAGE_DEBUG)
+        return image_url
+
+async def localize_image_urls(image_urls: list) -> list:
+    """
+    批量将图片URL本地化
+    
+    Args:
+        image_urls: 图片URL列表
+        
+    Returns:
+        本地化后的URL列表
+    """
+    if not Config.IMAGE_LOCALIZATION or not image_urls:
+        return image_urls
+    
+    if IMAGE_DEBUG:
+        logger.debug(f"本地化 {len(image_urls)} 个URL: {image_urls}")
+    else:
+        logger.info(f"本地化 {len(image_urls)} 个图片")
+    
+    localized_urls = []
+    for url in image_urls:
+        localized_url = await download_and_save_image(url)
+        localized_urls.append(localized_url)
+    
+    return localized_urls