U-Boot

# SPDX-License-Identifier:      GPL-2.0+
#
# Copyright 2025 Simon Glass <sjg@chromium.org>
# Author: Simon Glass <sjg@chromium.org>

"""Create test disk-images for VBE"""

import gzip
import os
import tempfile

import utils
from fs_helper import DiskHelper, FsHelper
from u_boot_pylib import tools


def dtb_for_compatible(log, kernpath, version, model):
    """Create a fake devicetree binary for a given model name

    Args:
        log (multiplexed_log.Logfile): Log to write to
        kernpath (str): Path to place the created DTB
        version (str): Version string to use with the dtb
        model (str): Model name

    Return:
        str: Filename of devicetree file written
    """
    dtb_file = os.path.join(kernpath, f'dtb-{version}-{model}')
    data = f'/dts-v1/; / {{ compatible = "{model}"; version = "{version}"; }};'
    utils.run_and_log_no_ubman(log, f'dtc -o {dtb_file}',
                               stdin=data.encode('utf-8'))
    return dtb_file


def create_extlinux(config, log, kernpath, dirpath, slot, version, has_oem):
    """Create a fake extlinux image

    Args:
        config (ArbitraryAttributeContainer): Configuration
        log (multiplexed_log.Logfile): Log to write to
        kernpath (str): Directory path for the kernel
        dirpath (str): Directory path to write to (created by this function)
        slot (str): Slot name (A, B or recovery)
        version (str): Linux version to use, e.g. '6.14.0-24'
        has_oem (bool): true to create an OEM FIT
    """
    ext_path = os.path.join(dirpath, 'extlinux')
    if not os.path.exists(dirpath):
        os.mkdir(dirpath)
    if not os.path.exists(ext_path):
        os.mkdir(ext_path)
    dtb_info = ' (no-dtb)' if has_oem else ''
    with open(os.path.join(ext_path, 'extlinux.conf'), 'w',
              encoding='ascii') as outf:
        print(f'''## /boot/extlinux/extlinux.conf
##
## IMPORTANT WARNING
##
## The configuration of this file is generated automatically.
## Do not edit this file manually, use: u-boot-update

default l0
menu title Ubuntu Menu ({slot})
prompt 1
timeout 50


label l0
	menu label Ubuntu 24.04.2 LTS {version}{dtb_info}

	fit /ubuntu-{version}.fit
	append root=/dev/disk/by-id/dm-uuid-LVM-lNvIY7Drx7RGUJpQZGPr2axs2ueSe57B7p7jtyJGIUnMJMHCSPVphCu9IFivkJhK ro earlycon
''', file=outf)

    # True to add devicetrees and multiple configurations
    add_dts = not has_oem

    # Create a FIT containing kernel, initrd and devicetree
    with tempfile.TemporaryDirectory(suffix='vbe') as tmp:
        kern = os.path.join(tmp, 'kern')
        tools.write_file(kern, gzip.compress(f'vmlinux-{slot}'.encode('utf-8')))
        mkimage = config.build_dir + '/tools/mkimage'

        initrd = os.path.join(tmp, f'initrd.img-{version}')
        tools.write_file(initrd, f'initrd {version}', binary=False)

        snow = dtb_for_compatible(log, tmp, version, 'snow')
        kevin = dtb_for_compatible(log, tmp, version, 'kevin')

        fit = os.path.join(kernpath, f'ubuntu-{version}.fit')
        cmd = f'{mkimage} -f auto -T kernel -A sandbox -O linux '
        dts = f'-b {snow} -b {kevin}' if add_dts else ''
        utils.run_and_log_no_ubman(log,
                                   cmd + f'-d {kern} -i {initrd} {dts} {fit}')

        # Create a FIT containing OEM devicetrees
        if has_oem:
            oem_snow = dtb_for_compatible(log, tmp, 'oem', 'snow')
            oem_kevin = dtb_for_compatible(log, tmp, 'oem', 'kevin')

            fit = os.path.join(dirpath, 'oem.fit')
            utils.run_and_log_no_ubman(
                log,
                f'{mkimage} -f auto  -A sandbox -O linux -T flat_dt --load-only'
                f' -b {oem_snow} -b {oem_kevin} {fit}')

def setup_vbe_image(config, log):
    """Create three partitions (fat, boot and root) for a VBE boot flow

    Args:
        config (ArbitraryAttributeContainer): Configuration
        log (multiplexed_log.Logfile): Log to write to

    The intent is to load either one or two FITs.

    Two separate images are created:

      - vbe0 has no OEM FIT; OS FIT contains the devicetrees
      - vbe1 has an OEM FIT containing the devicetrees; OS FIT does not
    """
    for seq in range(2):
        with (DiskHelper(config, seq, 'vbe') as img,
            FsHelper(config, 'fat', 1, 'efi') as vfat,
            FsHelper(config, 'ext4', 1, f'boot{seq}') as boot,
            FsHelper(config, 'ext4', 17, 'root') as root):
            with open(os.path.join(vfat.srcdir, 'u-boot.efi'), 'wb') as outf:
                outf.write(b'fake binary\n')
            vfat.mk_fs()
            img.add_fs(vfat, DiskHelper.VFAT, bootable=True)

            has_oem = bool(seq)

            dir_a = os.path.join(boot.srcdir, 'a')
            vera = '6.14.0-24-generic'
            create_extlinux(config, log, boot.srcdir, dir_a, 'A', vera, has_oem)

            dir_b = os.path.join(boot.srcdir, 'b')
            verb = '6.8.0-1028-intel'
            create_extlinux(config, log, boot.srcdir, dir_b, 'B', verb, has_oem)

            boot_slot = 'b' if seq else 'a'
            fname = os.path.join(config.result_dir, 'vbe-state.dts')
            tools.write_file(fname, b'''// VBE state file

/dts-v1/;

/ {
    compatible = "vbe,abrec-state";

    os {
        compatible = "vbe,os-state";

        next-boot {
            slot = "boot_slot";
        };
    };
};
'''.replace(b'boot_slot', boot_slot.encode('utf-8')))

            utils.run_and_log_no_ubman(
                log, ['dtc', fname, '-O', 'dtb', '-o',
                      f'{boot.srcdir}/vbe-state'])

            boot.mk_fs()
            img.add_fs(boot, DiskHelper.EXT4)

            root.mk_fs()
            img.add_fs(root, DiskHelper.EXT4)

            img.create()
            img.remove_img = False