lybic-guiagents 0.1.0__py3-none-any.whl

desktop_env/providers/virtualbox/manager.py

@@ -0,0 +1,463 @@
+import logging
+import os
+import platform
+import shutil
+import subprocess
+import threading
+import time
+import zipfile
+
+import psutil
+import requests
+from filelock import FileLock
+from tqdm import tqdm
+
+from desktop_env.providers.base import VMManager
+
+logger = logging.getLogger("desktopenv.providers.virtualbox.VirtualBoxVMManager")
+logger.setLevel(logging.INFO)
+
+MAX_RETRY_TIMES = 10
+RETRY_INTERVAL = 5
+UBUNTU_ARM_URL = "NOT_AVAILABLE"
+UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_x86_virtualbox/resolve/main/Ubuntu.zip"
+DOWNLOADED_FILE_NAME = "Ubuntu.zip"
+REGISTRY_PATH = '.virtualbox_vms'
+
+LOCK_FILE_NAME = '.virtualbox_lck'
+VMS_DIR = "./virtualbox_vm_data"
+update_lock = threading.Lock()
+
+if platform.system() == 'Windows':
+    vboxmanage_path = r"C:\Program Files\Oracle\VirtualBox"
+    os.environ["PATH"] += os.pathsep + vboxmanage_path
+
+
+def generate_new_vm_name(vms_dir, os_type):
+    registry_idx = 0
+    while True:
+        attempted_new_name = f"{os_type}{registry_idx}"
+        if os.path.exists(
+                os.path.join(vms_dir, attempted_new_name, attempted_new_name, attempted_new_name + ".vbox")):
+            registry_idx += 1
+        else:
+            return attempted_new_name
+
+
+def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu", bridged_adapter_name=None):
+    os.makedirs(vms_dir, exist_ok=True)
+
+    def __download_and_unzip_vm():
+        # Determine the platform and CPU architecture to decide the correct VM image to download
+        if platform.system() == 'Darwin':  # macOS
+            url = UBUNTU_ARM_URL
+            raise Exception("MacOS host is not currently supported for VirtualBox.")
+        elif platform.machine().lower() in ['amd64', 'x86_64']:
+            url = UBUNTU_X86_URL
+        else:
+            raise Exception("Unsupported platform or architecture.")
+
+        # Download the virtual machine image
+        logger.info("Downloading the virtual machine image...")
+        downloaded_size = 0
+
+        while True:
+            downloaded_file_path = os.path.join(vms_dir, downloaded_file_name)
+            headers = {}
+            if os.path.exists(downloaded_file_path):
+                downloaded_size = os.path.getsize(downloaded_file_path)
+                headers["Range"] = f"bytes={downloaded_size}-"
+
+            with requests.get(url, headers=headers, stream=True) as response:
+                if response.status_code == 416:
+                    # This means the range was not satisfiable, possibly the file was fully downloaded
+                    logger.info("Fully downloaded or the file size changed.")
+                    break
+
+                response.raise_for_status()
+                total_size = int(response.headers.get('content-length', 0))
+
+                with open(downloaded_file_path, "ab") as file, tqdm(
+                        desc="Progress",
+                        total=total_size,
+                        unit='iB',
+                        unit_scale=True,
+                        unit_divisor=1024,
+                        initial=downloaded_size,
+                        ascii=True
+                ) as progress_bar:
+                    try:
+                        for data in response.iter_content(chunk_size=1024):
+                            size = file.write(data)
+                            progress_bar.update(size)
+                    except (requests.exceptions.RequestException, IOError) as e:
+                        logger.error(f"Download error: {e}")
+                        time.sleep(RETRY_INTERVAL)
+                        logger.error("Retrying...")
+                    else:
+                        logger.info("Download succeeds.")
+                        break  # Download completed successfully
+
+        # Unzip the downloaded file
+        logger.info("Unzipping the downloaded file...☕️")
+        with zipfile.ZipFile(downloaded_file_path, 'r') as zip_ref:
+            zip_ref.extractall(vms_dir)
+        logger.info("Files have been successfully extracted to the directory: " + vms_dir)
+
+    def import_vm(vms_dir, target_vm_name, max_retries=1):
+        """Import the .ovf file into VirtualBox."""
+        logger.info(f"Starting to import VM {target_vm_name}...")
+        command = (
+            f"VBoxManage import {os.path.abspath(os.path.join(vms_dir, original_vm_name, original_vm_name + '.ovf'))} "
+            f"--vsys 0 "
+            f"--vmname {target_vm_name} "
+            f"--settingsfile {os.path.abspath(os.path.join(vms_dir, target_vm_name, target_vm_name + '.vbox'))} "
+            f"--basefolder {vms_dir} "
+            f"--unit 14 "
+            f"--disk {os.path.abspath(os.path.join(vms_dir, target_vm_name, target_vm_name + '_disk1.vmdk'))}")
+
+        for attempt in range(max_retries):
+            result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+                                    errors='ignore')
+            if result.returncode == 0:
+                logger.info("Successfully imported VM.")
+                return True
+            else:
+                if not result.stderr or "Error" in result.stderr:
+                    logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
+                else:
+                    logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
+
+                if attempt == max_retries - 1:
+                    logger.error("Maximum retry attempts reached, failed to import the virtual machine.")
+                    return False
+                
+    def configure_vm_network(vm_name, interface_name=None):
+        # Config of bridged network
+        command = f'VBoxManage modifyvm "{vm_name}" --nic1 bridged'
+        result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+                                    errors='ignore')
+        if not interface_name:
+            output = subprocess.check_output(f"VBoxManage list bridgedifs", shell=True, stderr=subprocess.STDOUT)
+            output = output.decode()
+            output = output.splitlines()
+            result = []
+            for line in output:
+                entries = line.split()
+                if entries and entries[0] == "Name:":
+                    name = ' '.join(entries[1:])
+                if entries and entries[0] == "IPAddress:":
+                    ip = entries[1]
+                    result.append((name, ip))
+            logger.info("Found the following network adapters, default to the first. If you want to change it, please set the argument -r to the name of the adapter.")
+            for i, (name, ip) in enumerate(result):
+                logger.info(f"{i+1}: {name} ({ip})")
+            interface_id = 1
+            interface_name = result[interface_id-1][0]
+        command = f'vboxmanage modifyvm "{vm_name}" --bridgeadapter1 "{interface_name}"'
+        result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+                                    errors='ignore')
+        if result.returncode == 0:
+            logger.info(f"Changed to bridge adapter {interface_name}.")
+            return True
+        else:
+            logger.error(f"Failed to change to bridge adapter {interface_name}: {result.stderr}")
+            return False
+        
+        # # Config of NAT network
+        # command = f"VBoxManage natnetwork add --netname natnet --network {nat_network} --dhcp on"
+        # result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+        #                             errors='ignore')
+        # if result.returncode == 0:
+        #     logger.info(f"Created NAT network {nat_network}.")
+        # else:
+        #     logger.error(f"Failed to create NAT network {nat_network}")
+        #     return False
+        # command = f"VBoxManage modifyvm {vm_name} --nic1 natnetwork"
+        # result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+        #                             errors='ignore')
+        # command = f"VBoxManage modifyvm {vm_name} --natnet1 natnet"
+        # result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+        #                             errors='ignore')
+        # if result.returncode == 0:
+        #     logger.info("Switched VM to the NAT network.")
+        # else:
+        #     logger.error("Failed to switch VM to the NAT network")
+        #     return False
+        # logger.info("Start to configure port forwarding...")
+        # command = f"VBoxManage modifyvm {vm_name} --natpf1 'server,tcp,,5000,,5000'"
+        # result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8",
+        #                     errors='ignore')
+        # if result.returncode == 0:
+        #     logger.info("Successfully created port forwarding rule.")
+        #     return True
+        # logger.error("Failed to create port forwarding rule.")
+        # return False
+
+
+    vm_path = os.path.join(vms_dir, vm_name, vm_name + ".vbox")
+
+    # Execute the function to download and unzip the VM, and update the vm metadata
+    if not os.path.exists(vm_path):
+        __download_and_unzip_vm()
+        import_vm(vms_dir, vm_name)
+        if not configure_vm_network(vm_name, bridged_adapter_name):
+            raise Exception("Failed to configure VM network!")
+    else:
+        logger.info(f"Virtual machine exists: {vm_path}")
+
+    # Start the virtual machine
+    def start_vm(vm_name, max_retries=20):
+        command = f'VBoxManage startvm "{vm_name}" --type headless'
+
+        for attempt in range(max_retries):
+            result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
+            if result.returncode == 0:
+                logger.info("Virtual machine started.")
+                return True
+            else:
+                if not result.stderr or "Error" in result.stderr:
+                    logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
+                else:
+                    logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
+
+                if attempt == max_retries - 1:
+                    logger.error("Maximum retry attempts reached, failed to start the virtual machine.")
+                    return False
+
+    if not start_vm(vm_name):
+        raise ValueError("Error encountered during installation, please rerun the code for retrying.")
+
+    def get_vm_ip(vm_name):
+        command = f'VBoxManage guestproperty get "{vm_name}" /VirtualBox/GuestInfo/Net/0/V4/IP'
+        result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
+        if result.returncode == 0:
+            return result.stdout.strip().split()[1]
+        else:
+            logger.error(f"Get VM IP failed: {result.stderr}")
+            return None
+        
+    def change_resolution(vm_name, resolution=(1920, 1080, 32)):
+        command = f'VBoxManage controlvm "{vm_name}" setvideomodehint {" ".join(map(str, resolution))}'
+        result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
+        if result.returncode == 0:
+            return True
+        else:
+            return False
+
+    # Function used to check whether the virtual machine is ready
+    def download_screenshot(vm_name):
+        ip = get_vm_ip(vm_name)
+        url = f"http://{ip}:5000/screenshot"
+        try:
+            # max trey times 1, max timeout 1
+            response = requests.get(url, timeout=(10, 10))
+            if response.status_code == 200:
+                return True
+        except Exception as e:
+            logger.error(f"Error: {e}")
+            logger.error(f"Type: {type(e).__name__}")
+            logger.error(f"Error detail: {str(e)}")
+        return False
+
+    # Try downloading the screenshot until successful
+    while not download_screenshot(vm_name):
+        logger.info("Check whether the virtual machine is ready...")
+        time.sleep(RETRY_INTERVAL)
+
+    if not change_resolution(vm_name):
+        logger.error(f"Change resolution failed.")
+        raise
+
+    logger.info("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")
+
+    def create_vm_snapshot(vm_name, max_retries=20):
+        logger.info("Saving VirtualBox VM state...")
+        command = f'VBoxManage snapshot "{vm_name}" take init_state'
+
+        for attempt in range(max_retries):
+            result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
+            if result.returncode == 0:
+                logger.info("Snapshot created.")
+                return True
+            else:
+                if "Error" in result.stderr:
+                    logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
+                else:
+                    logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
+
+                if attempt == max_retries - 1:
+                    logger.error("Maximum retry attempts reached, failed to create snapshot.")
+                    return False
+
+    # Create a snapshot of the virtual machine
+    if create_vm_snapshot(vm_name, max_retries=MAX_RETRY_TIMES):
+        return vm_path
+    else:
+        raise ValueError("Error encountered during installation, please rerun the code for retrying.")
+
+
+class VirtualBoxVMManager(VMManager):
+    def __init__(self, registry_path=REGISTRY_PATH):
+        self.registry_path = registry_path
+        self.lock = FileLock(LOCK_FILE_NAME, timeout=60)
+        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, lock_needed=True):
+        if lock_needed:
+            with self.lock:
+                self._add_vm(vm_path)
+        else:
+            self._add_vm(vm_path)
+
+    def _add_vm(self, vm_path, region=None):
+        assert region in [None, 'local'], "For VirtualBox provider, the region should be neither None or 'local'."
+        with self.lock:
+            with open(self.registry_path, 'r') as file:
+                lines = file.readlines()
+                new_lines = lines + [f'{vm_path}|free\n']
+            with open(self.registry_path, 'w') as file:
+                file.writelines(new_lines)
+
+    def occupy_vm(self, vm_path, pid, lock_needed=True):
+        if lock_needed:
+            with self.lock:
+                self._occupy_vm(vm_path, pid)
+        else:
+            self._occupy_vm(vm_path, pid)
+
+    def _occupy_vm(self, vm_path, pid, region=None):
+        assert region in [None, 'local'], "For VirtualBox provider, the region should be neither None or 'local'."
+        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 == vm_path:
+                        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 delete_vm(self, vm_path, lock_needed=True):
+        if lock_needed:
+            with self.lock:
+                self._delete_vm(vm_path)
+        else:
+            self._delete_vm(vm_path)
+
+    def _delete_vm(self, vm_path):
+        raise NotImplementedError
+
+    def check_and_clean(self, vms_dir, lock_needed=True):
+        if lock_needed:
+            with self.lock:
+                self._check_and_clean(vms_dir)
+        else:
+            self._check_and_clean(vms_dir)
+
+    def _check_and_clean(self, vms_dir):
+        with self.lock:  # Lock when cleaning up the registry and vms_dir
+            # Check and clean on the running vms, detect the released ones and mark then as 'free'
+            active_pids = {p.pid for p in psutil.process_iter()}
+            new_lines = []
+            vm_paths = []
+
+            with open(self.registry_path, 'r') as file:
+                lines = file.readlines()
+                for line in lines:
+                    vm_path, pid_str = line.strip().split('|')
+                    if not os.path.exists(vm_path):
+                        logger.info(f"VM {vm_path} not found, releasing it.")
+                        new_lines.append(f'{vm_path}|free\n')
+                        continue
+
+                    vm_paths.append(vm_path)
+                    if pid_str == "free":
+                        new_lines.append(line)
+                        continue
+
+                    if int(pid_str) in active_pids:
+                        new_lines.append(line)
+                    else:
+                        new_lines.append(f'{vm_path}|free\n')
+            with open(self.registry_path, 'w') as file:
+                file.writelines(new_lines)
+
+            # Check and clean on the files inside vms_dir, delete the unregistered ones
+            os.makedirs(vms_dir, exist_ok=True)
+            vm_names = os.listdir(vms_dir)
+            for vm_name in vm_names:
+                # skip the downloaded .zip file
+                if vm_name == DOWNLOADED_FILE_NAME:
+                    continue
+                # Skip the .DS_Store file on macOS
+                if vm_name == ".DS_Store":
+                    continue
+
+                flag = True
+                for vm_path in vm_paths:
+                    if vm_name + ".vbox" in vm_path:
+                        flag = False
+                if flag:
+                    shutil.rmtree(os.path.join(vms_dir, vm_name))
+
+    def list_free_vms(self, lock_needed=True):
+        if lock_needed:
+            with self.lock:
+                return self._list_free_vms()
+        else:
+            return self._list_free_vms()
+
+    def _list_free_vms(self):
+        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, pid_str = line.strip().split('|')
+                    if pid_str == "free":
+                        free_vms.append((vm_path, pid_str))
+            return free_vms
+
+    def get_vm_path(self, os_type, region=None, screen_size=(1920, 1080), **kwargs):
+        # Note: screen_size parameter is ignored for VirtualBox provider
+        # but kept for interface consistency with other providers
+        if os_type != "Ubuntu":
+            raise ValueError("Only support Ubuntu for now.")
+
+        with self.lock:
+            if not VirtualBoxVMManager.checked_and_cleaned:
+                VirtualBoxVMManager.checked_and_cleaned = True
+                self._check_and_clean(vms_dir=VMS_DIR)
+
+        allocation_needed = False
+        with self.lock:
+            free_vms_paths = self._list_free_vms()
+            if len(free_vms_paths) == 0:
+                # No free virtual machine available, generate a new one
+                allocation_needed = True
+            else:
+                # Choose the first free virtual machine
+                chosen_vm_path = free_vms_paths[0][0]
+                self._occupy_vm(chosen_vm_path, os.getpid())
+                return chosen_vm_path
+            
+        if allocation_needed:
+            logger.info("No free virtual machine available. Generating a new one, which would take a while...☕")
+            new_vm_name = generate_new_vm_name(vms_dir=VMS_DIR, os_type=os_type)
+            new_vm_path = _install_vm(new_vm_name, vms_dir=VMS_DIR,
+                                    downloaded_file_name=DOWNLOADED_FILE_NAME,
+                                    bridged_adapter_name=region)
+            with self.lock:
+                self._add_vm(new_vm_path)
+                self._occupy_vm(new_vm_path, os.getpid())
+            return new_vm_path

desktop_env/providers/virtualbox/provider.py

@@ -0,0 +1,124 @@
+import logging
+import platform
+import subprocess
+import time
+import os
+from desktop_env.providers.base import Provider
+import xml.etree.ElementTree as ET
+
+logger = logging.getLogger("desktopenv.providers.virtualbox.VirtualBoxProvider")
+logger.setLevel(logging.INFO)
+
+WAIT_TIME = 3
+
+# Note: Windows will not add command VBoxManage to PATH by default. Please add the folder where VBoxManage executable is in (Default should be "C:\Program Files\Oracle\VirtualBox" for Windows) to PATH.
+
+class VirtualBoxProvider(Provider):
+    @staticmethod
+    def _execute_command(command: list):
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60, text=True,
+                                encoding="utf-8")
+        if result.returncode != 0:
+            raise Exception("\033[91m" + result.stdout + result.stderr + "\033[0m")
+        return result.stdout.strip()
+    
+    @staticmethod
+    def _get_vm_uuid(path_to_vm: str):
+        try:
+            output = subprocess.check_output(f"VBoxManage list vms", shell=True, stderr=subprocess.STDOUT)
+            output = output.decode()
+            output = output.splitlines()
+            if path_to_vm.endswith('.vbox'):
+                # Load and parse the XML content from the file
+                tree = ET.parse(path_to_vm)
+                root = tree.getroot()
+                
+                # Find the <Machine> element and retrieve its 'uuid' attribute
+                machine_element = root.find('.//{http://www.virtualbox.org/}Machine')
+                if machine_element is not None:
+                    uuid = machine_element.get('uuid')[1:-1]
+                    return uuid
+                else:
+                    logger.error(f"UUID not found in file {path_to_vm}")
+                    raise
+            elif any(line.split()[1] == "{" + path_to_vm + "}" for line in output):
+                logger.info(f"Got valid UUID {path_to_vm}.")
+                return path_to_vm
+            else:
+                for line in output:
+                    if line.split()[0] == '"' + path_to_vm + '"':
+                        uuid = line.split()[1][1:-1]
+                        return uuid
+                logger.error(f"The path you provided does not match any of the \".vbox\" file, name, or UUID of VM.")
+                raise
+        except subprocess.CalledProcessError as e:
+                logger.error(f"Error executing command: {e.output.decode().strip()}")
+            
+
+    def start_emulator(self, path_to_vm: str, headless: bool, os_type: str = None, *args, **kwargs):
+        # Note: os_type parameter is ignored for VirtualBox provider
+        # but kept for interface consistency with other providers
+        logger.info("Starting VirtualBox VM...")
+
+        while True:
+            try:
+                uuid = VirtualBoxProvider._get_vm_uuid(path_to_vm)
+                output = subprocess.check_output(f"VBoxManage list runningvms", shell=True, stderr=subprocess.STDOUT)
+                output = output.decode()
+                output = output.splitlines()
+
+                if any(line.split()[1] == "{" + uuid + "}" for line in output):
+                    logger.info("VM is running.")
+                    break
+                else:
+                    logger.info("Starting VM...")
+                    VirtualBoxProvider._execute_command(["VBoxManage", "startvm", uuid]) if not headless else \
+                    VirtualBoxProvider._execute_command(
+                            ["VBoxManage", "startvm", uuid, "--type", "headless"])
+                    time.sleep(WAIT_TIME)
+
+            except subprocess.CalledProcessError as e:
+                logger.error(f"Error executing command: {e.output.decode().strip()}")
+
+    def get_ip_address(self, path_to_vm: str) -> str:
+        logger.info("Getting VirtualBox VM IP address...")
+        while True:
+            try:
+                uuid = VirtualBoxProvider._get_vm_uuid(path_to_vm)
+                output = VirtualBoxProvider._execute_command(
+                    ["VBoxManage", "guestproperty", "get", uuid, "/VirtualBox/GuestInfo/Net/0/V4/IP"]
+                )
+                result = output.split()[1]
+                if result != "value":
+                    logger.info(f"VirtualBox VM IP address: {result}")
+                    return result
+                else:
+                    logger.error("VM IP address not found. Have you installed the guest additions?")
+                    raise
+            except Exception as e:
+                logger.error(e)
+                time.sleep(WAIT_TIME)
+                logger.info("Retrying to get VirtualBox VM IP address...")
+
+    def save_state(self, path_to_vm: str, snapshot_name: str):
+        logger.info("Saving VirtualBox VM state...")
+        uuid = VirtualBoxProvider._get_vm_uuid(path_to_vm)
+        VirtualBoxProvider._execute_command(["VBoxManage", "snapshot", uuid, "take", snapshot_name])
+        time.sleep(WAIT_TIME)  # Wait for the VM to save
+
+    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
+        logger.info(f"Reverting VirtualBox VM to snapshot: {snapshot_name}...")
+        uuid = VirtualBoxProvider._get_vm_uuid(path_to_vm)
+        VirtualBoxProvider._execute_command(["VBoxManage", "controlvm", uuid, "savestate"])
+        time.sleep(WAIT_TIME)  # Wait for the VM to stop
+        VirtualBoxProvider._execute_command(["VBoxManage", "snapshot", uuid, "restore", snapshot_name])
+        time.sleep(WAIT_TIME)  # Wait for the VM to revert
+        return path_to_vm
+
+    def stop_emulator(self, path_to_vm: str, region=None, *args, **kwargs):
+        # Note: region parameter is ignored for VirtualBox provider
+        # but kept for interface consistency with other providers
+        logger.info("Stopping VirtualBox VM...")
+        uuid = VirtualBoxProvider._get_vm_uuid(path_to_vm)
+        VirtualBoxProvider._execute_command(["VBoxManage", "controlvm", uuid, "savestate"])
+        time.sleep(WAIT_TIME)  # Wait for the VM to stop