U-Boot

#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0+
#
"""Script to build/run U-Boot with QEMU

It assumes that

- you build U-Boot in ${ubdir}/<name> where <name> is the U-Boot
  board config
- your OS images are in ${imagedir}/{distroname}/

So far the script supports ARM, RISC-V and x86
"""

import argparse
import os
from pathlib import Path
import subprocess
import sys
import shlex
import time

import build_helper

OUR_PATH = os.path.dirname(os.path.realpath(__file__))
OUR1_PATH = os.path.dirname(OUR_PATH)

# Bring in the patman and dtoc libraries (but don't override the first path
# in PYTHONPATH)
sys.path.insert(2, os.path.join(OUR1_PATH, 'tools'))

# pylint: disable=C0413
from u_boot_pylib import command
from u_boot_pylib import tools
from u_boot_pylib import tout


def parse_args():
    """Parses command-line arguments"""
    parser = argparse.ArgumentParser(
        description='Build and/or run U-Boot with QEMU',
        formatter_class=argparse.RawTextHelpFormatter)
    build_helper.add_common_args(parser)
    parser.add_argument('-e', '--sct-run', action='store_true',
                        help='Run UEFI Self-Certification Test (SCT)')
    parser.add_argument('-E', '--use-tianocore', action='store_true',
                        help='Run Tianocore (OVMF) instead of U-Boot')
    parser.add_argument('-Q', '--use-qboot', action='store_true',
                        help='Run qboot instead of U-Boot')
    parser.add_argument('-x', '--xpl', action='store_true',
                        help='Use xPL image rather than U-Boot proper')
    parser.add_argument('-X', '--no-spl', action='store_true',
                        help='Use no-SPL build (x86_64 only)')
    parser.add_argument('-T', '--tkey', action='store_true',
                        help='Enable TKey USB passthrough for testing')
    parser.add_argument(
        '--sct-seq',
        help='SCT sequence-file to be written into the SCT image if -e')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Show executed commands')

    return parser.parse_args()


class BuildQemu:
    """Build and/or run U-Boot with QEMU based on command line arguments"""

    def __init__(self, args):
        """Set up arguments and configure paths"""
        self.args = args

        self.helper = build_helper.Helper(args)
        self.helper.read_settings()
        self.ubdir = Path(self.helper.get_setting('build_dir', '/tmp/b'))
        self.sctdir = Path(self.helper.get_setting('sct_dir', '~/dev/efi/sct'))
        self.tiano = Path(self.helper.get_setting('tianocore_dir',
                                                  '~/dev/tiano'))
        self.qboot = Path(self.helper.get_setting('qboot_dir', '~/dev/qboot'))
        self.mnt = Path(self.helper.get_setting('sct_mnt', '/mnt/sct'))

        self.qemu_extra = []
        self.helper.mem = '512M'  # Default QEMU memory

        if args.disk:
            self.helper.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])

        if args.sct_run:
            self.helper.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])
            # For potential interaction within SCT
            self.qemu_extra.extend(['-device', 'qemu-xhci'])
            self.qemu_extra.extend(['-device', 'usb-kbd'])
            sct_image_path = self.sctdir / 'sct.img'
            if not sct_image_path.exists():
                tout.fatal(f'Error: SCT image {sct_image_path} not found, '
                           'required for -e')
            self.qemu_extra.extend([
                '-drive', f'file={sct_image_path},format=raw,if=none,id=vda',
                '-device', 'virtio-blk-pci,drive=vda,bootindex=1'])
            # Basic networking for SCT, if needed
            self.qemu_extra.extend([
                '-device', 'virtio-net-pci,netdev=net0',
                '-netdev', 'user,id=net0'])
            args.serial_only = True  # SCT implies serial output

        if args.os:
            self.helper.mem = '4G'
            self.qemu_extra.extend(['-smp', '4'])

        self.kvm_params = []
        if args.kvm:
            self.kvm_params = ['-enable-kvm', '-cpu', 'host']

        bios_override = None
        if args.custom:
            bios_override = Path(args.custom)
            if not bios_override.exists():
                tout.fatal(
                    'Error: Custom BIOS specified (-c) but not found at '
                    f'{bios_override}')
        elif args.use_tianocore:
            if args.arch == 'arm':
                bios_override = Path(self.tiano, 'OVMF-pure-efi.aarch64.fd.64m')
            elif args.arch == 'riscv':
                bios_override = Path(self.tiano, 'RISCV_VIRT_CODE.fd')
            else:
                bios_override = Path(self.tiano, 'OVMF-pure-efi.x64.fd')
            if not bios_override.exists():
                tout.fatal(
                    'Error: Tianocore BIOS specified (-E) but not found at '
                    f'{bios_override}')
        elif args.use_qboot:
            bios_override = Path(self.qboot, 'bios.bin')
            if not bios_override.exists():
                tout.fatal(
                    'Error: qboot BIOS specified (-Q) but not found at '
                    f'{bios_override}')

        self.seq_fname = Path(args.sct_seq) if args.sct_seq else None

        # arch-specific setup
        if args.arch == 'arm':
            if args.xpl:
                self.board = 'qemu_arm_spl'
                default_bios = 'image.bin'
            else:
                self.board = 'qemu_arm'
                default_bios = 'u-boot.bin'
            self.helper.qemu = 'qemu-system-arm'
            self.qemu_extra.extend(['-machine', 'virt'])
            if not args.kvm:
                self.qemu_extra.extend(['-accel', 'tcg'])
            if self.helper.bitness == 64:
                if args.xpl:
                    self.board = 'qemu_arm64_spl'
                else:
                    self.board = 'qemu_arm64'
                self.helper.qemu = 'qemu-system-aarch64'
                self.qemu_extra.extend(['-cpu', 'cortex-a57'])
        elif args.arch == 'riscv':
            if args.xpl:
                self.board = 'qemu-riscv64_spl'
                default_bios = 'u-boot.bin'
            else:
                self.board = 'qemu-riscv64'
                default_bios = 'u-boot.bin'
            self.helper.qemu = 'qemu-system-riscv64'
            self.qemu_extra.extend(['-machine', 'virt'])
            if not args.kvm:
                self.qemu_extra.extend(['-accel', 'tcg'])
            if self.helper.bitness == 32:
                if args.xpl:
                    self.board = 'qemu-riscv32_spl'
                else:
                    self.board = 'qemu-riscv32'
                self.helper.qemu = 'qemu-system-riscv32'
        elif args.arch == 'x86':
            self.board = 'qemu-x86'
            default_bios = 'u-boot.rom'
            self.helper.qemu = 'qemu-system-i386'
            self.qemu_extra.extend(['-machine', 'q35'])
            if args.tkey:
                # Pass through TKey USB device to QEMU
                self.qemu_extra.extend(['-device', 'usb-host,vendorid=0x1207,productid=0x8887'])
            if args.no_spl and self.helper.bitness != 64:
                raise ValueError('-X/--no-spl requires 64-bit mode (cannot be used with -w)')
            if self.helper.bitness == 64:
                if args.no_spl:
                    self.board = 'qemu-x86_64_nospl'
                else:
                    self.board = 'qemu-x86_64'
                self.helper.qemu = 'qemu-system-x86_64'
        else:
            raise ValueError(f"Invalid arch '{args.arch}'")

        self.build_dir = self.ubdir / self.board
        self.bios = (bios_override if bios_override
                     else self.build_dir / default_bios)

    @staticmethod
    def execute_command(cmd_list, desc, check=True, **kwargs):
        """Execute a shell command and handle errors

        Args:
            cmd_list (list of str): The command and its arguments as a list
            desc (str): A description of the command being executed
            check (bool): Raise CalledProcessError on non-zero exit code
            kwargs: Additional arguments for subprocess.run

        Return:
            subprocess.CompletedProcess: The result of the subprocess.run call

        Raises:
            SystemExit: If the command is not found or fails and check is True
        """
        tout.info(f"Executing: {desc} -> {shlex.join(cmd_list)}")
        try:
            # Decode stdout/stderr by default if text=True
            if 'text' not in kwargs:
                kwargs['text'] = True
            return subprocess.run(cmd_list, check=check, **kwargs)
        except FileNotFoundError:
            tout.fatal(f"Error: Command '{cmd_list[0]}' not found")
        except subprocess.CalledProcessError as proc:
            tout.error(f'Error {desc}: Command failed with exit code '
                       f'{proc.returncode}')
            if proc.stdout:
                tout.error(f'Stdout:\n{proc.stdout}')
            if proc.stderr:
                tout.error(f'Stderr:\n{proc.stderr}')
            tout.fatal('Failed')

    def build_u_boot(self):
        """Build U-Boot using buildman
        """
        self.build_dir.mkdir(parents=True, exist_ok=True)
        cmd = ['buildman', '-w', '-o', str(self.build_dir), '--board',
               self.board, '-I']
        if self.args.no_pager:
            cmd += ['-a', '~CONSOLE_PAGER']

        self.execute_command(
            cmd,
            f'Building U-Boot for {self.board} in {self.build_dir}')

    def update_sct_sequence(self):
        """Update the SCT image with a specified sequence file

        Requires sudo for loop device setup and mounting
        """
        if not (self.args.sct_run and self.seq_fname and
                self.seq_fname.exists()):
            if (self.args.sct_run and self.seq_fname and
                    not self.seq_fname.exists()):
                tout.warning(f'Warning: SCT sequence file {self.seq_fname}'
                             'not found')
            return

        fname = self.sctdir / 'sct.img'
        if not fname.exists():
            tout.fatal(f'Error: SCT image {fname} not found')

        loopdev = None
        try:
            # Find free loop device and attach
            loopdev = command.output_one_line(
                'sudo', 'losetup', '--show', '-f', '-P', str(fname))
            partition_path_str = f'{loopdev}p1'

            uid, gid = os.getuid(), os.getgid()
            mount_cmd = ['sudo', 'mount', partition_path_str,
                         str(self.mnt), '-o', f'uid={uid},gid={gid},rw']
            mount_cmd.extend(['-t', 'vfat'])

            self.execute_command(mount_cmd,
                                 f'Mounting {partition_path_str} to {self.mnt}')

            target_sct_path = self.mnt / self.seq_fname.name
            self.execute_command(
                ['sudo', 'cp', str(self.seq_fname), str(target_sct_path)],
                f'Copying {self.seq_fname.name} to {self.mnt}'
            )
            tout.info(f"Copied {self.seq_fname} to {target_sct_path}")

        finally:
            if Path(self.mnt).is_mount():
                self.execute_command(['sudo', 'umount', str(self.mnt)],
                                     f'Unmounting {self.mnt}', check=False)
            if loopdev:
                self.execute_command(['sudo', 'losetup', '-d', loopdev],
                                     f'Detaching loop device {loopdev}',
                                     check=False)

    def run_qemu(self):
        """Construct and run the QEMU command"""
        if not self.bios.exists():
            tout.fatal(f"Error: BIOS file '{self.bios}' not found")

        qemu_cmd = [str(self.helper.qemu)]
        if self.bios:
            qemu_cmd.extend(['-bios', str(self.bios)])
        qemu_cmd.extend(self.kvm_params)
        qemu_cmd.extend(['-m', self.helper.mem])

        if not self.args.sct_run and not self.args.use_qboot:
            qemu_cmd.extend(['-netdev', 'user,id=net0,hostfwd=tcp::2222-:22',
                             '-device', 'virtio-net-pci,netdev=net0'])

        # SCT usually runs headlessly
        if self.args.serial_only or self.args.sct_seq:
            qemu_cmd.extend(['-display', 'none'])
        elif self.args.arch in ('arm', 'riscv'):
            qemu_cmd.extend(['-device', 'virtio-gpu-pci'])
            qemu_cmd.extend(['-device', 'qemu-xhci', '-device', 'usb-kbd',
                             '-device', 'usb-tablet', '-device', 'usb-mouse'])
            qemu_cmd.extend(['-display', 'default,show-cursor=on'])
        elif self.args.arch == 'x86':
            qemu_cmd.extend(['-device', 'qemu-xhci'])
            qemu_cmd.extend(['-device', 'usb-kbd', '-device', 'usb-tablet'])
            qemu_cmd.extend(['-display', 'default,show-cursor=on'])
        if not any(item.startswith('-serial') for item in self.qemu_extra):
            qemu_cmd.extend(['-serial', 'mon:stdio'])

        self.helper.add_qemu_args(self.args, qemu_cmd)

        # Add other parameters gathered from options
        qemu_cmd.extend(self.qemu_extra)

        self.helper.setup_share(qemu_cmd)

        self.helper.run(qemu_cmd)

    def start(self):
        """Build and run QEMU"""
        if not self.args.no_build and not self.args.use_tianocore:
            self.build_u_boot()

        # Update SCT sequence if -e and -S are given
        if self.args.sct_run and self.seq_fname:
            self.update_sct_sequence()

        self.run_qemu()


def main():
    """Parse arguments and start the program"""
    args = parse_args()
    tout.init(tout.INFO if args.verbose else tout.WARNING)

    qemu = BuildQemu(args)
    qemu.start()

if __name__ == '__main__':
    main()