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 os
import signal
import sys

import vmnet

CONNECTION_AUTO = {
    "vfkit": "fd",
    "krunkit": "socket",
    "qemu": "fd",
}

VMNET_OFFLOAD_AUTO = {
    "vfkit": "off",
    # libkrun enables offloading feautres by default, so we must enable vmnet
    # offloading.  See https://github.com/containers/libkrun/issues/264
    "krunkit": "on",
    "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=["auto", "fd", "socket"],
        default="auto",
        help="How to connect the helper and vm (auto)",
    )

    # Helper arguments

    p.add_argument(
        "--operation-mode",
        choices=vmnet.OPERATION_MODES,
        default=vmnet.OPERATION_MODES[0],
        help=f"vmnet operation mode ({vmnet.OPERATION_MODES[0]})",
    )
    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(
        "--vmnet-offload",
        choices=["auto", "on", "off"],
        default="auto",
        help="Enable vmnet tso and checksum-offload options (auto)",
    )

    # 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(
        "--krunkit-port",
        metavar="PORT",
        type=int,
        default=8081,
        help="krunkit restful-uri port (8081)",
    )
    p.add_argument(
        "--cpus",
        metavar="N",
        type=cpus,
        default=1,
        help="Number of vpus (1)",
    )
    p.add_argument(
        "--memory",
        metavar="M",
        type=int,
        default=1024,
        help="Memory size in MiB (1024)",
    )

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

    # Debugging

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

    args = p.parse_args()

    if args.connection == "auto":
        args.connection = CONNECTION_AUTO[args.driver]

    if args.operation_mode == "bridged" and not args.shared_interface:
        p.error("--shared-interface required for --operation-mode=bridged")

    if args.enable_isolation and args.operation_mode != "host":
        p.error("--enable-isolation requires --operation-mode=host")

    if args.vmnet_offload == "auto":
        args.vmnet_offload = VMNET_OFFLOAD_AUTO[args.driver]

    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)


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_mac_address"], fd=vm_sock.fileno())
        vm.start()
        try:
            vm.wait_for_ip_address()
            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_mac_address"], socket=socket)
        vm.start()
        try:
            vm.wait_for_ip_address()
            os.wait()
        finally:
            vm.stop()
    finally:
        helper.stop()


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()