Update remote_uploader.py

remote_uploader.py

@@ -1,27 +1,150 @@
 import requests
 import argparse
 import os
-import time
+import asyncio
+import websockets
+import json
+import base64
+import logging
+import sys
 
-class FileUploader:
-    def __init__(self, api_key, space_id):
-        self.server_url = "https://gkbtyo-rqvays-5001.preview.cloudstudio.work"
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+class WebSocketFileUploader:
+    def __init__(self, server_url, api_key, space_id):
+        self.server_url = server_url
         self.api_key = api_key
         self.space_id = space_id
         
-    def upload_file(self, file_path):
-        """上传单个文件"""
-        if not os.path.exists(file_path):
-            print(f"文件不存在: {file_path}")
+    async def connect_and_upload(self, upload_dir):
+        """连接到 WebSocket 服务器并上传文件"""
+        # 将 HTTP URL 转换为 WebSocket URL
+        ws_url = self.server_url.replace('http://', 'ws://').replace('https://', 'wss://').rstrip('/')
+        ws_url += f'/ws/upload/{self.space_id}?token={self.api_key}'
+        
+        logger.info(f"正在连接到 WebSocket: {ws_url}")
+        
+        try:
+            async with websockets.connect(ws_url, ping_interval=20, ping_timeout=60) as websocket:
+                logger.info("WebSocket 连接成功")
+                
+                # 发送认证信息
+                auth_msg = {
+                    "type": "auth",
+                    "api_key": self.api_key,
+                    "space_id": self.space_id
+                }
+                await websocket.send(json.dumps(auth_msg))
+                
+                # 等待认证响应
+                response = await websocket.recv()
+                auth_response = json.loads(response)
+                
+                if auth_response.get("type") == "auth_success":
+                    logger.info("认证成功，开始上传文件...")
+                    return await self._upload_files(websocket, upload_dir)
+                else:
+                    logger.error(f"认证失败: {auth_response.get('message', '未知错误')}")
+                    return False
+                    
+        except Exception as e:
+            logger.error(f"WebSocket 连接失败: {e}")
+            # 如果 WebSocket 失败，回退到 HTTP API
+            logger.info("回退到 HTTP API 上传...")
+            return self._fallback_http_upload(upload_dir)
+    
+    async def _upload_files(self, websocket, upload_dir):
+        """通过 WebSocket 上传目录中的所有文件"""
+        if not os.path.exists(upload_dir):
+            logger.error(f"目录不存在: {upload_dir}")
             return False
         
-        filename = os.path.basename(file_path)
-        print(f"正在上传文件: {filename}")
+        # 获取所有文件
+        all_files = []
+        for root, dirs, files in os.walk(upload_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                if os.path.isfile(file_path):
+                    all_files.append(file_path)
+        
+        if not all_files:
+            logger.info("📁 目录中没有找到任何文件")
+            return True
+        
+        logger.info(f"📁 找到 {len(all_files)} 个文件，开始上传...")
+        
+        success_count = 0
+        failed_count = 0
+        
+        for file_path in all_files:
+            try:
+                if await self._upload_single_file(websocket, file_path):
+                    success_count += 1
+                    logger.info(f"✅ 文件 '{os.path.basename(file_path)}' 上传成功!")
+                else:
+                    failed_count += 1
+                    logger.error(f"❌ 文件 '{os.path.basename(file_path)}' 上传失败")
+                
+                # 稍微延迟一下
+                await asyncio.sleep(0.5)
+                
+            except Exception as e:
+                logger.error(f"上传文件 {file_path} 时发生异常: {e}")
+                failed_count += 1
+        
+        logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+        return success_count > 0
+    
+    async def _upload_single_file(self, websocket, file_path):
+        """上传单个文件通过 WebSocket"""
+        try:
+            filename = os.path.basename(file_path)
+            logger.info(f"正在上传文件: {filename}")
+            
+            # 读取文件内容并编码为 base64
+            with open(file_path, 'rb') as f:
+                file_content = f.read()
+                file_b64 = base64.b64encode(file_content).decode('utf-8')
+            
+            # 构建上传消息
+            upload_msg = {
+                "type": "file_upload",
+                "filename": filename,
+                "content": file_b64,
+                "space_id": self.space_id
+            }
+            
+            # 发送文件
+            await websocket.send(json.dumps(upload_msg))
+            
+            # 等待响应
+            response = await websocket.recv()
+            result = json.loads(response)
+            
+            if result.get("type") == "upload_success":
+                return True
+            else:
+                logger.error(f"上传失败: {result.get('message', '未知错误')}")
+                return False
+                
+        except Exception as e:
+            logger.error(f"上传文件 {file_path} 时出错: {e}")
+            return False
+    
+    def _fallback_http_upload(self, upload_dir):
+        """HTTP API 回退方案"""
+        logger.info("使用 HTTP API 回退方案上传文件...")
+        
+        if not os.path.exists(upload_dir):
+            logger.error(f"目录不存在: {upload_dir}")
+            return False
         
         # 构建完整的 API 端点 URL
         upload_url = f"{self.server_url.rstrip('/')}/api/remote_upload"
         
-        # 准备请求头和数据
+        # 准备请求头
         headers = {
             'X-API-Key': self.api_key
         }
@@ -29,94 +152,77 @@ class FileUploader:
             'space_id': self.space_id
         }
         
-        # 读取文件内容并上传
-        try:
-            with open(file_path, 'rb') as f:
-                files = {
-                    'file': (filename, f)
-                }
-                
-                # 发送 POST 请求
-                response = requests.post(upload_url, headers=headers, data=data, files=files)
+        # 获取所有文件
+        all_files = []
+        for root, dirs, files in os.walk(upload_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                if os.path.isfile(file_path):
+                    all_files.append(file_path)
+        
+        if not all_files:
+            logger.info("📁 目录中没有找到任何文件")
+            return True
+        
+        logger.info(f"📁 找到 {len(all_files)} 个文件，开始HTTP上传...")
+        
+        success_count = 0
+        failed_count = 0
+        
+        for file_path in all_files:
+            try:
+                filename = os.path.basename(file_path)
+                logger.info(f"正在上传文件: {filename}")
                 
-                # 打印服务器响应
-                print(f"服务器响应 ({response.status_code}):")
-                try:
-                    result = response.json()
-                    print(result)
+                with open(file_path, 'rb') as f:
+                    files = {
+                        'file': (filename, f)
+                    }
+                    
+                    # 发送 POST 请求
+                    response = requests.post(upload_url, headers=headers, data=data, files=files, timeout=30)
+                    
                     if response.status_code == 200:
-                        print(f"✅ 文件 '{filename}' 上传成功!")
-                        return True
+                        logger.info(f"✅ 文件 '{filename}' 上传成功!")
+                        success_count += 1
                     else:
-                        print(f"❌ 上传失败: {result.get('error', '未知错误')}")
-                        return False
-                except requests.exceptions.JSONDecodeError:
-                    print(response.text)
-                    return False
-                    
-        except IOError as e:
-            print(f"读取文件时出错: {e}")
-            return False
-        except requests.exceptions.RequestException as e:
-            print(f"请求时出错: {e}")
-            return False
-
-def upload_directory_once(upload_dir, api_key, space_id):
-    """一次性扫描并上传目录中的所有文件"""
-    if not os.path.exists(upload_dir):
-        print(f"目录不存在: {upload_dir}")
-        return
-    
-    print(f"🔍 开始扫描目录: {upload_dir}")
-    print(f"🔑 Space ID: {space_id}")
-    print("-" * 50)
-    
-    uploader = FileUploader(api_key, space_id)
-    
-    # 获取所有文件
-    all_files = []
-    for root, dirs, files in os.walk(upload_dir):
-        for file in files:
-            file_path = os.path.join(root, file)
-            if os.path.isfile(file_path):
-                all_files.append(file_path)
-    
-    if not all_files:
-        print("📁 目录中没有找到任何文件")
-        return
-    
-    print(f"📁 找到 {len(all_files)} 个文件，开始上传...")
-    
-    success_count = 0
-    failed_count = 0
-    
-    for file_path in all_files:
-        try:
-            if uploader.upload_file(file_path):
-                success_count += 1
-            else:
+                        logger.error(f"❌ 上传失败: HTTP {response.status_code}")
+                        failed_count += 1
+                
+                time.sleep(0.5)
+                        
+            except Exception as e:
+                logger.error(f"上传文件 {file_path} 时发生异常: {e}")
                 failed_count += 1
-            # 稍微延迟一下，避免服务器压力过大
-            time.sleep(0.5)
-        except Exception as e:
-            print(f"上传文件 {file_path} 时发生异常: {e}")
-            failed_count += 1
+        
+        logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+        return success_count > 0
+
+async def upload_directory_websocket(upload_dir, server_url, api_key, space_id):
+    """使用 WebSocket 方式上传目录中的所有文件"""
+    logger.info(f"🔍 开始扫描目录: {upload_dir}")
+    logger.info(f"📡 服务器地址: {server_url}")
+    logger.info(f"🔑 Space ID: {space_id}")
+    logger.info("-" * 50)
     
-    print("-" * 50)
-    print(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+    uploader = WebSocketFileUploader(server_url, api_key, space_id)
+    success = await uploader.connect_and_upload(upload_dir)
     
-    if success_count > 0:
-        print("🎉 文件已成功上传到您的网盘!")
+    if success:
+        logger.info("🎉 文件已成功上传到您的网盘!")
+    else:
+        logger.error("❌ 文件上传失败")
     
-    return success_count, failed_count
+    return success
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="一次性文件上传器 - 扫描并上传指定文件夹中的所有文件")
+    parser = argparse.ArgumentParser(description="WebSocket 文件上传器 - 扫描并上传指定文件夹中的所有文件")
     parser.add_argument("api_key", help="您的 API 密钥")
     parser.add_argument("space_id", help="Space ID")
+    parser.add_argument("--server", default="https://gkbtyo-rqvays-5001.preview.cloudstudio.work", help="服务器的 URL 地址")
     parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
     
     args = parser.parse_args()
     
-    # 开始一次性上传
-    upload_directory_once(args.upload_dir, args.api_key, args.space_id)
+    # 运行异步上传
+    asyncio.run(upload_directory_websocket(args.upload_dir, args.server, args.api_key, args.space_id))