example

#!/usr/bin/env python3 -u

# SPDX-FileCopyrightText: The vmnet-helper authors
# SPDX-License-Identifier: Apache-2.0

"""
Run a vm with vmnet-helper.
"""

import argparse
import logging
import os
import signal
import sys

import vmnet

VMNET_OFFLOAD_AUTO = {
    "vfkit": "off",
    "krunkit": "off",
    "qemu": "off",
}


def main():
    shared_interfaces = vmnet.shared_interfaces()
    p = argparse.ArgumentParser(description="Run virtual machine using vmnet-helper")

    # Main arguments

    p.add_argument("vm_name", help="VM name")
    p.add_argument(
        "--connection",
        choices=["fd", "socket", "runner"],
        default="fd",
        help="How to connect the helper and vm (fd)",
    )

    # Helper arguments

    p.add_argument(
        "--operation-mode",
        choices=vmnet.OPERATION_MODES,
        help="vmnet operation mode",
    )
    p.add_argument(
        "--start-address",
        help=(
            "The starting IPv4 address (string) to use for the interface. This "
            "address is used as the gateway address. The subsequent address up "
            "to and including --end-address_key are placed in the DHCP pool. All "
            "other addresses are available for static assignment. "
            "The address must be in the private IP range (RFC 1918)."
        ),
    )
    p.add_argument(
        "--end-address",
        help=(
            "The DHCP IPv4 range end address (string) to use for the interface. "
            "The address must be in the private IP range (RFC 1918)."
        ),
    )
    p.add_argument(
        "--subnet-mask", help="The IPv4 subnet mask (string) to use on the interface."
    )
    p.add_argument(
        "--shared-interface",
        choices=shared_interfaces,
        default=shared_interfaces[0],
        help=f"vmnet shared interface for --operation-mode=bridged ({shared_interfaces[0]})",
    )
    p.add_argument(
        "--enable-isolation",
        action="store_true",
        help="Isolate the guest from other guests on the vmnet interface, requires --operation-mode=host",
    )
    p.add_argument(
        "--enable-offloading",
        action="store_true",
        help="Enable vmnet tso and checksum-offload options (krunkit driver only)",
    )
    p.add_argument(
        "--network",
        dest="network_name",
        metavar="NAME",
        help="Join a network managed by vmnet-broker (macOS 26 only)",
    )

    # VM arguments

    p.add_argument(
        "--driver",
        metavar="NAME",
        choices=vmnet.DRIVERS,
        default=vmnet.DRIVERS[0],
        help=f"VM driver ({vmnet.DRIVERS[0]})",
    )
    p.add_argument(
        "--driver-command",
        metavar="COMMAND",
        help=f"Use custom driver command",
    )
    p.add_argument(
        "--cpus",
        metavar="N",
        type=cpus,
        default=2,
        help="Number of vpus (2)",
    )
    p.add_argument(
        "--memory",
        metavar="M",
        type=int,
        default=1024,
        help="Memory size in MiB (1024)",
    )
    p.add_argument(
        "--timeout",
        metavar="SECONDS",
        type=int,
        default=60,
        help="Time to wait for VM IP address (60 seconds)",
    )

    distros = list(vmnet.IMAGES.keys())
    p.add_argument(
        "--distro",
        metavar="NAME",
        choices=distros,
        default=distros[0],
        help=f"Linux distro ({distros[0]})",
    )

    p.add_argument(
        "--dns-servers",
        metavar="ADDR",
        type=ip_list,
        default=["8.8.8.8", "1.1.1.1"],
        help="Comma-separated DNS servers for the VM (8.8.8.8,1.1.1.1)",
    )

    # Debugging

    p.add_argument("-v", "--verbose", action="store_true", help="Be more verbose")

    args = p.parse_args()
    setup_logging(args.verbose)

    if args.network_name:
        if args.operation_mode:
            p.error("--network cannot be used with --operation-mode")
        if args.start_address:
            p.error("--network cannot be used with --start-address")
        if args.end_address:
            p.error("--network cannot be used with --end-address")
        if args.subnet_mask:
            p.error("--network cannot be used with --subnet-mask")

    if args.operation_mode == "bridged":
        if not args.shared_interface:
            p.error("--shared-interface required for --operation-mode=bridged")
        if args.enable_isolation:
            p.error("--enable-isolation not compatible with --operation-mode=bridged")

    signal.signal(signal.SIGTERM, terminate)
    signal.signal(signal.SIGINT, terminate)

    if args.connection == "fd":
        run_with_fd(args)
    elif args.connection == "socket":
        run_with_socket(args)
    elif args.connection == "runner":
        run_with_runner(args)


class SecondsFormatter(logging.Formatter):
    def format(self, record):
        record.seconds = record.relativeCreated / 1000.0
        return super().format(record)


def setup_logging(verbose=False):
    formatter = SecondsFormatter("[%(seconds)8.3f] %(levelname)s %(message)s")
    handler = logging.StreamHandler(stream=sys.stdout)
    handler.setFormatter(formatter)
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(handlers=[handler], level=level)


def run_with_fd(args):
    """
    Use file descriptor passing to connect the vm and the helper.
    """
    # Create a socketpair for the child processes. The vm_sock will be used by
    # the vm, and the host_sock will be used by the helper.
    vm_sock, helper_sock = vmnet.socketpair()

    # Start vmnet-helper first, since we need the MAC address before starting
    # the VM.
    helper = vmnet.Helper(args, fd=helper_sock.fileno())
    helper.start()
    try:
        # Start the VM with the MAC address and the second socket.
        vm = vmnet.VM(
            args, helper.interface[vmnet.VMNET_MAC_ADDRESS], fd=vm_sock.fileno()
        )
        vm.start()
        try:
            vm.wait_until_ready(timeout=args.timeout)
            os.wait()
        finally:
            vm.stop()
    finally:
        helper.stop()


def run_with_socket(args):
    """
    Use unix socket to connect the vm and the helper.
    """
    # The helper will create a unix datagram socket at this path and wait for
    # client connection.
    socket = vmnet.vm_path(args.vm_name, "vmnet.sock")

    # Start vmnet-helper first, since we need the MAC address before starting
    # the VM.
    helper = vmnet.Helper(args, socket=socket)
    helper.start()
    try:
        # Start the VM with the MAC address and the socket path.
        vm = vmnet.VM(args, helper.interface[vmnet.VMNET_MAC_ADDRESS], socket=socket)
        vm.start()
        try:
            vm.wait_until_ready(timeout=args.timeout)
            os.wait()
        finally:
            vm.stop()
    finally:
        helper.stop()


def run_with_runner(args):
    """
    Use vmnet-run to run a vm with fd.
    """
    vm = vmnet.VM(args, vmnet.mac_address_from(args.vm_name), fd=4, runner=vmnet.RUNNER)
    vm.start()
    try:
        vm.wait_until_ready(timeout=args.timeout)
        os.wait()
    finally:
        vm.stop()


def ip_list(s):
    # TODO: validate that values are IP addresses.
    return s.split(",")


def cpus(s):
    n = int(s)
    if n < 1:
        raise ValueError(f"Invalid number of cpus: '{s}'")
    return n


def terminate(signo, frame):
    sys.exit(1)


if __name__ == "__main__":
    main()