#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK

import argparse
import logging
import re
import subprocess
import tempfile

def main(**kwargs):
    if kwargs.get('detach'):
        virsh_subcommand = 'detach-device'
        virsh_action = 'detach'
    else:
        virsh_subcommand = 'attach-device'
        virsh_action = 'attach'

    for device in kwargs.get('device_ids'):
        bus, dev = device.split(':')
        with tempfile.NamedTemporaryFile('w+') as f:
            f.write(f"""
                <hostdev mode='subsystem' type='usb'>
                    <source>
                        <vendor id='0x{bus}'/>
                        <product id='0x{dev}' />
                    </source>
                </hostdev>
            """)
            f.flush()
            logging.info('%s %s',virsh_action, device)
            subprocess.run(['virsh', virsh_subcommand, kwargs.get('domain'), f.name, '--live'], check=False)

def regex_arg_type(regex):
    pat = re.compile(regex)
    def regex_type(value):
        if pat.match(value):
            return value
        else:
            raise argparse.ArgumentTypeError('argument "%s" must match regex "%s"' % (value, regex))
    return regex_type

def get_libvirt_domains():
    sp = subprocess.run(['virsh', 'list', '--name'], capture_output=True, encoding='utf8', check=True)
    for line in sp.stdout.split('\n'):
        if line:
            yield line.strip()

def get_usb_devices():
    sp = subprocess.run(['lsusb'], capture_output=True, encoding='utf8', check=True)
    for line in sp.stdout.split('\n'):
        if line:
            yield line.split()[5]

if __name__ == '__main__':
    logging.basicConfig(format='%(message)s', level=logging.INFO)
    parser = argparse.ArgumentParser()
    action_group = parser.add_mutually_exclusive_group()
    action_group.add_argument('--attach',
                              action='store_true',
                              help='attach the usb devices (default)')
    action_group.add_argument('--detach',
                              action='store_true',
                              help='detach the usb devices')
    parser.add_argument('-d', '--domain',
                        choices=list(get_libvirt_domains()),
                        default=next(get_libvirt_domains(), None),
                        help='the libvirt domain on which to attach/detach the usb devices (default: %(default)s)')
    parser.add_argument('device_ids',
                        metavar='device_ids',
                        nargs='+',
                        choices=list(get_usb_devices()),
                        help='a list of usb vendor:product ids in hexadecimal to attach/detach to the vm')
    try:
        import argcomplete
        argcomplete.autocomplete(parser)
    except ImportError:
        pass

    args = parser.parse_args()
    main(**args.__dict__)