""" JADE Benchmark 环境控制器 负责VM的重置、文件注入/收集、截图获取等操作 """ import subprocess import time import os import requests from PIL import Image import io import logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class JadeEnv: """轻量级JADE虚拟机环境控制器""" def __init__(self, vmx_path, snapshot_name="Jade_Ready", vm_ip="192.168.116.129", vm_password=None, guest_username=None, guest_password=None): """ 初始化JADE环境 Args: vmx_path: 虚拟机.vmx文件路径 snapshot_name: 快照名称 vm_ip: 虚拟机IP地址（用于HTTP通信） vm_password: 虚拟机文件加密密码（-vp参数） guest_username: 虚拟机内操作系统用户名（-gu参数） guest_password: 虚拟机内操作系统密码（-gp参数） """ self.vmx_path = vmx_path self.snapshot_name = snapshot_name self.vm_ip = vm_ip self.vm_url = f"http://{vm_ip}:5000" # VMware认证参数 self.vm_password = vm_password self.guest_username = guest_username self.guest_password = guest_password # VMware Fusion路径（macOS） self.vmrun = "/Applications/VMware Fusion.app/Contents/Library/vmrun" # 虚拟机内路径 self.guest_desktop = r"C:\Users\lzy\Desktop" logger.info(f"JadeEnv初始化: VM={os.path.basename(vmx_path)}, Snapshot={snapshot_name}") logger.info(f" 认证配置: vm_password={'已设置' if vm_password else '未设置'}, " f"guest_user={'已设置' if guest_username else '未设置'}, " f"guest_pass={'已设置' if guest_password else '未设置'}") def _build_vmrun_cmd(self, *args): """构建vmrun命令""" cmd = [self.vmrun, "-T", "fusion"] # 添加认证参数 if self.vm_password: cmd.extend(["-vp", self.vm_password]) if self.guest_username: cmd.extend(["-gu", self.guest_username]) if self.guest_password: cmd.extend(["-gp", self.guest_password]) cmd.extend(args) return cmd def _run_vmrun(self, *args, check=True, timeout=30): """执行vmrun命令""" cmd = self._build_vmrun_cmd(*args) # 打印完整命令（隐藏密码） cmd_display = [] skip_next = False for i, part in enumerate(cmd): if skip_next: cmd_display.append("***") skip_next = False elif part in ["-vp", "-gp"]: cmd_display.append(part) skip_next = True else: cmd_display.append(part) logger.info(f"执行vmrun命令: {' '.join(cmd_display)}") try: result = subprocess.run( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=timeout ) if check and result.returncode != 0: error_msg = result.stderr or result.stdout raise RuntimeError(f"vmrun命令执行失败: {error_msg}") return result except subprocess.TimeoutExpired: logger.error(f"❌ vmrun命令超时（{timeout}秒）") raise RuntimeError(f"vmrun命令执行超时（{timeout}秒）") def _detect_and_update_ip(self): """ 检测VM的IP地址，如果变化则自动更新 Returns: bool: IP是否发生变化 """ logger.info("🔍 检测VM IP地址...") try: # 使用vmrun获取VM IP cmd = self._build_vmrun_cmd("getGuestIPAddress", self.vmx_path) result = subprocess.run( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=10 ) if result.returncode == 0: new_ip = result.stdout.strip() if new_ip and new_ip != "": if new_ip != self.vm_ip: logger.info(f"⚠️ IP地址已变化: {self.vm_ip} → {new_ip}") logger.info(f" 自动更新IP地址...") # 更新实例变量 self.vm_ip = new_ip self.vm_url = f"http://{new_ip}:5000" # 更新配置文件 try: import json from pathlib import Path # 获取项目根目录（jade_env.py在scripts/core/，向上3级到项目根目录） project_root = Path(__file__).parent.parent.parent config_path = project_root / "config.json" if config_path.exists(): with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) config["network"]["vm_ip"] = new_ip with open(config_path, 'w', encoding='utf-8') as f: json.dump(config, f, indent=2, ensure_ascii=False) logger.info(f"✅ 配置文件已更新: {config_path}") except Exception as e: logger.warning(f"⚠️ 更新配置文件失败: {e}（不影响使用）") return True else: logger.info(f"✅ IP地址未变化: {self.vm_ip}") return False else: logger.warning(f"⚠️ vmrun返回空IP地址") return False else: error_msg = result.stderr or result.stdout logger.warning(f"⚠️ 获取IP失败: {error_msg}（将使用配置中的IP）") return False except subprocess.TimeoutExpired: logger.warning(f"⚠️ 获取IP超时（将使用配置中的IP）") return False except Exception as e: logger.warning(f"⚠️ 检测IP异常: {e}（将使用配置中的IP）") return False def reset(self, wait_time=5): """ 重置环境：恢复快照并启动虚拟机 Args: wait_time: 启动后等待时间（秒） """ logger.info(f"正在恢复快照: {self.snapshot_name}...") try: # 1. 恢复快照 self._run_vmrun("revertToSnapshot", self.vmx_path, self.snapshot_name) logger.info("✅ 快照恢复成功") time.sleep(3) # 2. 启动虚拟机（如果未运行） logger.info("正在启动虚拟机...") result = self._run_vmrun("start", self.vmx_path, check=False) if result.returncode == 0: logger.info("✅ 虚拟机启动成功") else: # 可能已经在运行 if "is already running" in result.stderr.lower(): logger.info("✅ 虚拟机已在运行") else: logger.warning(f"启动虚拟机警告: {result.stderr}") # 3. 等待系统稳定（快照恢复后agent_server已在运行） logger.info(f"等待系统稳定 ({wait_time}秒)...") time.sleep(wait_time) # 4. 检测并更新IP地址（恢复快照后IP可能变化） self._detect_and_update_ip() # 5. 验证HTTP服务可用 self._wait_for_http_service() logger.info("✅ 环境重置完成") except Exception as e: logger.error(f"❌ 环境重置失败: {e}") raise def _wait_for_http_service(self, max_retries=10, retry_interval=3): """等待agent_server.py HTTP服务可用""" logger.info(f"等待虚拟机HTTP服务... (URL: {self.vm_url})") # 绕过代理（避免Clash等代理工具干扰局域网访问） proxies = { 'http': None, 'https': None } for i in range(max_retries): try: logger.debug(f"尝试连接: {self.vm_url}/screen_info (timeout=5秒, 不使用代理)") response = requests.get(f"{self.vm_url}/screen_info", timeout=5, proxies=proxies) logger.debug(f"收到响应: status_code={response.status_code}") if response.status_code == 200: logger.info("✅ HTTP服务已就绪") return True else: logger.warning(f"HTTP状态码异常: {response.status_code}") except requests.exceptions.Timeout as e: logger.info(f"HTTP服务未就绪（超时），重试 {i+1}/{max_retries}...") if i < max_retries - 1: time.sleep(retry_interval) except requests.exceptions.ConnectionError as e: logger.info(f"HTTP服务未就绪（连接失败: {str(e)[:50]}），重试 {i+1}/{max_retries}...") if i < max_retries - 1: time.sleep(retry_interval) except requests.exceptions.RequestException as e: logger.warning(f"HTTP请求异常: {type(e).__name__}: {str(e)[:100]}") if i < max_retries - 1: logger.info(f"重试 {i+1}/{max_retries}...") time.sleep(retry_interval) else: logger.error("❌ HTTP服务超时！") logger.error(f" 最后错误: {e}") return False logger.error("❌ HTTP服务超时！请检查agent_server.py是否在VM中运行") logger.info(" 在VM中运行: python agent_server.py") return False def inject_file(self, host_path, guest_filename=None): """ 将文件从主机注入到虚拟机桌面 Args: host_path: 主机文件路径 guest_filename: 虚拟机中的文件名（默认使用原文件名） """ if not os.path.exists(host_path): raise FileNotFoundError(f"源文件不存在: {host_path}") if guest_filename is None: guest_filename = os.path.basename(host_path) guest_path = f"{self.guest_desktop}\\{guest_filename}" # 获取文件大小 file_size = os.path.getsize(host_path) file_size_kb = file_size / 1024 logger.info(f"注入文件: {os.path.basename(host_path)} ({file_size_kb:.1f}KB) → 虚拟机桌面") logger.info(f" 源路径: {host_path}") logger.info(f" 目标路径: {guest_path}") try: start_time = time.time() # 使用vmrun传输（30秒超时） self._run_vmrun( "copyFileFromHostToGuest", self.vmx_path, host_path, guest_path, timeout=30 ) elapsed = time.time() - start_time logger.info(f"✅ 文件注入成功: {guest_filename} (耗时 {elapsed:.1f}秒)") except Exception as e: logger.error(f"❌ 文件注入失败: {e}") raise def collect_file(self, guest_filename, host_path): """ 从虚拟机桌面收集文件到主机 Args: guest_filename: 虚拟机桌面上的文件名 host_path: 主机保存路径 """ guest_path = f"{self.guest_desktop}\\{guest_filename}" logger.info(f"收集文件: {guest_filename} → {os.path.basename(host_path)}") try: # 确保目标目录存在 os.makedirs(os.path.dirname(host_path), exist_ok=True) # 方法1: 尝试使用vmrun try: self._run_vmrun( "copyFileFromGuestToHost", self.vmx_path, guest_path, host_path ) logger.info(f"✅ 文件收集成功（vmrun）: {guest_filename}") return except RuntimeError as e: logger.warning(f"vmrun收集失败，尝试HTTP方式: {e}") # 方法2: 尝试通过HTTP下载（备用） try: response = requests.get( f"{self.vm_url}/download/{guest_filename}", timeout=10 ) if response.status_code == 200: with open(host_path, 'wb') as f: f.write(response.content) logger.info(f"✅ 文件收集成功（HTTP）: {guest_filename}") else: raise RuntimeError(f"HTTP下载失败: {response.status_code}") except Exception as http_error: logger.error(f"❌ HTTP收集也失败: {http_error}") raise RuntimeError(f"文件收集失败（两种方法都失败）") except Exception as e: logger.error(f"❌ 文件收集失败: {e}") raise def get_screenshot(self, retry_with_ip_detect=True): """ 获取虚拟机截图 Args: retry_with_ip_detect: 如果连接失败，是否尝试检测IP并重试 Returns: PIL.Image对象 """ try: # 绕过代理 proxies = {'http': None, 'https': None} response = requests.get(f"{self.vm_url}/screenshot", timeout=5, proxies=proxies) if response.status_code == 200: return Image.open(io.BytesIO(response.content)) else: raise RuntimeError(f"截图失败: HTTP {response.status_code}") except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e: if retry_with_ip_detect: logger.warning(f"⚠️ 截图连接失败，尝试检测并更新IP...") if self._detect_and_update_ip(): # IP已更新，重试一次 logger.info(f"🔄 使用新IP重试截图...") return self.get_screenshot(retry_with_ip_detect=False) logger.error(f"❌ 获取截图失败: {e}") raise except Exception as e: logger.error(f"❌ 获取截图失败: {e}") raise def get_screen_info(self, retry_with_ip_detect=True): """ 获取虚拟机屏幕信息（分辨率、DPI等） Args: retry_with_ip_detect: 如果连接失败，是否尝试检测IP并重试 Returns: dict: 包含screen_width, screen_height, dpi_scale等信息 """ try: # 绕过代理 proxies = {'http': None, 'https': None} response = requests.get(f"{self.vm_url}/screen_info", timeout=5, proxies=proxies) if response.status_code == 200: return response.json() else: raise RuntimeError(f"获取屏幕信息失败: HTTP {response.status_code}") except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e: if retry_with_ip_detect: logger.warning(f"⚠️ 屏幕信息连接失败，尝试检测并更新IP...") if self._detect_and_update_ip(): # IP已更新，重试一次 logger.info(f"🔄 使用新IP重试获取屏幕信息...") return self.get_screen_info(retry_with_ip_detect=False) logger.error(f"❌ 获取屏幕信息失败: {e}") raise except Exception as e: logger.error(f"❌ 获取屏幕信息失败: {e}") raise def list_desktop_files(self): """ 列出虚拟机桌面文件（用于调试） Returns: list: 文件名列表 """ try: # 绕过代理 proxies = {'http': None, 'https': None} response = requests.get(f"{self.vm_url}/list_desktop", timeout=5, proxies=proxies) if response.status_code == 200: return response.json().get('files', []) else: raise RuntimeError(f"列出文件失败: HTTP {response.status_code}") except Exception as e: logger.warning(f"⚠️ 列出桌面文件失败: {e}") return [] def send_action(self, action_type, **params): """ 发送动作到虚拟机（用于未来的Agent自动执行） Args: action_type: 动作类型 (click/type/hotkey) **params: 动作参数 """ try: # 绕过代理 proxies = {'http': None, 'https': None} payload = {"type": action_type, **params} response = requests.post( f"{self.vm_url}/action", json=payload, timeout=5, proxies=proxies ) if response.status_code == 200: logger.debug(f"动作执行成功: {action_type}") else: raise RuntimeError(f"动作执行失败: HTTP {response.status_code}") except Exception as e: logger.error(f"❌ 发送动作失败: {e}") raise def get_mouse_pos(self): """ 从虚拟机获取当前鼠标物理坐标 Returns: tuple: (x, y) 物理坐标，失败返回 (None, None) """ try: # 绕过代理 proxies = {'http': None, 'https': None} response = requests.get(f"{self.vm_url}/mouse_pos", timeout=2, proxies=proxies) if response.status_code == 200: data = response.json() return data['x'], data['y'] return None, None except Exception as e: logger.debug(f"获取VM鼠标位置失败: {e}") return None, None if __name__ == "__main__": # 测试代码 print("JadeEnv 测试") print("=" * 60) # 配置（需要根据实际情况修改） VMX_PATH = "/Volumes/Castor/虚拟机/Jade_Win_11.vmwarevm/Windows 11 64 位 ARM 2.vmx" SNAPSHOT = "Jade_Ready" VM_PASSWORD = "lizhanyuan" try: env = JadeEnv( vmx_path=VMX_PATH, snapshot_name=SNAPSHOT, vm_password=VM_PASSWORD, guest_username="lzy", guest_password="LIZHANYUAN" ) # 测试重置 print("\n测试1: 重置环境") env.reset() # 测试获取屏幕信息 print("\n测试2: 获取屏幕信息") info = env.get_screen_info() print(f" 分辨率: {info['screen_width']}x{info['screen_height']}") print(f" DPI缩放: {info['dpi_scale']}") # 测试列出桌面文件 print("\n测试3: 列出桌面文件") files = env.list_desktop_files() print(f" 桌面文件: {files[:5]}..." if len(files) > 5 else f" 桌面文件: {files}") print("\n✅ 所有测试通过！") except Exception as e: print(f"\n❌ 测试失败: {e}") import traceback traceback.print_exc()