Upload image_uploader.py

src/image_uploader.py

@@ -0,0 +1,242 @@
+import os
+import requests
+import aiohttp
+import logging
+import json
+import aiofiles
+import base64
+from urllib.parse import urlparse
+from .config import Config
+
+# 初始化日志
+logger = logging.getLogger("sora-api.image_uploader")
+
+class PicGoUploader:
+    """PicGo图床上传器"""
+    
+    def __init__(self, picgo_url=None, api_key=None):
+        """
+        初始化PicGo上传器
+        
+        Args:
+            picgo_url: PicGo服务器URL，默认为 http://127.0.0.1:36677
+            api_key: PicGo API密钥（如果启用了）
+        """
+        self.picgo_url = picgo_url or os.getenv("PICGO_URL", "http://127.0.0.1:36677")
+        self.api_key = api_key or os.getenv("PICGO_API_KEY", "")
+        self.enabled = os.getenv("ENABLE_PICGO", "").lower() in ("true", "1", "yes")
+        
+        # 日志输出
+        if self.enabled:
+            logger.info(f"PicGo上传已启用，服务器地址: {self.picgo_url}")
+        else:
+            logger.info("PicGo上传未启用")
+    
+    async def upload(self, file_path):
+        """
+        上传图片到PicGo
+        
+        Args:
+            file_path: 图片文件路径
+            
+        Returns:
+            成功: 图片URL
+            失败: None
+        """
+        if not self.enabled:
+            logger.debug("PicGo上传未启用，跳过上传")
+            return None
+            
+        if not os.path.exists(file_path):
+            logger.error(f"文件不存在: {file_path}")
+            return None
+            
+        try:
+            # 读取图片文件
+            async with aiofiles.open(file_path, "rb") as f:
+                file_content = await f.read()
+                
+            # 编码为Base64
+            base64_content = base64.b64encode(file_content).decode('utf-8')
+            
+            # 准备上传数据
+            data = {
+                "list": [{
+                    "data": f"data:image/{os.path.splitext(file_path)[1][1:]};base64,{base64_content}"
+                }]
+            }
+            
+            # 准备请求头
+            headers = {"Content-Type": "application/json"}
+            if self.api_key:
+                headers["picgo-api-key"] = self.api_key
+                
+            # 发送请求
+            timeout = aiohttp.ClientTimeout(total=60)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(
+                    f"{self.picgo_url}/upload",
+                    headers=headers,
+                    json=data
+                ) as response:
+                    result = await response.json()
+                    
+            # 处理响应
+            if result.get("success"):
+                image_url = result.get("result")[0]
+                logger.info(f"图片已成功上传到PicGo: {image_url}")
+                return image_url
+            else:
+                logger.error(f"PicGo上传失败: {result.get('message', '未知错误')}")
+                return None
+                
+        except Exception as e:
+            logger.error(f"PicGo上传过程出错: {str(e)}")
+            return None
+
+class WebDAVUploader:
+    """WebDAV上传器"""
+    
+    def __init__(self, webdav_url=None, username=None, password=None):
+        """
+        初始化WebDAV上传器
+        
+        Args:
+            webdav_url: WebDAV服务器URL
+            username: WebDAV用户名
+            password: WebDAV密码
+        """
+        self.webdav_url = webdav_url or os.getenv("WEBDAV_URL", "")
+        self.username = username or os.getenv("WEBDAV_USERNAME", "")
+        self.password = password or os.getenv("WEBDAV_PASSWORD", "")
+        self.base_path = os.getenv("WEBDAV_BASE_PATH", "/sora-images")
+        self.enabled = os.getenv("ENABLE_WEBDAV", "").lower() in ("true", "1", "yes")
+        self.public_url = os.getenv("WEBDAV_PUBLIC_URL", "")
+        
+        # 确保base_path以/开头
+        if self.base_path and not self.base_path.startswith('/'):
+            self.base_path = f"/{self.base_path}"
+            
+        # 日志输出
+        if self.enabled:
+            if not self.webdav_url:
+                logger.warning("WebDAV上传已启用，但未设置WebDAV URL")
+                self.enabled = False
+            else:
+                logger.info(f"WebDAV上传已启用，服务器地址: {self.webdav_url}")
+        else:
+            logger.info("WebDAV上传未启用")
+    
+    async def upload(self, file_path):
+        """
+        上传图片到WebDAV服务器
+        
+        Args:
+            file_path: 图片文件路径
+            
+        Returns:
+            成功: 图片URL
+            失败: None
+        """
+        if not self.enabled:
+            logger.debug("WebDAV上传未启用，跳过上传")
+            return None
+            
+        if not os.path.exists(file_path):
+            logger.error(f"文件不存在: {file_path}")
+            return None
+            
+        try:
+            # 读取图片文件
+            async with aiofiles.open(file_path, "rb") as f:
+                file_content = await f.read()
+                
+            # 准备上传路径
+            filename = os.path.basename(file_path)
+            remote_path = f"{self.base_path}/{filename}"
+            
+            # 确保目录存在
+            await self._ensure_directory(self.base_path)
+            
+            # 上传文件
+            timeout = aiohttp.ClientTimeout(total=60)
+            auth = aiohttp.BasicAuth(self.username, self.password)
+            async with aiohttp.ClientSession(timeout=timeout, auth=auth) as session:
+                async with session.put(
+                    f"{self.webdav_url}{remote_path}",
+                    data=file_content
+                ) as response:
+                    if response.status in (200, 201, 204):
+                        # 生成访问URL
+                        if self.public_url:
+                            url = f"{self.public_url.rstrip('/')}{remote_path}"
+                        else:
+                            url = f"{self.webdav_url}{remote_path}"
+                            
+                        logger.info(f"图片已成功上传到WebDAV: {url}")
+                        return url
+                    else:
+                        logger.error(f"WebDAV上传失败，状态码: {response.status}")
+                        return None
+                        
+        except Exception as e:
+            logger.error(f"WebDAV上传过程出错: {str(e)}")
+            return None
+    
+    async def _ensure_directory(self, directory):
+        """确保WebDAV目录存在"""
+        try:
+            timeout = aiohttp.ClientTimeout(total=30)
+            auth = aiohttp.BasicAuth(self.username, self.password)
+            async with aiohttp.ClientSession(timeout=timeout, auth=auth) as session:
+                # 检查目录是否存在
+                async with session.request(
+                    "PROPFIND", 
+                    f"{self.webdav_url}{directory}",
+                    headers={"Depth": "0"},
+                ) as response:
+                    if response.status == 207:
+                        return True
+                        
+                # 目录不存在，创建目录
+                async with session.request(
+                    "MKCOL", 
+                    f"{self.webdav_url}{directory}"
+                ) as response:
+                    if response.status in (201, 405):  # 405表示目录已存在
+                        return True
+                    else:
+                        logger.error(f"创建WebDAV目录失败: {directory}, 状态码: {response.status}")
+                        return False
+        except Exception as e:
+            logger.error(f"确保WebDAV目录存在时出错: {str(e)}")
+            return False
+
+# 创建上传器实例
+picgo_uploader = PicGoUploader()
+webdav_uploader = WebDAVUploader()
+
+async def upload_image(file_path):
+    """
+    上传图片到配置的服务
+    
+    Args:
+        file_path: 图片文件路径
+        
+    Returns:
+        元组 (成功标志, 图片URL或None)
+    """
+    # 尝试PicGo上传
+    if picgo_uploader.enabled:
+        picgo_url = await picgo_uploader.upload(file_path)
+        if picgo_url:
+            return True, picgo_url
+    
+    # 尝试WebDAV上传
+    if webdav_uploader.enabled:
+        webdav_url = await webdav_uploader.upload(file_path)
+        if webdav_url:
+            return True, webdav_url
+    
+    # 都上传失败
+    return False, None