usb_boot

#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2021-2022 Eli Array Minkoff
#
# SPDX-License-Identifier: 0BSD

"""
Use QEMU to run a bootable USB thumb drive
"""
import argparse
import subprocess


def lsusb():
    """Call the lsusb system command, and return its output as list of strings
    lsusb outputs a list of USB devices with the following format:
        Bus 001 Device 001: ID 1d6b:0002 Linux Foundation 2.0 root hub
        Bus 001 Device 002: ID 239a:801e Adafruit Trinket M0
        Bus 001 Device 005: ID 057e:2006 Nintendo Co., Ltd Joy-Con L
        Bus 001 Device 006: ID 057e:2007 Nintendo Co., Ltd Joy-Con R
        Bus 001 Device 007: ID 05ac:0202 Apple, Inc. Keyboard [ALPS]
    A list of device ids can be found at http://www.linux-usb.org/usb.ids
    """
    return subprocess.check_output(["lsusb"]).decode().splitlines()


def parse_lsusb_entry(lsusb_entry):
    """Parse a line outputed by lsusb, and return a dict with its attributes

    Return
    (dict) {
        'bus': (int) the numeric id of the bus the device is connected to
        'addr': (int) the numeric id of the device's address on the bus
        'id': (str) the identifier, manufacturer, and description of the device
    }
    """
    #
    bus_dev_str, id = [i.strip() for i in lsusb_entry.split(":", 1)]
    # bus_dev_str should look something like 'Bus MMM Device NNN'.
    bus, addr = [int(i) for i in bus_dev_str.split(" ")[1::2]]
    return {"bus": bus, "addr": addr, "id": id}


def get_args():
    """Parse the command-line flags, return a Namespace object."""
    parser = argparse.ArgumentParser(
        description="Test bootable USB using QEMU",
        allow_abbrev=False,
    )
    # add basic arguments
    parser.add_argument(
        "--qemucmd",
        "-qemucmd",
        "-q",
        help="The QEMU binary to use",
    )
    parser.add_argument(
        "--efi-path",
        "-efi-path",
        "-O",
        help="Path to the EFI binary. Use in conjunction with --efi.",
        default="/usr/share/ovmf/OVMF.fd",
    )
    parser.add_argument(
        "--memory",
        "-memory",
        "-m",
        help="The amount of memory to allocate",
    )
    # add mutually exclusive --kvm and --nokvm flags and their variants
    kvm_status = parser.add_mutually_exclusive_group()
    kvm_status.add_argument(
        "--kvm", "-kvm", "-k", action="store_true", dest="use_kvm"
    )
    kvm_status.add_argument(
        "--nokvm", "-nokvm", "-K", action="store_false", dest="use_kvm"
    )
    parser.set_defaults(use_kvm=True)

    # add mutually exclusive --efi and --noefi flags and their variants
    use_efi = parser.add_mutually_exclusive_group()
    use_efi.add_argument(
        "--efi", "-efi", "-e", action="store_true", dest="use_efi"
    )
    use_efi.add_argument(
        "--noefi", "-noefi", "-E", action="store_false", dest="use_efi"
    )
    parser.set_defaults(use_efi=False)

    # parse the arguments using the parser object created
    return parser.parse_args()


def main():
    print("usb_boot Copyright (C) 2021, 2022 Eli Array Minkoff.")
    print(
        "This program comes with ABSOLUTELY NO WARRANTY;",
        "for details, see the LICENCE file.",
    )
    print(
        "This is free software,",
        "and you are welcome to redistribute it under certain conditions;",
        "for details, see the LICENCE file.\n\n",
    )
    args = get_args()
    usb_devices = list(map(parse_lsusb_entry, lsusb()))
    # print out the devices in a numbered list, and ask for selection
    for i in range(len(usb_devices)):
        print(f'{i} | {usb_devices[i]["id"]}')
    selection = -1
    while selection < 0:
        selection = int(input("Enter the number of your selection: "))
        cmd = [
            "sudo",
            args.qemucmd,
            "-m",
            args.memory,
            "-usb",
            "-device",
            "usb-host,hostbus={},hostaddr={}".format(
                usb_devices[selection]["bus"],
                usb_devices[selection]["addr"],
            ),
        ]
        if args.use_kvm:
            cmd.append("-enable-kvm")
        if args.use_efi:
            cmd += ["-smbios", "type=0,uefi=on", "-bios", args.efi_path]
        subprocess.run(cmd)


if __name__ == "__main__":
    main()