Update remote_uploader.py

remote_uploader.py

@@ -1,95 +1,157 @@
 import asyncio
 import websockets
-import subprocess
 import json
 import logging
 import sys
-import urllib.parse
 import base64
 import os
-import threading
-import requests
 from urllib.parse import urlparse
-import argparse
 
-# Configure logging
+# 配置日志
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# 全局变量，用于在用户输入时访问 WebSocket
-global_websocket = None
-
-async def connect(card_id, machine_secret, upload_files=None, server="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/"):
-    global global_websocket
-
-    # 使用查询参数传递 machine_secret
-    encoded_secret = urllib.parse.quote(machine_secret)
-    uri = f"{server}{card_id}?secret={encoded_secret}"
-    logger.info(f"Attempting to connect to {uri}")
-
-    try:
-        async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
-            global_websocket = websocket
-            logger.info("Connected to WebSocket")
-
-            machine_info = {"type": "machine", "card_id": card_id, "token": machine_secret}  # 添加用户token
-            await websocket.send(json.dumps(machine_info))
-            logger.debug(f"Sent machine registration: {machine_info}")
-
-            # 如果指定了上传文件列表，发送文件内容
-            if upload_files:
-                for upload_file in upload_files:
-                    if os.path.exists(upload_file):
-                        try:
-                            with open(upload_file, 'rb') as f:
-                                file_content = f.read()
-                            file_b64 = base64.b64encode(file_content).decode('utf-8')
-                            file_msg = {
-                                "type": "file_upload",
-                                "filename": os.path.basename(upload_file),
-                                "content": file_b64
-                                # 去除 machine_id
-                            }
-                            await websocket.send(json.dumps(file_msg))
-                            logger.info(f"Uploaded file: {upload_file}")
-                        except Exception as e:
-                            logger.error(f"Error uploading file {upload_file}: {e}")
-                    else:
-                        logger.error(f"File not found: {upload_file}")
-
-            # 等待一段时间以接收任何响应，然后关闭
-            try:
-                await asyncio.sleep(5)  # 等待服务器响应
-            except:
-                pass
-
-    except Exception as e:
-        logger.error(f"Failed to connect or maintain connection: {e}")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="WebSocket-based File Uploader for remote control file transfer")
-    parser.add_argument("api_key", help="API Key (used as machine_secret and token)")
-    parser.add_argument("space_id", help="Space ID (used as card_id)")
-    parser.add_argument("--server", default="wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/", help="WebSocket server base URL (default: wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/)")
-    parser.add_argument("--upload-dir", default="output", help="Directory to upload files from (default: output)")
-
-    args = parser.parse_args()
-
-    if not os.path.exists(args.upload_dir):
-        print(f"目录不存在: {args.upload_dir}")
-        sys.exit(1)
+class WebSocketFileUploader:
+    def __init__(self, server_url, api_key, space_id):
+        # 将HTTP URL转换为WebSocket URL
+        parsed_url = urlparse(server_url)
+        scheme = 'wss' if parsed_url.scheme == 'https' else 'ws'
+        self.ws_url = f"{scheme}://{parsed_url.netloc}/ws/upload/{space_id}"
+        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, ping_interval=20, ping_timeout=60) as 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":
+                    print(f"❌ 认证失败: {auth_response.get('message', '未知错误')}")
+                    return False
+                
+                # 发送文件
+                file_msg = {
+                    "type": "file_upload",
+                    "filename": filename,
+                    "content": file_b64,
+                    "space_id": self.space_id
+                }
+                await websocket.send(json.dumps(file_msg))
+                
+                # 等待上传响应
+                response = await websocket.recv()
+                upload_response = json.loads(response)
+                
+                if upload_response.get("type") == "upload_success":
+                    print(f"✅ 文件 '{filename}' 上传成功!")
+                    return True
+                else:
+                    print(f"❌ 上传失败: {upload_response.get('message', '未知错误')}")
+                    return False
+                    
+        except Exception as e:
+            print(f"上传文件时出错: {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
+    
+    print(f"🔍 开始扫描目录: {upload_dir}")
+    print(f"📡 服务器地址: {server_url}")
+    print(f"🔑 Space ID: {space_id}")
+    print("-" * 50)
+    
+    uploader = WebSocketFileUploader(server_url, api_key, space_id)
+    
+    # 获取所有文件
     all_files = []
-    for root, dirs, files in os.walk(args.upload_dir):
+    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("目录中没有找到任何文件")
-        sys.exit(0)
+        print("📁 目录中没有找到任何文件")
+        return
+    
+    print(f"📁 找到 {len(all_files)} 个文件，开始上传...")
+    
+    success_count = 0
+    failed_count = 0
+    
+    for file_path in all_files:
+        try:
+            if await uploader.upload_file(file_path):
+                success_count += 1
+            else:
+                failed_count += 1
+            # 稍微延迟一下，避免服务器压力过大
+            await asyncio.sleep(0.5)
+        except Exception as e:
+            print(f"上传文件 {file_path} 时发生异常: {e}")
+            failed_count += 1
+    
+    print("-" * 50)
+    print(f"📊 上传完成! 成功: {success_count}, 失败: {failed_count}")
+    
+    if success_count > 0:
+        print("🎉 文件已成功上传到您的网盘!")
+    
+    return success_count, failed_count
+
+async def main():
+    if len(sys.argv) < 3:
+        print("使用方法: python3 remote_uploader.py <api_key> <space_id> [--server <server_url>] [--upload-dir <directory>]")
+        sys.exit(1)
+    
+    api_key = sys.argv[1]
+    space_id = sys.argv[2]
+    
+    # 解析可选参数
+    server_url = "https://gkbtyo-rqvays-5001.preview.cloudstudio.work"
+    upload_dir = "output"
+    
+    i = 3
+    while i < len(sys.argv):
+        if sys.argv[i] == "--server" and i + 1 < len(sys.argv):
+            server_url = sys.argv[i + 1]
+            i += 2
+        elif sys.argv[i] == "--upload-dir" and i + 1 < len(sys.argv):
+            upload_dir = sys.argv[i + 1]
+            i += 2
+        else:
+            i += 1
+    
+    # 开始异步上传
+    await upload_directory_once(upload_dir, server_url, api_key, space_id)
 
-    print(f"找到 {len(all_files)} 个文件，开始上传...")
-    asyncio.run(connect(args.space_id, args.api_key, upload_files=all_files, server=args.server))
-    print("上传完成!")
+if __name__ == "__main__":
+    asyncio.run(main())