sci-gui-agent-benchmark/scripts/tools/detect_vm_ip.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动检测VM的IP地址
使用vmrun getGuestIPAddress命令获取VM的当前IP
"""
import subprocess
import sys
import json
import os
from pathlib import Path

try:
    import requests
except ImportError:
    requests = None

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from scripts.utils.config_loader import load_config


def get_vm_ip(vmx_path, vm_password=None):
    """
    使用vmrun获取VM的IP地址
    
    Args:
        vmx_path: 虚拟机.vmx文件路径
        vm_password: 虚拟机文件加密密码（可选）
    
    Returns:
        str: VM的IP地址，如果失败返回None
    """
    vmrun = "/Applications/VMware Fusion.app/Contents/Library/vmrun"
    
    # 构建命令
    cmd = [vmrun, "-T", "fusion"]
    if vm_password:
        cmd.extend(["-vp", vm_password])
    cmd.extend(["getGuestIPAddress", vmx_path])
    
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=10
        )
        
        if result.returncode == 0:
            ip = result.stdout.strip()
            if ip and ip != "":
                return ip
            else:
                print(f"⚠️  vmrun返回空IP地址")
                return None
        else:
            error_msg = result.stderr or result.stdout
            print(f"❌ 获取IP失败: {error_msg}")
            return None
            
    except subprocess.TimeoutExpired:
        print(f"❌ 获取IP超时")
        return None
    except Exception as e:
        print(f"❌ 获取IP异常: {e}")
        return None


def update_config_ip(new_ip, config_path="config.json"):
    """
    更新config.json中的IP地址
    
    Args:
        new_ip: 新的IP地址
        config_path: 配置文件路径
    """
    config_path = os.path.join(project_root, config_path)
    
    if not os.path.exists(config_path):
        print(f"❌ 配置文件不存在: {config_path}")
        return False
    
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        old_ip = config.get("network", {}).get("vm_ip", "未知")
        
        if old_ip == new_ip:
            print(f"✅ IP地址未变化: {new_ip}")
            return True
        
        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)
        
        print(f"✅ 已更新IP地址: {old_ip} → {new_ip}")
        return True
        
    except Exception as e:
        print(f"❌ 更新配置文件失败: {e}")
        return False


def test_http_connection(ip, port=5000, timeout=5):
    """
    测试HTTP连接是否可用
    
    Args:
        ip: VM的IP地址
        port: 端口号
        timeout: 超时时间（秒）
    
    Returns:
        bool: 连接是否成功
    """
    try:
        import requests
    except ImportError:
        print(f"⚠️  requests模块未安装，跳过HTTP连接测试")
        return False
    
    url = f"http://{ip}:{port}/screen_info"
    proxies = {'http': None, 'https': None}  # 绕过代理
    
    try:
        response = requests.get(url, timeout=timeout, proxies=proxies)
        if response.status_code == 200:
            print(f"✅ HTTP服务连接成功: {url}")
            return True
        else:
            print(f"⚠️  HTTP服务响应异常: 状态码 {response.status_code}")
            return False
    except requests.exceptions.Timeout:
        print(f"⚠️  HTTP服务连接超时: {url}")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"⚠️  HTTP服务连接失败: {url}")
        print(f"   错误: {str(e)[:100]}")
        return False
    except Exception as e:
        print(f"⚠️  HTTP连接异常: {e}")
        return False


def main():
    """主函数"""
    print("=" * 60)
    print("🔍 检测VM IP地址")
    print("=" * 60)
    
    # 加载配置
    try:
        config = load_config()
        vmx_path = config["vmware"]["vmx_path"]
        vm_password = config["vmware"].get("vm_password")
        current_ip = config["network"].get("vm_ip", "未知")
    except Exception as e:
        print(f"❌ 加载配置失败: {e}")
        sys.exit(1)
    
    print(f"\n📋 当前配置:")
    print(f"   VM路径: {os.path.basename(vmx_path)}")
    print(f"   当前IP: {current_ip}")
    print(f"   端口: {config['network'].get('agent_server_port', 5000)}")
    
    # 获取VM IP
    print(f"\n🔍 正在获取VM IP地址...")
    vm_ip = get_vm_ip(vmx_path, vm_password)
    
    if not vm_ip:
        print("\n❌ 无法获取VM IP地址")
        print("   可能原因:")
        print("   1. VM未运行")
        print("   2. VM网络未配置")
        print("   3. vmrun命令执行失败")
        sys.exit(1)
    
    print(f"✅ 检测到VM IP: {vm_ip}")
    
    # 测试HTTP连接
    port = config["network"].get("agent_server_port", 5000)
    print(f"\n🔗 测试HTTP连接 (端口 {port})...")
    http_ok = test_http_connection(vm_ip, port)
    
    # 询问是否更新配置
    if vm_ip != current_ip:
        print(f"\n⚠️  IP地址已变化: {current_ip} → {vm_ip}")
        
        if http_ok:
            print(f"\n❓ 是否更新配置文件? (y/n): ", end="")
            try:
                choice = input().strip().lower()
                if choice == 'y':
                    if update_config_ip(vm_ip):
                        print(f"\n✅ 配置已更新！")
                    else:
                        print(f"\n❌ 配置更新失败")
                else:
                    print(f"\n⏭️  跳过更新")
            except KeyboardInterrupt:
                print(f"\n\n⚠️  用户取消")
        else:
            print(f"\n⚠️  HTTP服务不可用，请检查:")
            print(f"   1. VM中是否运行了 agent_server.py?")
            print(f"   2. 端口 {port} 是否被占用?")
            print(f"   3. 防火墙是否阻止了连接?")
    else:
        print(f"\n✅ IP地址未变化")
        if not http_ok:
            print(f"\n⚠️  但HTTP服务不可用，请检查VM中的agent_server.py")
    
    print("=" * 60)


if __name__ == "__main__":
    main()