515 lines
20 KiB
Python
515 lines
20 KiB
Python
"""
|
||
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()
|
||
|