#!/usr/bin/env python3

"""
This small script changes the boot order of a virsh guest domain.
It assumes that a guest can be booted from either an HDD or a CD, defines
the right boot order according to the arguments, and reboots the guest.
"""

import argparse
import subprocess
import sys
import tempfile
import time
from xml.dom import minidom


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--vm', type=str, help='Name or ID of a guest domain (VM)', required=True)
    order_group = parser.add_mutually_exclusive_group(required=True)
    order_group.add_argument('--hd', action='store_true', help='Boot from HDD first', default=False)
    order_group.add_argument('--cd', action='store_true', help='Boot from CD first', default=False)
    return parser.parse_args()


def virsh(*varargs):
    cmd = list(varargs)
    cmd.insert(0, 'virsh')
    return subprocess.check_output(cmd, timeout=20).decode('utf-8').strip()


def get_name(vm_id):
    return virsh('domname', vm_id)


def shutdown(vm_name):
    virsh('shutdown', vm_name)

    for _ in range(0, 24):
        if virsh('domstate', vm_name) == 'shut off':
            return
        time.sleep(5)

    raise TimeoutError(f'Shutting down \'{vm_name}\' took too long')


def update_boot_order(definition, cd_first=False):

    dom = minidom.parseString(definition)
    os_element = dom.getElementsByTagName('os')[0]

    for child in os_element.getElementsByTagName('boot'):
        dev = child.getAttribute('dev')
        if dev in ['cdrom', 'hd']:
            os_element.removeChild(child)
        else:
            raise ValueError(f'Found unexpected boot device: \'{dev}\'')

    first = dom.createElement('boot')
    first.setAttribute('dev', 'cdrom' if cd_first else 'hd')
    os_element.appendChild(first)

    second = dom.createElement('boot')
    second.setAttribute('dev', 'hd' if cd_first else 'cdrom')
    os_element.appendChild(second)

    return dom.toprettyxml()


def read_definition(vm_name):
    return virsh('dumpxml', '--inactive', '--security-info', vm_name)


def define(definition):
    with tempfile.NamedTemporaryFile() as f:
        f.write(definition.encode())
        f.seek(0)
        virsh('define', f.name)


def boot(vm_name):
    return virsh('start', vm_name)


if __name__ == '__main__':
    args = get_args()
    try:
        vm = get_name(args.vm) if args.vm.isdigit() else args.vm
        shutdown(vm)
        old_def = read_definition(vm)
        new_def = update_boot_order(old_def, args.cd)
        define(new_def)
        boot(vm)
        print('Done')
    except Exception as e:
        print(e, file=sys.stderr)
        exit(1)
