Support more distros in the example

example

@@ -27,10 +27,19 @@ import uuid
 
 import yaml
 
+
 HOME = os.path.expanduser("~/.vmnet-helper")
 IMAGES = {
-    "arm64": "https://cloud-images.ubuntu.com/releases/24.10/release/ubuntu-24.10-server-cloudimg-arm64.img",
-    "x86_64": "https://cloud-images.ubuntu.com/releases/24.10/release/ubuntu-24.10-server-cloudimg-amd64.img",
+    "ubuntu": {
+        "arm64": "https://cloud-images.ubuntu.com/releases/24.10/release/ubuntu-24.10-server-cloudimg-arm64.img",
+        "x86_64": "https://cloud-images.ubuntu.com/releases/24.10/release/ubuntu-24.10-server-cloudimg-amd64.img",
+    },
+    "alpine": {
+        "arm64": "https://dl-cdn.alpinelinux.org/alpine/v3.21/releases/cloud/nocloud_alpine-3.21.2-aarch64-uefi-cloudinit-r0.qcow2",
+    },
+    "fedora": {
+        "arm64": "https://download.fedoraproject.org/pub/fedora/linux/releases/41/Cloud/aarch64/images/Fedora-Cloud-Base-Generic-41-1.4.aarch64.qcow2",
+    },
 }
 
 # Apple recommends that the receive buffer is 4 times bigger than the send
@@ -60,6 +69,12 @@ def main():
     p.add_argument(
         "--bridged", metavar="INTERFACE", help="Create brigded network using interface"
     )
+    p.add_argument(
+        "--distro",
+        choices=list(IMAGES.keys()),
+        default="ubuntu",
+        help="Linux distro (ubuntu)",
+    )
     p.add_argument("-v", "--verbose", action="store_true", help="Be more verbose")
     args = p.parse_args()
 
@@ -75,30 +90,26 @@ def main():
 
     # Start vmnet-helper first, since we need the MAC address for starting the VM.
     helper, mac_address = start_helper(
-        args.vm_name,
+        args,
         helper_sock.fileno(),
-        bridged=args.bridged,
-        verbose=args.verbose,
         log=helper_log,
     )
     try:
         # Start the VM with the second socket and the MAC address.
         serial = vm_path(args.vm_name, "serial.log")
         vm = start_vm(
-            args.vm_name,
+            args,
             vm_sock.fileno(),
             mac_address,
             vm_log,
             serial,
-            driver=args.driver,
-            cpus=args.cpus,
         )
 
-        lookup_ip_address(args.vm_name, serial)
+        lookup_ip_address(args, serial)
 
         # Wait until the VM terminate. We can also exit and use pid files to
         # terminate the processes.
-        print(f"Waiting until virtual machine stops")
+        print("Waiting until virtual machine stops")
         vm.wait()
     except KeyboardInterrupt:
         print("Terminating")
@@ -116,8 +127,8 @@ def cpus(s):
     return n
 
 
-def create_image():
-    image_url = IMAGES[platform.machine()]
+def create_image(args):
+    image_url = IMAGES[args.distro][platform.machine()]
     image_hash = hashlib.sha256(image_url.encode()).hexdigest()
     path = cache_path("images", image_hash, "disk.img")
     if not os.path.exists(path):
@@ -178,12 +189,14 @@ def create_socketpair():
     return pair
 
 
-def start_helper(vm_name, helper_fd, bridged=None, verbose=False, log=None):
+def start_helper(args, helper_fd, log=None):
     """
     Starts vmnet-helper with helper_fd.
     """
-    interface_id = interface_id_from(vm_name)
-    print(f"Starting vmnet-helper for '{vm_name}' with interface id '{interface_id}'")
+    interface_id = interface_id_from(args.vm_name)
+    print(
+        f"Starting vmnet-helper for '{args.vm_name}' with interface id '{interface_id}'"
+    )
     cmd = [
         "sudo",
         "--non-interactive",
@@ -192,10 +205,10 @@ def start_helper(vm_name, helper_fd, bridged=None, verbose=False, log=None):
         f"--fd={helper_fd}",
         f"--interface-id={interface_id}",
     ]
-    if bridged:
+    if args.bridged:
         cmd.append("--operation-mode=bridged")
-        cmd.append(f"--shared-interface={bridged}")
-    if verbose:
+        cmd.append(f"--shared-interface={args.bridged}")
+    if args.verbose:
         cmd.append("--verbose")
 
     helper = subprocess.Popen(
@@ -228,31 +241,33 @@ def interface_id_from(name):
     return str(uuid.UUID(bytes=md[:16], version=4))
 
 
-def start_vm(vm_name, vm_fd, mac_address, log, serial, driver="vfkit", cpus=1):
+def start_vm(args, vm_fd, mac_address, log, serial):
     """
     Starts a vfkit VM using the vm_fd and mac_address.
     """
-    image = create_image()
-    disk = create_disk(vm_name, image)
-    cidata = create_cidata(vm_name, mac_address)
-    if driver == "vfkit":
-        cmd = vfkit_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial)
-    elif driver == "qemu":
-        cmd = qemu_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial)
+    image = create_image(args)
+    disk = create_disk(args, image)
+    cidata = create_cidata(args, mac_address)
+    if args.driver == "vfkit":
+        cmd = vfkit_command(args, vm_fd, mac_address, image, disk, cidata, serial)
+    elif args.driver == "qemu":
+        cmd = qemu_command(args, vm_fd, mac_address, image, disk, cidata, serial)
     else:
-        raise ValueError(f"Invalid driver '{driver}'")
-    print(f"Starting '{driver}' virtual machine '{vm_name}' with mac address '{mac_address}'")
+        raise ValueError(f"Invalid driver '{args.driver}'")
+    print(
+        f"Starting '{args.driver}' virtual machine '{args.vm_name}' with mac address '{mac_address}'"
+    )
     silent_remove(serial)
     vm = subprocess.Popen(cmd, stderr=log, pass_fds=[vm_fd])
     return vm
 
 
-def vfkit_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial):
-    efi_store = vm_path(vm_name, "efi-variable-store")
+def vfkit_command(args, vm_fd, mac_address, image, disk, cidata, serial):
+    efi_store = vm_path(args.vm_name, "efi-variable-store")
     return [
         "vfkit",
         "--memory=2048",
-        f"--cpus={cpus}",
+        f"--cpus={args.cpus}",
         f"--bootloader=efi,variable-store={efi_store},create",
         f"--device=usb-mass-storage,path={cidata},readonly",
         f"--device=virtio-blk,path={disk}",
@@ -262,19 +277,19 @@ def vfkit_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial
     ]
 
 
-def qemu_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial):
+def qemu_command(args, vm_fd, mac_address, image, disk, cidata, serial):
     return [
         "qemu-system-aarch64",
         "-name",
-        vm_name,
+        args.vm_name,
         "-m",
         "2048",
         "-cpu",
         "host",
         "-machine",
         "virt,accel=hvf",
         "-smp",
-        f"{cpus},sockets=1,cores={cpus},threads=1",
+        f"{args.cpus},sockets=1,cores={args.cpus},threads=1",
         "-drive",
         "if=pflash,format=raw,readonly=on,file=/opt/homebrew/share/qemu/edk2-aarch64-code.fd",
         "-drive",
@@ -297,56 +312,56 @@ def qemu_command(vm_name, vm_fd, mac_address, cpus, image, disk, cidata, serial)
     ]
 
 
-def lookup_ip_address(vm_name, serial, timeout=10):
+def lookup_ip_address(args, serial):
     """
     Find the ip address in the serial log and print it.
     """
-    address_prefix = f"{vm_name} address: "
-    login_prefix = f"{vm_name} login: "
+    address_prefix = f"{args.vm_name} address: "
+    login_prefix = f"{args.vm_name} login: "
 
     p = subprocess.Popen(
         ["tail", "-F", serial],
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
     )
     try:
-        for line in iter(p.stdout.readline, ''):
+        for line in iter(p.stdout.readline, ""):
             line = line.strip().decode()
             if line.startswith(address_prefix):
                 _, ip_address = line.split(":", 1)
                 print(f"Virtual machine IP address: {ip_address}")
                 break
             if line.startswith(login_prefix):
-                print(f"Failed to look up virtual machine IP address")
+                print("Failed to look up virtual machine IP address")
                 break
     finally:
         p.kill()
         p.wait()
 
 
-def create_disk(vm_name, image):
+def create_disk(args, image):
     """
     Create a disk from image using copy-on-write.
     """
-    disk = vm_path(vm_name, "disk.img")
+    disk = vm_path(args.vm_name, "disk.img")
     if not os.path.isfile(disk):
         print(f"Creating disk '{disk}'")
         subprocess.run(["cp", "-c", image, disk], check=True)
     return disk
 
 
-def create_cidata(vm_name, mac_address):
+def create_cidata(args, mac_address):
     """
     Create cloud-init iso image.
 
     We always create cidata.iso to run users scripts and update configuration
     on every start.
     """
-    vm_home = vm_path(vm_name)
+    vm_home = vm_path(args.vm_name)
     cidata = os.path.join(vm_home, "cidata.iso")
-    create_user_data(vm_name)
-    create_meta_data(vm_name)
-    create_network_config(vm_name, mac_address)
+    create_user_data(args)
+    create_meta_data(args)
+    create_network_config(args, mac_address)
     cmd = [
         "mkisofs",
         "-output",
@@ -370,11 +385,11 @@ def create_cidata(vm_name, mac_address):
     return cidata
 
 
-def create_user_data(vm_name):
+def create_user_data(args):
     """
     Create cloud-init user-data file.
     """
-    path = vm_path(vm_name, "user-data")
+    path = vm_path(args.vm_name, "user-data")
     data = {
         "password": "ubuntu",
         "chpasswd": {
@@ -383,32 +398,32 @@ def create_user_data(vm_name):
         "ssh_authorized_keys": public_keys(),
         "runcmd": [
             "ip_address=$(ip -4 -j addr show dev vmnet0 | jq -r '.[0].addr_info[0].local')",
-            f"sudo echo {vm_name} address: $ip_address > /etc/issue",
-        ]
+            f"sudo echo {args.vm_name} address: $ip_address > /etc/issue",
+        ],
     }
     with open(path, "w") as f:
         f.write("#cloud-config\n")
         yaml.dump(data, f)
 
 
-def create_meta_data(vm_name):
+def create_meta_data(args):
     """
     Create cloud-init meta-data file.
     """
-    path = vm_path(vm_name, "meta-data")
+    path = vm_path(args.vm_name, "meta-data")
     data = {
         "instance-id": str(uuid.uuid4()),
-        "local-hostname": vm_name,
+        "local-hostname": args.vm_name,
     }
     with open(path, "w") as f:
         yaml.dump(data, f)
 
 
-def create_network_config(vm_name, mac_address):
+def create_network_config(args, mac_address):
     """
     Create cloud-init network-config file.
     """
-    path = vm_path(vm_name, "network-config")
+    path = vm_path(args.vm_name, "network-config")
     data = f"""\
 version: 2
 ethernets: