lybic-guiagents 0.1.0__py3-none-any.whl

desktop_env/providers/aws/proxy_pool.py

@@ -0,0 +1,193 @@
+import random
+import requests
+import logging
+import time
+from typing import List, Dict, Optional
+from dataclasses import dataclass
+from threading import Lock
+import json
+
+logger = logging.getLogger("desktopenv.providers.aws.ProxyPool")
+logger.setLevel(logging.INFO)
+
+@dataclass
+class ProxyInfo:
+    host: str
+    port: int
+    username: Optional[str] = None
+    password: Optional[str] = None
+    protocol: str = "http"  # http, https, socks5
+    failed_count: int = 0
+    last_used: float = 0
+    is_active: bool = True
+
+class ProxyPool:
+    def __init__(self, config_file: str = None):
+        self.proxies: List[ProxyInfo] = []
+        self.current_index = 0
+        self.lock = Lock()
+        self.max_failures = 3  # 最大失败次数
+        self.cooldown_time = 300  # 5分钟冷却时间
+        
+        if config_file:
+            self.load_proxies_from_file(config_file)
+    
+    def load_proxies_from_file(self, config_file: str):
+        """从配置文件加载代理列表"""
+        try:
+            with open(config_file, 'r') as f:
+                proxy_configs = json.load(f)
+                
+            for config in proxy_configs:
+                proxy = ProxyInfo(
+                    host=config['host'],
+                    port=config['port'],
+                    username=config.get('username'),
+                    password=config.get('password'),
+                    protocol=config.get('protocol', 'http')
+                )
+                self.proxies.append(proxy)
+                
+            logger.info(f"Loaded {len(self.proxies)} proxies from {config_file}")
+        except Exception as e:
+            logger.error(f"Failed to load proxies from {config_file}: {e}")
+    
+    def add_proxy(self, host: str, port: int, username: str = None, 
+                  password: str = None, protocol: str = "http"):
+        """添加代理到池中"""
+        proxy = ProxyInfo(host=host, port=port, username=username, 
+                         password=password, protocol=protocol)
+        with self.lock:
+            self.proxies.append(proxy)
+        logger.info(f"Added proxy {host}:{port}")
+    
+    def get_next_proxy(self) -> Optional[ProxyInfo]:
+        """获取下一个可用的代理"""
+        with self.lock:
+            if not self.proxies:
+                return None
+            
+            # 过滤掉失败次数过多的代理
+            active_proxies = [p for p in self.proxies if self._is_proxy_available(p)]
+            
+            if not active_proxies:
+                logger.warning("No active proxies available")
+                return None
+            
+            # 轮询选择代理
+            proxy = active_proxies[self.current_index % len(active_proxies)]
+            self.current_index += 1
+            proxy.last_used = time.time()
+            
+            return proxy
+    
+    def _is_proxy_available(self, proxy: ProxyInfo) -> bool:
+        """检查代理是否可用"""
+        if not proxy.is_active:
+            return False
+        
+        if proxy.failed_count >= self.max_failures:
+            # 检查是否过了冷却时间
+            if time.time() - proxy.last_used < self.cooldown_time:
+                return False
+            else:
+                # 重置失败计数
+                proxy.failed_count = 0
+        
+        return True
+    
+    def mark_proxy_failed(self, proxy: ProxyInfo):
+        """标记代理失败"""
+        with self.lock:
+            proxy.failed_count += 1
+            if proxy.failed_count >= self.max_failures:
+                logger.warning(f"Proxy {proxy.host}:{proxy.port} marked as failed "
+                             f"(failures: {proxy.failed_count})")
+    
+    def mark_proxy_success(self, proxy: ProxyInfo):
+        """标记代理成功"""
+        with self.lock:
+            proxy.failed_count = 0
+    
+    def test_proxy(self, proxy: ProxyInfo, test_url: str = "http://httpbin.org/ip", 
+                   timeout: int = 10) -> bool:
+        """测试代理是否正常工作"""
+        try:
+            proxy_url = self._format_proxy_url(proxy)
+            proxies = {
+                'http': proxy_url,
+                'https': proxy_url
+            }
+            
+            response = requests.get(test_url, proxies=proxies, timeout=timeout)
+            if response.status_code == 200:
+                self.mark_proxy_success(proxy)
+                return True
+            else:
+                self.mark_proxy_failed(proxy)
+                return False
+                
+        except Exception as e:
+            logger.debug(f"Proxy test failed for {proxy.host}:{proxy.port}: {e}")
+            self.mark_proxy_failed(proxy)
+            return False
+    
+    def _format_proxy_url(self, proxy: ProxyInfo) -> str:
+        """格式化代理URL"""
+        if proxy.username and proxy.password:
+            return f"{proxy.protocol}://{proxy.username}:{proxy.password}@{proxy.host}:{proxy.port}"
+        else:
+            return f"{proxy.protocol}://{proxy.host}:{proxy.port}"
+    
+    def get_proxy_dict(self, proxy: ProxyInfo) -> Dict[str, str]:
+        """获取requests库使用的代理字典"""
+        proxy_url = self._format_proxy_url(proxy)
+        return {
+            'http': proxy_url,
+            'https': proxy_url
+        }
+    
+    def test_all_proxies(self, test_url: str = "http://httpbin.org/ip"):
+        """测试所有代理"""
+        logger.info("Testing all proxies...")
+        working_count = 0
+        
+        for proxy in self.proxies:
+            if self.test_proxy(proxy, test_url):
+                working_count += 1
+                logger.info(f"✓ Proxy {proxy.host}:{proxy.port} is working")
+            else:
+                logger.warning(f"✗ Proxy {proxy.host}:{proxy.port} failed")
+        
+        logger.info(f"Proxy test completed: {working_count}/{len(self.proxies)} working")
+        return working_count
+    
+    def get_stats(self) -> Dict:
+        """获取代理池统计信息"""
+        with self.lock:
+            total = len(self.proxies)
+            active = len([p for p in self.proxies if self._is_proxy_available(p)])
+            failed = len([p for p in self.proxies if p.failed_count >= self.max_failures])
+            
+            return {
+                'total': total,
+                'active': active,
+                'failed': failed,
+                'success_rate': active / total if total > 0 else 0
+            }
+
+# 全局代理池实例
+_proxy_pool = None
+
+def get_global_proxy_pool() -> ProxyPool:
+    """获取全局代理池实例"""
+    global _proxy_pool
+    if _proxy_pool is None:
+        _proxy_pool = ProxyPool()
+    return _proxy_pool
+
+def init_proxy_pool(config_file: str = None):
+    """初始化全局代理池"""
+    global _proxy_pool
+    _proxy_pool = ProxyPool(config_file)
+    return _proxy_pool 

desktop_env/providers/azure/manager.py

@@ -0,0 +1,87 @@
+import os
+import threading
+import boto3
+import psutil
+
+import logging
+
+from desktop_env.providers.base import VMManager
+
+logger = logging.getLogger("desktopenv.providers.azure.AzureVMManager")
+logger.setLevel(logging.INFO)
+
+REGISTRY_PATH = '.azure_vms'
+
+
+def _allocate_vm(region):
+    raise NotImplementedError
+
+
+class AzureVMManager(VMManager):
+    def __init__(self, registry_path=REGISTRY_PATH):
+        self.registry_path = registry_path
+        self.lock = threading.Lock()
+        self.initialize_registry()
+
+    def initialize_registry(self):
+        with self.lock:  # Locking during initialization
+            if not os.path.exists(self.registry_path):
+                with open(self.registry_path, 'w') as file:
+                    file.write('')
+
+    def add_vm(self, vm_path, region):
+        with self.lock:
+            with open(self.registry_path, 'r') as file:
+                lines = file.readlines()
+                vm_path_at_vm_region = "{}@{}".format(vm_path, region)
+                new_lines = lines + [f'{vm_path_at_vm_region}|free\n']
+            with open(self.registry_path, 'w') as file:
+                file.writelines(new_lines)
+
+    def occupy_vm(self, vm_path, pid, region):
+        with self.lock:
+            new_lines = []
+            with open(self.registry_path, 'r') as file:
+                lines = file.readlines()
+                for line in lines:
+                    registered_vm_path, _ = line.strip().split('|')
+                    if registered_vm_path == "{}@{}".format(vm_path, region):
+                        new_lines.append(f'{registered_vm_path}|{pid}\n')
+                    else:
+                        new_lines.append(line)
+            with open(self.registry_path, 'w') as file:
+                file.writelines(new_lines)
+
+    def check_and_clean(self):
+        raise NotImplementedError
+
+    def list_free_vms(self, region):
+        with self.lock:  # Lock when reading the registry
+            free_vms = []
+            with open(self.registry_path, 'r') as file:
+                lines = file.readlines()
+                for line in lines:
+                    vm_path_at_vm_region, pid_str = line.strip().split('|')
+                    vm_path, vm_region = vm_path_at_vm_region.split("@")
+                    if pid_str == "free" and vm_region == region:
+                        free_vms.append((vm_path, pid_str))
+            return free_vms
+
+    def get_vm_path(self, region, screen_size=(1920, 1080), **kwargs):
+        # Note: screen_size parameter is ignored for Azure provider
+        # but kept for interface consistency with other providers
+        self.check_and_clean()
+        free_vms_paths = self.list_free_vms(region)
+        if len(free_vms_paths) == 0:
+            # No free virtual machine available, generate a new one
+            logger.info("No free virtual machine available. Generating a new one, which would take a while...☕")
+            new_vm_path = _allocate_vm(region)
+            self.add_vm(new_vm_path, region)
+            self.occupy_vm(new_vm_path, os.getpid(), region)
+            return new_vm_path
+        else:
+            # Choose the first free virtual machine
+            chosen_vm_path = free_vms_paths[0][0]
+            self.occupy_vm(chosen_vm_path, os.getpid(), region)
+            return chosen_vm_path
+

desktop_env/providers/azure/provider.py

@@ -0,0 +1,207 @@
+import os
+import time
+from azure.identity import DefaultAzureCredential
+from azure.mgmt.compute import ComputeManagementClient
+from azure.mgmt.network import NetworkManagementClient
+from azure.core.exceptions import ResourceNotFoundError
+
+import logging
+
+from desktop_env.providers.base import Provider
+
+logger = logging.getLogger("desktopenv.providers.azure.AzureProvider")
+logger.setLevel(logging.INFO)
+
+WAIT_DELAY = 15
+MAX_ATTEMPTS = 10
+
+# To use the Azure provider, download azure-cli by https://learn.microsoft.com/en-us/cli/azure/install-azure-cli,
+# use "az login" to log into you Azure account,
+# and set environment variable "AZURE_SUBSCRIPTION_ID" to your subscription ID.
+# Provide your resource group name and VM name in the format "RESOURCE_GROUP_NAME/VM_NAME" and pass as an argument for "-p".
+
+class AzureProvider(Provider):
+    def __init__(self, region: str = None):
+        super().__init__(region)
+        credential = DefaultAzureCredential()
+        try:
+            self.subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
+        except:
+            logger.error("Azure subscription ID not found. Please set environment variable \"AZURE_SUBSCRIPTION_ID\".")
+            raise
+        self.compute_client = ComputeManagementClient(credential, self.subscription_id)
+        self.network_client = NetworkManagementClient(credential, self.subscription_id)
+
+    def start_emulator(self, path_to_vm: str, headless: bool, os_type: str = None, *args, **kwargs):
+        # Note: os_type parameter is ignored for Azure provider
+        # but kept for interface consistency with other providers
+        logger.info("Starting Azure VM...")
+        resource_group_name, vm_name = path_to_vm.split('/')
+
+        vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name, expand='instanceView')
+        power_state = vm.instance_view.statuses[-1].code
+        if power_state == "PowerState/running":
+            logger.info("VM is already running.")
+            return
+        
+        try:
+            # Start the instance
+            for _ in range(MAX_ATTEMPTS):
+                async_vm_start = self.compute_client.virtual_machines.begin_start(resource_group_name, vm_name)
+                logger.info(f"VM {path_to_vm} is starting...")
+                # Wait for the instance to start
+                async_vm_start.wait(timeout=WAIT_DELAY)
+                vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name, expand='instanceView')
+                power_state = vm.instance_view.statuses[-1].code
+                if power_state == "PowerState/running":
+                    logger.info(f"VM {path_to_vm} is already running.")
+                    break
+        except Exception as e:
+            logger.error(f"Failed to start the Azure VM {path_to_vm}: {str(e)}")
+            raise
+
+    def get_ip_address(self, path_to_vm: str) -> str:
+        logger.info("Getting Azure VM IP address...")
+        resource_group_name, vm_name = path_to_vm.split('/')
+
+        vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name)
+
+        for interface in vm.network_profile.network_interfaces:
+            name=" ".join(interface.id.split('/')[-1:])
+            sub="".join(interface.id.split('/')[4])
+
+            try:
+                thing=self.network_client.network_interfaces.get(sub, name).ip_configurations
+
+                network_card_id = thing[0].public_ip_address.id.split('/')[-1]
+                public_ip_address = self.network_client.public_ip_addresses.get(resource_group_name, network_card_id)
+                logger.info(f"VM IP address is {public_ip_address.ip_address}")
+                return public_ip_address.ip_address
+
+            except Exception as e:
+                logger.error(f"Cannot get public IP for VM {path_to_vm}")
+                raise
+
+    def save_state(self, path_to_vm: str, snapshot_name: str):
+        print("Saving Azure VM state...")
+        resource_group_name, vm_name = path_to_vm.split('/')
+
+        vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name)
+
+        try:
+            # Backup each disk attached to the VM
+            for disk in vm.storage_profile.data_disks + [vm.storage_profile.os_disk]:
+                # Create a snapshot of the disk
+                snapshot = {
+                    'location': vm.location,
+                    'creation_data': {
+                        'create_option': 'Copy',
+                        'source_uri': disk.managed_disk.id
+                    }
+                }
+                async_snapshot_creation = self.compute_client.snapshots.begin_create_or_update(resource_group_name, snapshot_name, snapshot)
+                async_snapshot_creation.wait(timeout=WAIT_DELAY)
+
+            logger.info(f"Successfully created snapshot {snapshot_name} for VM {path_to_vm}.")
+        except Exception as e:
+            logger.error(f"Failed to create snapshot {snapshot_name} of the Azure VM {path_to_vm}: {str(e)}")
+            raise
+
+    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
+        logger.info(f"Reverting VM to snapshot: {snapshot_name}...")
+        resource_group_name, vm_name = path_to_vm.split('/')
+
+        vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name)
+
+        # Stop the VM for disk creation
+        logger.info(f"Stopping VM: {vm_name}")
+        async_vm_stop = self.compute_client.virtual_machines.begin_deallocate(resource_group_name, vm_name)
+        async_vm_stop.wait(timeout=WAIT_DELAY)  # Wait for the VM to stop
+
+        try:
+            # Get the snapshot
+            snapshot = self.compute_client.snapshots.get(resource_group_name, snapshot_name)
+
+            # Get the original disk information
+            original_disk_id = vm.storage_profile.os_disk.managed_disk.id
+            disk_name = original_disk_id.split('/')[-1]
+            if disk_name[-1] in ['0', '1']:
+                new_disk_name = disk_name[:-1] + str(int(disk_name[-1])^1)
+            else:
+                new_disk_name = disk_name + "0"
+
+            # Delete the disk if it exists
+            self.compute_client.disks.begin_delete(resource_group_name, new_disk_name).wait(timeout=WAIT_DELAY)
+
+            # Make sure the disk is deleted before proceeding to the next step
+            disk_deleted = False
+            polling_interval = 10
+            attempts = 0
+            while not disk_deleted and attempts < MAX_ATTEMPTS:
+                try:
+                    self.compute_client.disks.get(resource_group_name, new_disk_name)
+                    # If the above line does not raise an exception, the disk still exists
+                    time.sleep(polling_interval)
+                    attempts += 1
+                except ResourceNotFoundError:
+                    disk_deleted = True
+
+            if not disk_deleted:
+                logger.error(f"Disk {new_disk_name} deletion timed out.")
+                raise
+
+            # Create a new managed disk from the snapshot
+            snapshot = self.compute_client.snapshots.get(resource_group_name, snapshot_name)
+            disk_creation = {
+                'location': snapshot.location,
+                'creation_data': {
+                    'create_option': 'Copy',
+                    'source_resource_id': snapshot.id
+                },
+                'zones': vm.zones if vm.zones else None  # Preserve the original disk's zone
+            }
+            async_disk_creation = self.compute_client.disks.begin_create_or_update(resource_group_name, new_disk_name, disk_creation)
+            restored_disk = async_disk_creation.result()  # Wait for the disk creation to complete
+
+            vm.storage_profile.os_disk = {
+                'create_option': vm.storage_profile.os_disk.create_option,
+                'managed_disk': {
+                    'id': restored_disk.id
+                }
+            }
+
+            async_vm_creation = self.compute_client.virtual_machines.begin_create_or_update(resource_group_name, vm_name, vm)
+            async_vm_creation.wait(timeout=WAIT_DELAY)
+
+            # Delete the original disk
+            self.compute_client.disks.begin_delete(resource_group_name, disk_name).wait()
+
+            logger.info(f"Successfully reverted to snapshot {snapshot_name}.")
+        except Exception as e:
+            logger.error(f"Failed to revert the Azure VM {path_to_vm} to snapshot {snapshot_name}: {str(e)}")
+            raise
+
+    def stop_emulator(self, path_to_vm, region=None):
+        logger.info(f"Stopping Azure VM {path_to_vm}...")
+        resource_group_name, vm_name = path_to_vm.split('/')
+
+        vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name, expand='instanceView')
+        power_state = vm.instance_view.statuses[-1].code
+        if power_state == "PowerState/deallocated":
+            print("VM is already stopped.")
+            return
+
+        try:
+            for _ in range(MAX_ATTEMPTS):
+                async_vm_deallocate = self.compute_client.virtual_machines.begin_deallocate(resource_group_name, vm_name)
+                logger.info(f"Stopping VM {path_to_vm}...")
+                # Wait for the instance to start
+                async_vm_deallocate.wait(timeout=WAIT_DELAY)
+                vm = self.compute_client.virtual_machines.get(resource_group_name, vm_name, expand='instanceView')
+                power_state = vm.instance_view.statuses[-1].code
+                if power_state == "PowerState/deallocated":
+                    logger.info(f"VM {path_to_vm} is already stopped.")
+                    break
+        except Exception as e:
+            logger.error(f"Failed to stop the Azure VM {path_to_vm}: {str(e)}")
+            raise

desktop_env/providers/base.py

@@ -0,0 +1,97 @@
+from abc import ABC, abstractmethod
+
+
+class Provider(ABC):
+    def __init__(self, region: str = None):
+        """
+        Region of the cloud service.
+        """
+        self.region = region
+
+    @abstractmethod
+    def start_emulator(self, path_to_vm: str, headless: bool):
+        """
+        Method to start the emulator.
+        """
+        pass
+
+    @abstractmethod
+    def get_ip_address(self, path_to_vm: str) -> str:
+        """
+        Method to get the private IP address of the VM. Private IP means inside the VPC.
+        """
+        pass
+
+    @abstractmethod
+    def save_state(self, path_to_vm: str, snapshot_name: str):
+        """
+        Method to save the state of the VM.
+        """
+        pass
+
+    @abstractmethod
+    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str) -> str:
+        """
+        Method to revert the VM to a given snapshot.
+        """
+        pass
+
+    @abstractmethod
+    def stop_emulator(self, path_to_vm: str):
+        """
+        Method to stop the emulator.
+        """
+        pass
+
+
+class VMManager(ABC):
+    checked_and_cleaned = False
+
+    @abstractmethod
+    def initialize_registry(self, **kwargs):
+        """
+        Initialize registry.
+        """
+        pass
+
+    @abstractmethod
+    def add_vm(self, vm_path, **kwargs):
+        """
+        Add the path of new VM to the registration.
+        """
+        pass
+
+    @abstractmethod
+    def delete_vm(self, vm_path, **kwargs):
+        """
+        Delete the registration of VM by path.
+        """
+        pass
+
+    @abstractmethod
+    def occupy_vm(self, vm_path, pid, **kwargs):
+        """
+        Mark the path of VM occupied by the pid.
+        """
+        pass
+
+    @abstractmethod
+    def list_free_vms(self, **kwargs):
+        """
+        List the paths of VM that are free to use allocated.
+        """
+        pass
+
+    @abstractmethod
+    def check_and_clean(self, **kwargs):
+        """
+        Check the registration list, and remove the paths of VM that are not in use.
+        """
+        pass
+
+    @abstractmethod
+    def get_vm_path(self, **kwargs):
+        """
+        Get a virtual machine that is not occupied, generate a new one if no free VM.
+        """
+        pass