remote_uploader.py

websocket_uploader.py

@@ -1,227 +1,233 @@
 import asyncio
 import websockets
+import subprocess
 import json
 import logging
 import sys
+import urllib.parse
 import base64
 import os
-import time
-import ssl
+import threading
+import requests
 from urllib.parse import urlparse
 
-# 配置日志
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+# Configure logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-def build_websocket_url(website_url):
-    """构建WebSocket URL"""
-    try:
-        # 解析原始URL
-        parsed = urlparse(website_url)
-        
-        # 构建WebSocket URL
-        if parsed.scheme == 'https':
-            ws_scheme = 'wss'
-        else:
-            ws_scheme = 'ws'
-        
-        # 提取主机名，移除端口号
-        hostname = parsed.hostname
-        
-        # 构建WebSocket URL，使用端口8765
-        ws_url = f"{ws_scheme}://{hostname}:8765/ws"
-        
-        logger.info(f"原始URL: {website_url}")
-        logger.info(f"WebSocket URL: {ws_url}")
-        
-        return ws_url
-    
-    except Exception as e:
-        logger.error(f"构建WebSocket URL时出错: {e}")
-        # 回退方案：简单替换
-        if website_url.startswith('https://'):
-            return website_url.replace('https://', 'wss://').replace(':5001', ':8765') + '/ws'
-        else:
-            return website_url.replace('http://', 'ws://').replace(':5001', ':8765') + '/ws'
+# 全局变量，用于在用户输入时访问 WebSocket
+global_websocket = None
 
-async def upload_files_via_websocket(space_id, user_token, website_url):
-    """通过WebSocket上传output目录中的所有文件"""
-    
-    output_dir = "output"
-    if not os.path.exists(output_dir):
-        logger.error(f"输出目录不存在: {output_dir}")
-        return False
-    
-    # 获取所有文件
-    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):
-                files_to_upload.append(file_path)
-    
-    if not files_to_upload:
-        logger.info("输出目录中没有找到任何文件")
-        return True
-    
-    logger.info(f"找到 {len(files_to_upload)} 个文件需要上传")
-    
-    # 构建WebSocket URL
-    ws_url = build_websocket_url(website_url)
-    
+async def connect(space_id, machine_secret, token, upload_file=None, upload_dir=None):
+    global global_websocket
+    # 使用查询参数传递 machine_secret 和 token
+    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}"
+    logger.info(f"Attempting to connect to {uri}")
     try:
-        # 连接到WebSocket服务器
-        logger.info(f"正在连接到 WebSocket 服务器: {ws_url}")
-        
-        # 创建SSL上下文（如果是wss://）
-        ssl_context = None
-        if ws_url.startswith('wss://'):
-            ssl_context = ssl.create_default_context()
-            ssl_context.check_hostname = False
-            ssl_context.verify_mode = ssl.CERT_NONE
-        
-        # 设置连接超时
-        connect_timeout = 10
-        
-        async with websockets.connect(
-            ws_url, 
-            ssl=ssl_context,
-            ping_interval=20,
-            ping_timeout=10,
-            close_timeout=10
-        ) as websocket:
-            logger.info("WebSocket 连接成功")
-            
-            # 发送注册消息
-            register_msg = {
-                'type': 'register',
-                'space_id': space_id,
-                'user_token': user_token
-            }
-            await websocket.send(json.dumps(register_msg))
-            logger.info("已发送注册消息")
-            
-            # 等待注册确认
-            try:
-                response = await asyncio.wait_for(websocket.recv(), timeout=10)
-                response_data = json.loads(response)
-                logger.info(f"收到响应: {response_data}")
-            except asyncio.TimeoutError:
-                logger.error("等待注册确认超时")
-                return False
-            
-            if response_data.get('type') == 'registered':
-                logger.info("注册成功，开始上传文件")
-                
-                # 上传文件
-                success_count = 0
-                failed_count = 0
-                
-                for file_path in files_to_upload:
+        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}
+            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:
-                        filename = os.path.basename(file_path)
-                        logger.info(f"正在上传文件: {filename}")
-                        
-                        # 读取文件内容并编码为base64
-                        with open(file_path, 'rb') as f:
+                        with open(upload_file, 'rb') as f:
                             file_content = f.read()
-                            file_b64 = base64.b64encode(file_content).decode('utf-8')
-                        
-                        file_size = len(file_content)
-                        logger.info(f"文件 {filename} 大小: {file_size} 字节")
-                        
-                        # 发送文件上传消息
-                        upload_msg = {
-                            'type': 'file_upload',
-                            'space_id': space_id,
-                            'user_token': user_token,
-                            'filename': filename,
-                            'content': file_b64
+                        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(upload_msg))
-                        logger.info(f"已发送文件 {filename}")
-                        
-                        # 等待上传结果
-                        try:
-                            upload_response = await asyncio.wait_for(websocket.recv(), timeout=30)
-                            upload_result = json.loads(upload_response)
-                            logger.info(f"上传响应: {upload_result}")
-                        except asyncio.TimeoutError:
-                            logger.error(f"等待文件 {filename} 上传结果超时")
-                            failed_count += 1
-                            continue
-                        
-                        if upload_result.get('type') == 'upload_success':
-                            logger.info(f"文件 {filename} 上传成功")
-                            success_count += 1
-                        else:
-                            logger.error(f"文件 {filename} 上传失败: {upload_result}")
-                            failed_count += 1
-                        
-                        # 稍微延迟避免服务器压力过大
-                        await asyncio.sleep(0.5)
-                        
+                        await websocket.send(json.dumps(file_msg))
+                        logger.info(f"Uploaded file: {upload_file}")
                     except Exception as e:
-                        logger.error(f"上传文件 {file_path} 时发生错误: {e}")
-                        failed_count += 1
-                
-                logger.info(f"上传完成! 成功: {success_count}, 失败: {failed_count}")
-                
-                # 发送完成消息
+                        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:
-                    complete_msg = {'type': 'upload_complete'}
-                    await websocket.send(json.dumps(complete_msg))
-                    logger.info("已发送完成消息")
+                    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"发送完成消息时出错: {e}")
-                
-                return success_count > 0
-                
-            elif response_data.get('type') == 'error':
-                logger.error(f"注册失败: {response_data.get('message', '未知错误')}")
-                return False
-            else:
-                logger.error(f"未知的响应类型: {response_data}")
-                return False
-        
-    except websockets.exceptions.InvalidURI as e:
-        logger.error(f"无效的WebSocket URI: {e}")
-        return False
-    except websockets.exceptions.ConnectionClosed as e:
-        logger.error(f"WebSocket连接被关闭: {e}")
-        return False
+                    logger.error(f"Error processing message: {e}")
+                    await websocket.send(json.dumps({"type": "error", "data": str(e)}))
     except Exception as e:
-        logger.error(f"WebSocket 连接或上传过程中发生错误: {e}")
-        return False
+        logger.error(f"Failed to connect or maintain connection: {e}")
 
-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]
-    
-    logger.info(f"开始 WebSocket 文件上传")
-    logger.info(f"Space ID: {space_id}")
-    logger.info(f"User Token: {user_token[:8]}...")
-    logger.info(f"Website URL: {website_url}")
-    
-    # 运行异步上传
-    try:
-        result = asyncio.run(upload_files_via_websocket(space_id, user_token, website_url))
-        
-        if result:
-            logger.info("所有文件上传成功")
-            sys.exit(0)
-        else:
-            logger.error("文件上传失败")
-            sys.exit(1)
-    except Exception as e:
-        logger.error(f"运行上传程序时出错: {e}")
-        sys.exit(1)
+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):
+                        try:
+                            with open(filename, '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),
+                                "content": file_b64
+                            }
+                            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")
+                        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}")
+                    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)
 
 if __name__ == "__main__":
-    main()
+    if len(sys.argv) < 4:
+        print("Usage: python remote_client.py <space_id> <machine_secret> <token> [--upload <file_path>] [--upload-dir <dir_path>]")
+        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))