388 lines
16 KiB
Python
388 lines
16 KiB
Python
import os
|
|
import platform
|
|
import random
|
|
import re
|
|
import subprocess
|
|
import threading
|
|
import uuid
|
|
import zipfile
|
|
from time import sleep
|
|
import shutil
|
|
import psutil
|
|
import requests
|
|
from tqdm import tqdm
|
|
|
|
__version__ = "0.1.12"
|
|
|
|
MAX_RETRY_TIMES = 10
|
|
UBUNTU_ARM_URL = "https://huggingface.co/datasets/xlangai/ubuntu_arm/resolve/main/Ubuntu.zip"
|
|
UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_x86/resolve/main/Ubuntu.zip"
|
|
DOWNLOADED_FILE_NAME = "Ubuntu.zip"
|
|
REGISTRY_PATH = '.vms'
|
|
VMS_DIR = "./vm_data"
|
|
update_lock = threading.Lock()
|
|
|
|
|
|
class VirtualMachineManager:
|
|
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):
|
|
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):
|
|
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 release_vm(self, vm_path):
|
|
with self.lock: # Lock when modifying the registry
|
|
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(line)
|
|
else:
|
|
new_lines.append(f'{registered_vm_path}|free\n')
|
|
with open(self.registry_path, 'w') as file:
|
|
file.writelines(new_lines)
|
|
|
|
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('|')
|
|
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 + ".vmx" in vm_path:
|
|
flag = False
|
|
if flag:
|
|
shutil.rmtree(os.path.join(vms_dir, vm_name))
|
|
|
|
def list_vms(self):
|
|
with self.lock: # Lock when reading the registry
|
|
all_vms = []
|
|
with open(self.registry_path, 'r') as file:
|
|
lines = file.readlines()
|
|
for line in lines:
|
|
vm_path, pid_str = line.strip().split('|')
|
|
all_vms.append((vm_path, pid_str))
|
|
return all_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 generate_new_vm_name(self, vms_dir):
|
|
registry_idx = 0
|
|
while True:
|
|
attempted_new_name = f"Ubuntu{registry_idx}"
|
|
if os.path.exists(
|
|
os.path.join(vms_dir, attempted_new_name, attempted_new_name, attempted_new_name + ".vmx")):
|
|
registry_idx += 1
|
|
else:
|
|
return attempted_new_name
|
|
|
|
|
|
def _update_vm(vmx_path, target_vm_name):
|
|
"""Update the VMX file with the new VM name and other parameters, so that the VM can be started successfully without conflict with the original VM."""
|
|
with update_lock:
|
|
dir_path, vmx_file = os.path.split(vmx_path)
|
|
|
|
def _generate_mac_address():
|
|
# VMware MAC address range starts with 00:0c:29
|
|
mac = [0x00, 0x0c, 0x29,
|
|
random.randint(0x00, 0x7f),
|
|
random.randint(0x00, 0xff),
|
|
random.randint(0x00, 0xff)]
|
|
return ':'.join(map(lambda x: "%02x" % x, mac))
|
|
|
|
# Backup the original file
|
|
with open(vmx_path, 'r') as file:
|
|
original_content = file.read()
|
|
|
|
# Generate new values
|
|
new_uuid_bios = str(uuid.uuid4())
|
|
new_uuid_location = str(uuid.uuid4())
|
|
new_mac_address = _generate_mac_address()
|
|
new_vmci_id = str(random.randint(-2147483648, 2147483647)) # Random 32-bit integer
|
|
|
|
# Update the content
|
|
updated_content = re.sub(r'displayName = ".*?"', f'displayName = "{target_vm_name}"', original_content)
|
|
updated_content = re.sub(r'uuid.bios = ".*?"', f'uuid.bios = "{new_uuid_bios}"', updated_content)
|
|
updated_content = re.sub(r'uuid.location = ".*?"', f'uuid.location = "{new_uuid_location}"', updated_content)
|
|
updated_content = re.sub(r'ethernet0.generatedAddress = ".*?"',
|
|
f'ethernet0.generatedAddress = "{new_mac_address}"',
|
|
updated_content)
|
|
updated_content = re.sub(r'vmci0.id = ".*?"', f'vmci0.id = "{new_vmci_id}"', updated_content)
|
|
|
|
# Write the updated content back to the file
|
|
with open(vmx_path, 'w') as file:
|
|
file.write(updated_content)
|
|
|
|
print(".vmx file updated successfully.")
|
|
|
|
vmx_file_base_name = os.path.splitext(vmx_file)[0]
|
|
|
|
assert vmx_file == "Ubuntu.vmx", "The VMX file should be named 'Ubuntu.vmx'."
|
|
files_to_rename = ['vmx', 'nvram', 'vmsd', 'vmxf']
|
|
|
|
for ext in files_to_rename:
|
|
original_file = os.path.join(dir_path, f"{vmx_file_base_name}.{ext}")
|
|
target_file = os.path.join(dir_path, f"{target_vm_name}.{ext}")
|
|
os.rename(original_file, target_file)
|
|
|
|
# Update the dir_path to the target vm_name, only replace the last character
|
|
# Split the path into parts up to the last folder
|
|
path_parts = dir_path.rstrip(os.sep).split(os.sep)
|
|
path_parts[-1] = target_vm_name
|
|
target_dir_path = os.sep.join(path_parts)
|
|
os.rename(dir_path, target_dir_path)
|
|
|
|
print("VM files renamed successfully.")
|
|
|
|
|
|
def _install_virtual_machine(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu"):
|
|
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
|
|
# if os.uname().machine == 'arm64': # Apple Silicon
|
|
url = UBUNTU_ARM_URL
|
|
# else:
|
|
# url = UBUNTU_X86_URL
|
|
elif platform.machine().lower() in ['amd64', 'x86_64']:
|
|
url = UBUNTU_X86_URL
|
|
else:
|
|
raise Exception("Unsupported platform or architecture")
|
|
|
|
# Download the virtual machine image
|
|
print("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
|
|
print("Fully downloaded or the file sized 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:
|
|
print(f"Download error: {e}")
|
|
sleep(1) # Wait for 1 second before retrying
|
|
print("Retrying...")
|
|
else:
|
|
print("Download succeeds.")
|
|
break # Download completed successfully
|
|
|
|
# Unzip the downloaded file
|
|
print("Unzipping the downloaded file...☕️")
|
|
with zipfile.ZipFile(downloaded_file_path, 'r') as zip_ref:
|
|
zip_ref.extractall(os.path.join(vms_dir, vm_name))
|
|
print("Files have been successfully extracted to the directory:", os.path.join(vms_dir, vm_name))
|
|
|
|
vm_path = os.path.join(vms_dir, vm_name, vm_name, vm_name + ".vmx")
|
|
|
|
# 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()
|
|
_update_vm(os.path.join(vms_dir, vm_name, original_vm_name, original_vm_name + ".vmx"), vm_name)
|
|
else:
|
|
print(f"Virtual machine exists: {vm_path}")
|
|
|
|
# Determine the platform of the host machine and decide the parameter for vmrun
|
|
def get_vmrun_type():
|
|
if platform.system() == 'Windows' or platform.system() == 'Linux':
|
|
return '-T ws'
|
|
elif platform.system() == 'Darwin': # Darwin is the system name for macOS
|
|
return '-T fusion'
|
|
else:
|
|
raise Exception("Unsupported operating system")
|
|
|
|
# Start the virtual machine
|
|
def start_vm(vm_path, max_retries=20):
|
|
command = f'vmrun {get_vmrun_type()} start "{vm_path}" nogui'
|
|
for attempt in range(max_retries):
|
|
result = subprocess.run(command, shell=True, text=True, capture_output=True)
|
|
if result.returncode == 0:
|
|
print("Virtual machine started.")
|
|
return True
|
|
else:
|
|
if "Error" in result.stderr:
|
|
print(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
|
|
else:
|
|
print(f"Attempt {attempt + 1} failed: {result.stderr}")
|
|
|
|
if attempt == max_retries - 1:
|
|
print("Maximum retry attempts reached, failed to start the virtual machine.")
|
|
return False
|
|
|
|
if not start_vm(vm_path):
|
|
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
|
|
|
|
def get_vm_ip(vm_path, max_retries=20):
|
|
command = f'vmrun {get_vmrun_type()} getGuestIPAddress "{vm_path}" -wait'
|
|
for attempt in range(max_retries):
|
|
result = subprocess.run(command, shell=True, text=True, capture_output=True)
|
|
if result.returncode == 0:
|
|
return result.stdout.strip()
|
|
else:
|
|
if "Error" in result.stderr:
|
|
print(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
|
|
else:
|
|
print(f"Attempt {attempt + 1} failed: {result.stderr}")
|
|
|
|
if attempt == max_retries - 1:
|
|
print("Maximum retry attempts reached, failed to get the IP of virtual machine.")
|
|
return None
|
|
|
|
vm_ip = get_vm_ip(vm_path)
|
|
if not vm_ip:
|
|
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
|
|
|
|
# Function used to check whether the virtual machine is ready
|
|
def download_screenshot(ip):
|
|
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:
|
|
print(f"Error: {e}")
|
|
print(f"Type: {type(e).__name__}")
|
|
print(f"Error detail: {str(e)}")
|
|
sleep(2)
|
|
return False
|
|
|
|
# Try downloading the screenshot until successful
|
|
while not download_screenshot(vm_ip):
|
|
print("Check whether the virtual machine is ready...")
|
|
|
|
print("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")
|
|
|
|
def create_vm_snapshot(vm_path, max_retries=20):
|
|
command = f'vmrun {get_vmrun_type()} snapshot "{vm_path}" "init_state"'
|
|
for attempt in range(max_retries):
|
|
result = subprocess.run(command, shell=True, text=True, capture_output=True)
|
|
if result.returncode == 0:
|
|
print("Snapshot created.")
|
|
return True
|
|
else:
|
|
if "Error" in result.stderr:
|
|
print(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
|
|
else:
|
|
print(f"Attempt {attempt + 1} failed: {result.stderr}")
|
|
|
|
if attempt == max_retries - 1:
|
|
print("Maximum retry attempts reached, failed to create snapshot.")
|
|
return False
|
|
|
|
# Create a snapshot of the virtual machine
|
|
if create_vm_snapshot(vm_path, max_retries=MAX_RETRY_TIMES):
|
|
return vm_path
|
|
else:
|
|
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
|
|
|
|
|
|
def _get_vm_path():
|
|
vm_manager = VirtualMachineManager(REGISTRY_PATH)
|
|
vm_manager.check_and_clean(vms_dir=VMS_DIR)
|
|
free_vms_paths = vm_manager.list_free_vms()
|
|
if len(free_vms_paths) == 0:
|
|
# No free virtual machine available, generate a new one
|
|
print("No free virtual machine available. Generating a new one, which would take a while...☕")
|
|
new_vm_name = vm_manager.generate_new_vm_name(vms_dir=VMS_DIR)
|
|
new_vm_path = _install_virtual_machine(new_vm_name, vms_dir=VMS_DIR, downloaded_file_name=DOWNLOADED_FILE_NAME)
|
|
vm_manager.add_vm(new_vm_path)
|
|
vm_manager.occupy_vm(new_vm_path, os.getpid())
|
|
return new_vm_path
|
|
else:
|
|
# Choose the first free virtual machine
|
|
chosen_vm_path = free_vms_paths[0][0]
|
|
vm_manager.occupy_vm(chosen_vm_path, os.getpid())
|
|
return chosen_vm_path
|