Update remote_uploader.py

remote_uploader.py

@@ -5,130 +5,138 @@ import logging
 import sys
 import base64
 import os
-import argparse
+import time
 
 # 配置日志
 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):
-        # 将HTTP URL转换为WebSocket URL
-        if server_url.startswith('http://'):
-            self.ws_url = server_url.replace('http://', 'ws://') + '/ws/upload'
-        elif server_url.startswith('https://'):
-            self.ws_url = server_url.replace('https://', 'wss://') + '/ws/upload'
-        else:
-            self.ws_url = f"ws://{server_url}/ws/upload"
-            
-        self.api_key = api_key
-        self.space_id = space_id
-        
-    async def upload_file(self, file_path):
-        """使用WebSocket上传单个文件"""
-        if not os.path.exists(file_path):
-            print(f"文件不存在: {file_path}")
-            return False
-        
-        filename = os.path.basename(file_path)
-        print(f"正在上传文件: {filename}")
-        
-        try:
-            # 读取文件内容并编码为base64
-            with open(file_path, 'rb') as f:
-                file_content = f.read()
-                file_b64 = base64.b64encode(file_content).decode('utf-8')
-            
-            # 连接到WebSocket服务器
-            async with websockets.connect(self.ws_url) as websocket:
-                # 发送认证和文件数据
-                upload_msg = {
-                    "type": "file_upload",
-                    "api_key": self.api_key,
-                    "space_id": self.space_id,
-                    "filename": filename,
-                    "content": file_b64
-                }
-                
-                await websocket.send(json.dumps(upload_msg))
-                logger.debug(f"发送文件上传请求: {filename}")
-                
-                # 等待服务器响应
-                response = await websocket.recv()
-                data = json.loads(response)
-                
-                if data.get("type") == "success":
-                    print(f"✅ 文件 '{filename}' 上传成功!")
-                    return True
-                else:
-                    print(f"❌ 上传失败: {data.get('error', '未知错误')}")
-                    return False
-                    
-        except Exception as e:
-            print(f"上传文件 {filename} 时出错: {e}")
-            return False
-
-async def upload_directory_once(upload_dir, server_url, api_key, space_id):
-    """一次性扫描并上传目录中的所有文件"""
-    if not os.path.exists(upload_dir):
-        print(f"目录不存在: {upload_dir}")
-        return
+async def upload_files_via_websocket(space_id, user_token, website_url):
+    """通过WebSocket上传output目录中的所有文件"""
     
-    print(f"🔍 开始扫描目录: {upload_dir}")
-    print(f"📡 服务器地址: {server_url}")
-    print(f"🔑 Space ID: {space_id}")
-    print("-" * 50)
-    
-    uploader = WebSocketFileUploader(server_url, api_key, space_id)
+    output_dir = "output"
+    if not os.path.exists(output_dir):
+        logger.error(f"输出目录不存在: {output_dir}")
+        return False
     
     # 获取所有文件
-    all_files = []
-    for root, dirs, files in os.walk(upload_dir):
+    files_to_upload = []
+    for root, dirs, files in os.walk(output_dir):
         for file in files:
             file_path = os.path.join(root, file)
             if os.path.isfile(file_path):
-                all_files.append(file_path)
+                files_to_upload.append(file_path)
+    
+    if not files_to_upload:
+        logger.info("输出目录中没有找到任何文件")
+        return True
     
-    if not all_files:
-        print("📁 目录中没有找到任何文件")
-        return
+    logger.info(f"找到 {len(files_to_upload)} 个文件需要上传")
     
-    print(f"📁 找到 {len(all_files)} 个文件，开始上传...")
+    # 构建WebSocket URL
+    ws_url = website_url.replace('http://', 'ws://').replace('https://', 'wss://') + '/socket.io/?EIO=4&transport=websocket'
     
+    try:
+        # 连接到WebSocket服务器
+        logger.info(f"正在连接到 WebSocket 服务器: {ws_url}")
+        
+        # 使用socket.io客户端协议
+        import socketio
+        
+        sio = socketio.AsyncClient()
+        
+        @sio.event
+        async def connect():
+            logger.info("WebSocket 连接成功")
+            # 注册客户端
+            await sio.emit('register_client', {
+                'space_id': space_id,
+                'user_token': user_token
+            })
+        
+        @sio.event
+        async def registered(data):
+            logger.info(f"客户端注册成功: {data}")
+            # 开始上传文件
+            await upload_files(sio, files_to_upload, space_id, user_token)
+        
+        @sio.event
+        async def upload_success(data):
+            logger.info(f"文件上传成功: {data}")
+        
+        @sio.event
+        async def upload_error(data):
+            logger.error(f"文件上传失败: {data}")
+        
+        @sio.event
+        async def error(data):
+            logger.error(f"WebSocket 错误: {data}")
+        
+        @sio.event
+        async def disconnect():
+            logger.info("WebSocket 连接断开")
+        
+        # 连接到服务器
+        await sio.connect(website_url)
+        
+        # 等待所有文件上传完成
+        await sio.wait()
+        
+    except Exception as e:
+        logger.error(f"WebSocket 连接或上传过程中发生错误: {e}")
+        return False
+
+async def upload_files(sio, files_to_upload, space_id, user_token):
+    """上传文件列表"""
     success_count = 0
     failed_count = 0
     
-    for file_path in all_files:
+    for file_path in files_to_upload:
         try:
-            if await uploader.upload_file(file_path):
-                success_count += 1
-            else:
-                failed_count += 1
-            # 稍微延迟一下，避免服务器压力过大
+            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')
+            
+            # 发送文件上传请求
+            await sio.emit('file_upload', {
+                'space_id': space_id,
+                'user_token': user_token,
+                'filename': filename,
+                'content': file_b64
+            })
+            
+            success_count += 1
+            # 稍微延迟避免服务器压力过大
             await asyncio.sleep(0.5)
+            
         except Exception as e:
-            print(f"上传文件 {file_path} 时发生异常: {e}")
+            logger.error(f"上传文件 {file_path} 时发生错误: {e}")
             failed_count += 1
     
-    print("-" * 50)
-    print(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+    logger.info(f"上传完成! 成功: {success_count}, 失败: {failed_count}")
     
-    if success_count > 0:
-        print("🎉 文件已成功上传到您的网盘!")
-    
-    return success_count, failed_count
+    # 上传完成后断开连接
+    await sio.disconnect()
 
-async def main():
-    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)")
+def main():
+    if len(sys.argv) != 4:
+        print("用法: python3 websocket_uploader.py <space_id> <user_token> <website_url>")
+        sys.exit(1)
+    
+    space_id = sys.argv[1]
+    user_token = sys.argv[2] 
+    website_url = sys.argv[3]
     
-    args = parser.parse_args()
+    logger.info(f"开始 WebSocket 文件上传")
+    logger.info(f"Space ID: {space_id}")
+    logger.info(f"Website URL: {website_url}")
     
-    # 开始上传
-    await upload_directory_once(args.upload_dir, args.server, args.api_key, args.space_id)
+    # 运行异步上传
+    asyncio.run(upload_files_via_websocket(space_id, user_token, website_url))
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()