Update remote_uploader.py

remote_uploader.py

@@ -10,6 +10,7 @@ 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')
@@ -18,216 +19,77 @@ logger = logging.getLogger(__name__)
 # 全局变量，用于在用户输入时访问 WebSocket
 global_websocket = None
 
-async def connect(space_id, machine_secret, token, upload_file=None, upload_dir=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 和 token
+
+    # 使用查询参数传递 machine_secret
     encoded_secret = urllib.parse.quote(machine_secret)
-    encoded_token = urllib.parse.quote(token)
-    uri = f"wss://remote-terminal-worker.nianxi4563.workers.dev/terminal/{space_id}?secret={encoded_secret}&token={encoded_token}"
+    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", "space_id": space_id, "token": token}
+
+            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_file:
-                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,
-                            "token": token
-                        }
-                        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}")
-
-            # 如果指定了上传目录，扫描并上传所有文件
-            if upload_dir:
-                if os.path.exists(upload_dir):
-                    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 all_files:
-                        logger.info(f"Found {len(all_files)} files in {upload_dir}, starting upload...")
-                        for file_path in all_files:
-                            try:
-                                with open(file_path, '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(file_path),
-                                    "content": file_b64,
-                                    "token": token
-                                }
-                                await websocket.send(json.dumps(file_msg))
-                                logger.info(f"Uploaded file from dir: {file_path}")
-                            except Exception as e:
-                                logger.error(f"Error uploading file {file_path}: {e}")
-                    else:
-                        logger.info(f"No files found in {upload_dir}")
-                else:
-                    logger.error(f"Directory not found: {upload_dir}")
-
-            # 创建下载目录
-            os.makedirs("./downloads", exist_ok=True)
-
-            # 启动用户输入监听（非阻塞）
-            asyncio.create_task(listen_user_input())
-
-            while True:
-                try:
-                    message = await websocket.recv()
-                    logger.debug(f"Received message: {message}")
-                    data = json.loads(message)
-                    if data["type"] == "command":
-                        command = data["command"]
-                        logger.info(f"Executing command: {command}")
-                        try:
-                            process = subprocess.run(command, shell=True, capture_output=True, text=True)
-                            output = process.stdout + process.stderr
-                            logger.debug(f"Command output: {output}")
-                            if process.returncode == 0:
-                                await websocket.send(json.dumps({"type": "output", "data": output}))
-                            else:
-                                await websocket.send(json.dumps({"type": "error", "data": output}))
-                        except Exception as e:
-                            error_message = f"Error executing command: {e}"
-                            logger.error(error_message)
-                            await websocket.send(json.dumps({"type": "error", "data": error_message}))
-                    elif data["type"] == "file_download_url":
-                        # 处理从服务器发来的文件下载URL
-                        url = data["url"]
-                        filename = data["filename"]
-                        logger.info(f"Received file download URL: {url}")
-                        try:
-                            download_dir = "./downloads"
-                            os.makedirs(download_dir, exist_ok=True)
-                            response = requests.get(url)
-                            if response.status_code == 200:
-                                file_path = os.path.join(download_dir, filename)
-                                with open(file_path, 'wb') as f:
-                                    f.write(response.content)
-                                success_msg = f"File downloaded successfully: {file_path}"
-                                logger.info(success_msg)
-                                await websocket.send(json.dumps({
-                                    "type": "output",
-                                    "data": success_msg
-                                }))
-                            else:
-                                error_msg = f"Failed to download file. Status code: {response.status_code}"
-                                logger.error(error_msg)
-                                await websocket.send(json.dumps({
-                                    "type": "error",
-                                    "data": error_msg
-                                }))
-                        except Exception as e:
-                            error_msg = f"Error downloading file from URL: {e}"
-                            logger.error(error_msg)
-                            await websocket.send(json.dumps({
-                                "type": "error",
-                                "data": error_msg
-                            }))
-                    elif data["type"] == "ping":
-                        logger.debug("Received ping, sending pong")
-                        await websocket.send(json.dumps({"type": "pong"}))
-                except websockets.exceptions.ConnectionClosed:
-                    logger.error("WebSocket connection closed")
-                    break
-                except Exception as e:
-                    logger.error(f"Error processing message: {e}")
-                    await websocket.send(json.dumps({"type": "error", "data": str(e)}))
-    except Exception as e:
-        logger.error(f"Failed to connect or maintain connection: {e}")
-
-async def listen_user_input():
-    global global_websocket
-    logger.info("Started listening for user input. Use ' --upload' to upload files to server.")
-    while True:
-        try:
-            user_input = await asyncio.to_thread(input, "Enter command or file path to upload: ")
-            if user_input.strip():
-                parts = user_input.strip().split()
-                if len(parts) >= 2 and parts[-1] == "--upload":
-                    filename = " ".join(parts[:-1])
-                    if os.path.exists(filename):
+            # 如果指定了上传文件列表，发送文件内容
+            if upload_files:
+                for upload_file in upload_files:
+                    if os.path.exists(upload_file):
                         try:
-                            with open(filename, 'rb') as f:
+                            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(filename),
+                                "filename": os.path.basename(upload_file),
                                 "content": file_b64
+                                # 去除 machine_id
                             }
-                            if global_websocket:
-                                await global_websocket.send(json.dumps(file_msg))
-                                logger.info(f"Uploaded file via input: {filename}")
-                            else:
-                                logger.error("WebSocket not connected")
+                            await websocket.send(json.dumps(file_msg))
+                            logger.info(f"Uploaded file: {upload_file}")
                         except Exception as e:
-                            logger.error(f"Error uploading file {filename}: {e}")
-                    else:
-                        logger.error(f"File not found: {filename}")
-                elif len(parts) >= 2 and parts[-1] == "--download":
-                    url = " ".join(parts[:-1])
-                    try:
-                        parsed_url = urlparse(url)
-                        filename = os.path.basename(parsed_url.path)
-                        download_dir = "./downloads"
-                        os.makedirs(download_dir, exist_ok=True)
-                        logger.info(f"Downloading file from URL: {url}")
-                        response = requests.get(url)
-                        if response.status_code == 200:
-                            file_path = os.path.join(download_dir, filename)
-                            with open(file_path, 'wb') as f:
-                                f.write(response.content)
-                            logger.info(f"File downloaded successfully: {file_path}")
-                        else:
-                            logger.error(f"Failed to download file. Status code: {response.status_code}")
-                    except Exception as e:
-                        logger.error(f"Error downloading from URL {url}: {e}")
-                else:
-                    if global_websocket:
-                        await global_websocket.send(json.dumps({
-                            "type": "command",
-                            "command": user_input
-                        }))
-                        logger.info(f"Sent command: {user_input}")
+                            logger.error(f"Error uploading file {upload_file}: {e}")
                     else:
-                        logger.error("WebSocket not connected")
-            await asyncio.sleep(0.1)
-        except Exception as e:
-            logger.error(f"Error in user input listener: {e}")
-            await asyncio.sleep(1)
+                        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__":
-    if len(sys.argv) < 4:
-        print("Usage: python remote_client.py <space_id> <machine_secret> <token> [--upload <file_path>] [--upload-dir <dir_path>]")
+    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)
-    space_id = sys.argv[1]
-    machine_secret = sys.argv[2]
-    token = sys.argv[3]
-    upload_file = None
-    upload_dir = None
-    if len(sys.argv) > 4:
-        if sys.argv[4] == "--upload" and len(sys.argv) > 5:
-            upload_file = sys.argv[5]
-        elif sys.argv[4] == "--upload-dir" and len(sys.argv) > 5:
-            upload_dir = sys.argv[5]
-    asyncio.run(connect(space_id, machine_secret, token, upload_file, upload_dir))
+
+    all_files = []
+    for root, dirs, files in os.walk(args.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(f"找到 {len(all_files)} 个文件，开始上传...")
+    asyncio.run(connect(args.space_id, args.api_key, upload_files=all_files, server=args.server))
+    print("上传完成!")