Qubes utility to reboot (halt, then start) a target VM

qvm-reboot

#!/usr/bin/env python3
"""
Utility script to reboot Qubes domains. Attempts
to perform a graceful shutdown, kills if shutdown fails,
then starts up. Inspiration for the timeout logic taken
from qubesadmin.tools.qvm_shutdown.main.
"""
import argparse
import time
from functools import partial
import concurrent.futures
import qubesadmin


q = qubesadmin.Qubes()


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "vms", nargs="+", type=str, action="store", help="VMs to reboot"
    )
    parser.add_argument(
        "--only-if-outdated",
        default=False,
        action="store_true",
        help=("Whether to skip reboot if VM is up to date" "(all updates applied)"),
    )
    parser.add_argument(
        "--allow-kill",
        default=True,
        action="store_true",
        help="Whether to kill VM if shutdown fails",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=30,
        action="store",
        help="How long to wait for shutdown before killing",
    )
    args = parser.parse_args()
    return args


def _vm_needs_reboot(vm):
    """
    Determine whether VM should be rebooted in order to apply updates.
    Mostly relevant for an AppVM, to check whether its TemplateVM has
    been updated.

    Adapted from:
    https://github.com/QubesOS/qubes-manager/blob/da2826db20fa852403240a45b3906a6c54b2fe33/qubesmanager/table_widgets.py#L402-L406
    """
    is_outdated = False
    for vol in vm.volumes.values():
        if vol.is_outdated():
            is_outdated = True
    return is_outdated


def reboot_vm(vm, timeout=60, allow_kill=False, only_if_outdated=False):
    """
    Attempts to halt gracefully, then restart, a domain.
    If timeout is reached without confirmed shutdown,
    domain will be killed, then booted.
    """
    if only_if_outdated and not _vm_needs_reboot(vm):
        return

    if vm.is_running():
        vm.shutdown()

    waited_so_far = 0
    while waited_so_far < timeout:
        if vm.get_power_state() != "Halted":
            time.sleep(1)
            waited_so_far += 1
        else:
            break

    if vm.get_power_state() != "Halted":
        vm.kill()

    vm.start()


if __name__ == "__main__":
    args = parse_args()
    vms = [q.domains[x] for x in args.vms]
    n_proc = len(vms)
    # use functools.partial to pass kwargs to map()
    # using a manager class would suffice too
    f_args = dict(
        timeout=args.timeout,
        only_if_outdated=args.only_if_outdated,
        allow_kill=args.allow_kill,
    )
    with concurrent.futures.ThreadPoolExecutor(max_workers=n_proc) as executor:
        executor.map(partial(reboot_vm, **f_args), vms)