remote_client.py

#!/usr/bin/env python3
"""
WebSocket远程客户端 - 在inferless环境中运行
用于连接到主服务器并执行远程命令和文件传输
"""

import asyncio
import websockets
import subprocess
import json
import logging
import sys
import urllib.parse
import base64
import os
import threading
import time
from pathlib import Path

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class RemoteClient:
    def __init__(self, server_url, user_token, space_id):
        self.server_url = server_url
        self.user_token = user_token
        self.space_id = space_id
        self.websocket = None
        self.is_connected = False
        
    async def connect(self):
        """连接到WebSocket服务器"""
        # 构建WebSocket URL
        ws_url = self.server_url.replace('http', 'ws').replace('https', 'wss')
        uri = f"{ws_url}/ws/client/{self.space_id}?token={urllib.parse.quote(self.user_token)}"
        
        logger.info(f"正在连接到服务器: {uri}")
        
        try:
            async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
                self.websocket = websocket
                self.is_connected = True
                logger.info("✅ 已连接到WebSocket服务器")
                
                # 发送客户端注册信息
                await self.register_client()
                
                # 启动消息监听循环
                await self.message_loop()
                
        except Exception as e:
            logger.error(f"❌ 连接失败: {e}")
            self.is_connected = False

    async def register_client(self):
        """注册客户端到服务器"""
        registration = {
            "type": "register",
            "space_id": self.space_id,
            "client_info": {
                "environment": "inferless",
                "python_version": sys.version,
                "working_directory": os.getcwd()
            }
        }
        await self.websocket.send(json.dumps(registration))
        logger.info("📝 已发送客户端注册信息")

    async def message_loop(self):
        """主消息处理循环"""
        while self.is_connected:
            try:
                message = await self.websocket.recv()
                data = json.loads(message)
                await self.handle_message(data)
                
            except websockets.exceptions.ConnectionClosed:
                logger.error("🔌 WebSocket连接已关闭")
                self.is_connected = False
                break
            except Exception as e:
                logger.error(f"❌ 处理消息时出错: {e}")

    async def handle_message(self, data):
        """处理收到的消息"""
        message_type = data.get("type")
        
        if message_type == "command":
            await self.execute_command(data)
        elif message_type == "upload_files":
            await self.upload_files(data)
        elif message_type == "ping":
            await self.send_pong()
        else:
            logger.warning(f"⚠️ 未知消息类型: {message_type}")

    async def execute_command(self, data):
        """执行远程命令"""
        command = data.get("command", "")
        logger.info(f"🚀 执行命令: {command}")
        
        try:
            # 执行命令
            process = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            
            # 发送执行结果
            result = {
                "type": "command_result",
                "command": command,
                "returncode": process.returncode,
                "stdout": process.stdout,
                "stderr": process.stderr,
                "success": process.returncode == 0
            }
            
            await self.websocket.send(json.dumps(result))
            
            if process.returncode == 0:
                logger.info("✅ 命令执行成功")
            else:
                logger.error(f"❌ 命令执行失败，返回码: {process.returncode}")
                
        except subprocess.TimeoutExpired:
            error_result = {
                "type": "command_result",
                "command": command,
                "error": "命令执行超时",
                "success": False
            }
            await self.websocket.send(json.dumps(error_result))
            logger.error("⏰ 命令执行超时")
            
        except Exception as e:
            error_result = {
                "type": "command_result",
                "command": command,
                "error": str(e),
                "success": False
            }
            await self.websocket.send(json.dumps(error_result))
            logger.error(f"❌ 命令执行异常: {e}")

    async def upload_files(self, data):
        """上传指定目录中的所有文件"""
        upload_dir = data.get("directory", "output")
        
        logger.info(f"📁 开始上传目录: {upload_dir}")
        
        if not os.path.exists(upload_dir):
            error_msg = f"目录不存在: {upload_dir}"
            logger.error(error_msg)
            await self.send_upload_result(False, error_msg)
            return

        try:
            # 扫描目录中的所有文件
            files_to_upload = []
            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):
                        files_to_upload.append(file_path)
            
            if not files_to_upload:
                msg = f"目录 {upload_dir} 中没有找到文件"
                logger.info(msg)
                await self.send_upload_result(True, msg)
                return
            
            logger.info(f"📦 找到 {len(files_to_upload)} 个文件，开始上传...")
            
            success_count = 0
            failed_count = 0
            
            for file_path in files_to_upload:
                try:
                    await self.upload_single_file(file_path)
                    success_count += 1
                    logger.info(f"✅ 上传成功: {os.path.basename(file_path)}")
                except Exception as e:
                    failed_count += 1
                    logger.error(f"❌ 上传失败: {file_path} - {e}")
                
                # 稍作延迟避免服务器压力
                await asyncio.sleep(0.1)
            
            result_msg = f"上传完成! 成功: {success_count}, 失败: {failed_count}"
            logger.info(f"📊 {result_msg}")
            await self.send_upload_result(True, result_msg)
            
        except Exception as e:
            error_msg = f"上传过程中出错: {e}"
            logger.error(error_msg)
            await self.send_upload_result(False, error_msg)

    async def upload_single_file(self, file_path):
        """上传单个文件"""
        filename = os.path.basename(file_path)
        
        with open(file_path, 'rb') as f:
            file_content = f.read()
            file_b64 = base64.b64encode(file_content).decode('utf-8')
        
        file_message = {
            "type": "file_upload",
            "filename": filename,
            "content": file_b64,
            "file_size": len(file_content)
        }
        
        await self.websocket.send(json.dumps(file_message))

    async def send_upload_result(self, success, message):
        """发送上传结果"""
        result = {
            "type": "upload_result",
            "success": success,
            "message": message
        }
        await self.websocket.send(json.dumps(result))

    async def send_pong(self):
        """响应ping消息"""
        pong = {"type": "pong"}
        await self.websocket.send(json.dumps(pong))

async def main():
    if len(sys.argv) < 4:
        print("使用方法: python remote_client.py <server_url> <user_token> <space_id>")
        print("示例: python remote_client.py https://gkbtyo-rqvays-5001.preview.cloudstudio.work abc123 space456")
        sys.exit(1)
    
    server_url = sys.argv[1]
    user_token = sys.argv[2]
    space_id = sys.argv[3]
    
    client = RemoteClient(server_url, user_token, space_id)
    
    while True:
        try:
            await client.connect()
        except KeyboardInterrupt:
            logger.info("👋 收到中断信号，正在退出...")
            break
        except Exception as e:
            logger.error(f"❌ 连接异常: {e}")
            logger.info("🔄 5秒后重新尝试连接...")
            await asyncio.sleep(5)

if __name__ == "__main__":
    asyncio.run(main())