Update app.py

app.py

@@ -6,61 +6,74 @@ 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 FileUploader:
-    def __init__(self, server_url, user_token, space_id):
-        self.server_url = server_url
+class WebSocketFileUploader:
+    def __init__(self, websocket_url, user_token, space_id):
+        self.websocket_url = websocket_url
         self.user_token = user_token
         self.space_id = space_id
-        self.websocket = None
         
-    async def connect_and_upload(self, upload_dir):
-        """连接WebSocket并上传文件"""
-        # 构建WebSocket URL
-        ws_url = self.server_url.replace('http://', 'ws://').replace('https://', 'wss://')
-        uri = f"{ws_url}/ws/upload/{self.space_id}"
+    async def upload_file(self, websocket, file_path):
+        """上传单个文件"""
+        if not os.path.exists(file_path):
+            logger.error(f"文件不存在: {file_path}")
+            return False
         
-        logger.info(f"正在连接到 {uri}")
+        filename = os.path.basename(file_path)
+        logger.info(f"正在上传文件: {filename}")
         
         try:
-            async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
-                self.websocket = websocket
-                logger.info("WebSocket连接成功")
+            with open(file_path, 'rb') as f:
+                file_content = f.read()
+                file_b64 = base64.b64encode(file_content).decode('utf-8')
                 
-                # 发送认证信息
-                auth_message = {
-                    "type": "auth",
-                    "token": self.user_token,
-                    "space_id": self.space_id
+                file_msg = {
+                    "type": "file_upload",
+                    "filename": filename,
+                    "content": file_b64,
+                    "space_id": self.space_id,
+                    "user_token": self.user_token
                 }
-                await websocket.send(json.dumps(auth_message))
-                logger.info("已发送认证信息")
-                
-                # 等待认证响应
-                response = await websocket.recv()
-                auth_result = json.loads(response)
                 
-                if auth_result.get("type") == "auth_success":
-                    logger.info("认证成功，开始上传文件")
-                    await self.upload_directory(upload_dir)
-                else:
-                    logger.error(f"认证失败: {auth_result.get('message', '未知错误')}")
-                    return
+                await websocket.send(json.dumps(file_msg))
+                logger.info(f"已发送文件: {filename}")
                 
+                # 等待服务器响应
+                try:
+                    response = await asyncio.wait_for(websocket.recv(), timeout=30.0)
+                    data = json.loads(response)
+                    
+                    if data.get("type") == "upload_success":
+                        logger.info(f"✅ 文件 '{filename}' 上传成功!")
+                        return True
+                    elif data.get("type") == "upload_error":
+                        logger.error(f"❌ 上传失败: {data.get('message', '未知错误')}")
+                        return False
+                    else:
+                        logger.warning(f"收到未知响应: {data}")
+                        return False
+                        
+                except asyncio.TimeoutError:
+                    logger.error(f"上传文件 {filename} 超时")
+                    return False
+                    
         except Exception as e:
-            logger.error(f"连接失败: {e}")
+            logger.error(f"上传文件 {file_path} 时出错: {e}")
+            return False
     
     async def upload_directory(self, upload_dir):
-        """扫描并上传目录中的所有文件"""
+        """上传目录中的所有文件"""
         if not os.path.exists(upload_dir):
             logger.error(f"目录不存在: {upload_dir}")
-            return
+            return 0, 0
         
         logger.info(f"🔍 开始扫描目录: {upload_dir}")
+        logger.info(f"📡 WebSocket服务器: {self.websocket_url}")
         logger.info(f"🔑 Space ID: {self.space_id}")
         logger.info("-" * 50)
         
@@ -74,95 +87,87 @@ class FileUploader:
         
         if not all_files:
             logger.info("📁 目录中没有找到任何文件")
-            # 发送完成消息
-            await self.websocket.send(json.dumps({
-                "type": "upload_complete",
-                "success_count": 0,
-                "failed_count": 0
-            }))
-            return
+            return 0, 0
         
         logger.info(f"📁 找到 {len(all_files)} 个文件，开始上传...")
         
-        success_count = 0
-        failed_count = 0
-        
-        for file_path in all_files:
-            try:
-                if await self.upload_file(file_path):
-                    success_count += 1
-                else:
-                    failed_count += 1
-                # 稍微延迟一下，避免发送过快
-                await asyncio.sleep(0.1)
-            except Exception as e:
-                logger.error(f"上传文件 {file_path} 时发生异常: {e}")
-                failed_count += 1
-        
-        # 发送上传完成消息
-        await self.websocket.send(json.dumps({
-            "type": "upload_complete",
-            "success_count": success_count,
-            "failed_count": failed_count
-        }))
-        
-        logger.info("-" * 50)
-        logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
-        
-        if success_count > 0:
-            logger.info("🎉 文件已成功上传到您的网盘!")
-        
-        return success_count, failed_count
-    
-    async def upload_file(self, file_path):
-        """上传单个文件"""
-        if not os.path.exists(file_path):
-            logger.error(f"文件不存在: {file_path}")
-            return False
-        
-        filename = os.path.basename(file_path)
-        logger.info(f"正在上传文件: {filename}")
-        
         try:
-            with open(file_path, 'rb') as f:
-                file_content = f.read()
-                file_b64 = base64.b64encode(file_content).decode('utf-8')
+            async with websockets.connect(
+                self.websocket_url, 
+                ping_interval=20, 
+                ping_timeout=60,
+                close_timeout=10
+            ) as websocket:
+                logger.info("WebSocket连接已建立")
                 
-                file_message = {
-                    "type": "file_upload",
-                    "filename": filename,
-                    "content": file_b64,
-                    "file_path": os.path.relpath(file_path)
+                # 发送认证信息
+                auth_msg = {
+                    "type": "auth",
+                    "user_token": self.user_token,
+                    "space_id": self.space_id
                 }
+                await websocket.send(json.dumps(auth_msg))
                 
-                await self.websocket.send(json.dumps(file_message))
+                # 等待认证响应
+                try:
+                    auth_response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
+                    auth_data = json.loads(auth_response)
+                    
+                    if auth_data.get("type") != "auth_success":
+                        logger.error(f"认证失败: {auth_data.get('message', '未知错误')}")
+                        return 0, len(all_files)
+                        
+                    logger.info("认证成功，开始上传文件")
+                    
+                except asyncio.TimeoutError:
+                    logger.error("认证超时")
+                    return 0, len(all_files)
                 
-                # 等待上传响应
-                response = await websocket.recv()
-                result = json.loads(response)
+                success_count = 0
+                failed_count = 0
+                
+                for file_path in all_files:
+                    try:
+                        if await self.upload_file(websocket, file_path):
+                            success_count += 1
+                        else:
+                            failed_count += 1
+                        # 稍微延迟一下，避免服务器压力过大
+                        await asyncio.sleep(0.5)
+                    except Exception as e:
+                        logger.error(f"上传文件 {file_path} 时发生异常: {e}")
+                        failed_count += 1
+                
+                logger.info("-" * 50)
+                logger.info(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+                
+                if success_count > 0:
+                    logger.info("🎉 文件已成功上传到您的网盘!")
+                
+                return success_count, failed_count
                 
-                if result.get("type") == "upload_success":
-                    logger.info(f"✅ 文件 '{filename}' 上传成功!")
-                    return True
-                else:
-                    logger.error(f"❌ 上传失败: {result.get('message', '未知错误')}")
-                    return False
-                    
         except Exception as e:
-            logger.error(f"上传文件 {file_path} 时出错: {e}")
-            return False
+            logger.error(f"WebSocket连接失败: {e}")
+            return 0, len(all_files)
+
+async def upload_directory_websocket(upload_dir, websocket_url, user_token, space_id):
+    """使用WebSocket上传目录中的所有文件"""
+    uploader = WebSocketFileUploader(websocket_url, user_token, space_id)
+    return await uploader.upload_directory(upload_dir)
 
-async def main():
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="WebSocket文件上传器 - 扫描并上传指定文件夹中的所有文件")
-    parser.add_argument("user_token", help="您的用户令牌(API密钥)")
+    parser.add_argument("user_token", help="用户令牌 (API Key)")
     parser.add_argument("space_id", help="Space ID")
-    parser.add_argument("--server", default="ws://127.0.0.1:5001", help="服务器的 WebSocket URL 地址 (默认: ws://127.0.0.1:5001)")
+    parser.add_argument("--websocket-url", default="ws://127.0.0.1:5001/ws", help="WebSocket服务器地址 (默认: ws://127.0.0.1:5001/ws)")
     parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
     
     args = parser.parse_args()
     
-    uploader = FileUploader(args.server, args.user_token, args.space_id)
-    await uploader.connect_and_upload(args.upload_dir)
-
-if __name__ == "__main__":
-    asyncio.run(main())
+    # 开始WebSocket上传
+    asyncio.run(upload_directory_websocket(
+        args.upload_dir, 
+        args.websocket_url, 
+        args.user_token, 
+        args.space_id
+    ))