#! /usr/bin/env python3

import argparse
import logging
import os
import platform
import struct
import subprocess
import sys
import tempfile

from mmap import mmap
from pathlib import Path

from depthcharge_tools import __version__
from depthcharge_tools.utils.argparse import (
    Command,
    Argument,
    Group,
)
from depthcharge_tools.utils.pathlib import (
    copy,
    decompress,
)
from depthcharge_tools.utils.platform import (
    Architecture,
    vboot_keys,
)
from depthcharge_tools.utils.string import (
    parse_bytesize,
)
from depthcharge_tools.utils.subprocess import (
    mkimage,
    vbutil_kernel,
    lz4,
    lzma,
    fdtget,
    fdtput,
)


class mkdepthcharge(
    Command,
    prog="mkdepthcharge",
    usage="%(prog)s [options] -o FILE [--] [VMLINUZ] [INITRAMFS] [DTB ...]",
    add_help=False,
):
    """Build boot images for the ChromeOS bootloader."""

    _logger = logging.getLogger(__name__)

    @property
    def logger(self):
        # Set verbosity before logging messages
        self.verbosity
        return self._logger

    # Inputs can have the same name and cause collisions in tmpdir.
    def _tempfile(self, name):
        f = self.tmpdir / name
        if not f.exists():
            return f

        for i in range(9999):
            f = f.with_name("{}-idx{:04}".format(name, i))
            if not f.exists():
                return f

        raise FileExistsError(self.tmpdir / name)

    # Debian packs the arm64 kernel uncompressed, but the bindeb-pkg
    # kernel target packs it as gzip. So we'll try to decompress inputs.
    def _decompress(self, f):
        self.logger.info(
            "Trying to decompress file '{}'."
            .format(f)
        )

        decomp = decompress(f, self._tempfile(f.name))
        if decomp is not None:
            self.logger.info(
                "Decompressed input '{}' as '{}'."
                .format(f, decomp)
            )

        return decomp or f

    # Copy inputs to tmpdir because mkimage wants modifiable files.
    def _copy(self, f):
        if self.tmpdir not in f.parents:
            f = copy(f, self._tempfile(f.name))
        f.chmod(0o755)

        return f

    @Group
    def input_files(self):
        """Input files"""

    @input_files.add
    @Argument(help=argparse.SUPPRESS)
    def files(self, *files):
        """Input files (vmlinuz, initramfs, dtbs)"""
        vmlinuz = []
        initramfs = []
        dtbs = []

        files = [Path(f).resolve() for f in files]

        for f in files:
            # Decompress files to run detection on content.
            decomp = self._decompress(f)
            with decomp.open("rb") as f_:
                head = f_.read(4096)

            # Portable Executable and ELF files
            if head.startswith(b"MZ") or head.startswith(b"ELF"):
                self.logger.info(
                    "File '{}' identified as a vmlinuz."
                    .format(f)
                )
                vmlinuz.append(decomp)

            # Cpio files
            elif (
                head.startswith(b"070701")
                or head.startswith(b"070702")
                or head.startswith(b"070707")
            ):
                self.logger.info(
                    "File '{}' identified as an initramfs."
                    .format(f)
                )
                # Keep initramfs compressed
                initramfs.append(f)

                # Avoid name collision when copying initramfs to tmpdir
                if decomp != f:
                    decomp.unlink()

            # Device-tree blobs
            elif head.startswith(b"\xd0\x0d\xfe\xed"):
                self.logger.info(
                    "File '{}' identified as a device-tree blob."
                    .format(f)
                )
                dtbs.append(decomp)

            # Failed to detect, assume in the order in usage string
            elif len(vmlinuz) == 0:
                self.logger.info(
                    "Assuming file '{}' is a vmlinuz."
                    .format(f)
                )
                vmlinuz.append(decomp)

            elif len(initramfs) == 0:
                self.logger.info(
                    "Assuming file '{}' is an initramfs."
                    .format(f)
                )
                # Keep initramfs compressed
                initramfs.append(f)

            else:
                self.logger.info(
                    "Assuming file '{}' is a device-tree blob."
                    .format(f)
                )
                dtbs.append(decomp)

        return {
            "vmlinuz": vmlinuz,
            "initramfs": initramfs,
            "dtbs": dtbs,
        }

    @input_files.add
    @Argument("-d", "--vmlinuz", nargs=1)
    def vmlinuz(self, vmlinuz=None):
        """Kernel executable"""
        files = self.files["vmlinuz"]

        if vmlinuz is not None:
            vmlinuz = Path(vmlinuz).resolve()
            self.logger.info(
                "Using file '{}' as a vmlinuz."
                .format(vmlinuz)
            )
            files = [self._decompress(vmlinuz), *files]

        if len(files) == 0:
            raise ValueError(
                "vmlinuz argument is required."
            )

        elif len(files) > 1:
            raise ValueError(
                "Can't build with multiple kernels"
            )

        vmlinuz = self._copy(files[0])

        return vmlinuz

    @input_files.add
    @Argument("-i", "--initramfs", nargs=1)
    def initramfs(self, initramfs=None):
        """Ramdisk image"""
        files = self.files["initramfs"]

        if initramfs is not None:
            initramfs = Path(initramfs).resolve()
            self.logger.info(
                "Using file '{}' as an initramfs."
                .format(initramfs)
            )
            files = [initramfs, *files]

        if len(files) > 1:
            raise ValueError(
                "Can't build with multiple initramfs"
            )

        if files:
            initramfs = self._copy(files[0])
        else:
            initramfs = None

        return initramfs

    @input_files.add
    @Argument("-b", "--dtbs", metavar="DTB", nargs="+")
    def dtbs(self, *dtbs):
        """Device-tree binary file"""
        files = self.files["dtbs"]

        dtbs = [Path(dtb).resolve() for dtb in dtbs]
        for dtb in dtbs:
            self.logger.info(
                "Using file '{}' as a device-tree blob."
                .format(dtb)
            )

        dtbs = [self._decompress(dtb) for dtb in dtbs]
        dtbs = [self._copy(dtb) for dtb in (*dtbs, *files)]

        return dtbs

    @Group
    def options(self):
        """Options"""
        # Check incompatible combinations
        if self.image_format == "zimage":
            if self.compress not in (None, "none"):
                raise ValueError(
                    "Compress argument not supported with zimage format."
                )
            if self.name is not None:
                raise ValueError(
                    "Name argument not supported with zimage format."
                )
            if self.dtbs:
                raise ValueError(
                    "Device tree files not supported with zimage format."
                )

    @options.add
    @Argument("-h", "--help", action="help")
    def print_help(self):
        """Show this help message."""
        # type(self).parser.print_help()

    @options.add
    @Argument(
        "-V", "--version",
        action="version",
        version="depthcharge-tools %(prog)s {}".format(__version__),
    )
    def version(self):
        """Print program version."""
        return type(self).version.version % {"prog": type(self).prog}

    @options.add
    @Argument("-v", "--verbose", count=True)
    def verbosity(self, verbosity=0):
        """Print more detailed output."""
        level = logging.WARNING - int(verbosity) * 10
        self._logger.setLevel(level)
        return verbosity

    @options.add
    @Argument("-o", "--output", required=True)
    def output(self, file_):
        """Write resulting image to FILE."""

        # Output path is obviously required
        if file_ is None:
            raise ValueError(
                "Output argument is required."
            )

        return Path(file_).resolve()

    @options.add
    @Argument("--tmpdir", nargs=1)
    def tmpdir(self, dir_=None):
        """Directory to keep temporary files."""
        if dir_ is None:
            dir_ = tempfile.TemporaryDirectory(
                prefix="mkdepthcharge-",
            )
            dir_ = self.exitstack.enter_context(dir_)

        dir_ = Path(dir_)
        os.makedirs(dir_, exist_ok=True)

        self.logger.debug("Working in temp dir '{}'.".format(dir_))

        return dir_

    @options.add
    @Argument("-A", "--arch", nargs=1)
    def arch(self, arch=None):
        """Architecture to build for."""

        # We should be able to make an image for other architectures, but
        # the default should be whatever board the kernel is for.
        if arch is None:
            with self.vmlinuz.open("rb") as f:
                head = f.read(4096)

            if head[0x202:0x206] == b"HdrS":
                arch = Architecture("x86")
            elif head[0x38:0x3c] == b"ARM\x64":
                arch = Architecture("arm64")
            elif head[0x34:0x38] == b"\x45\x45\x45\x45":
                arch = Architecture("arm")

            self.logger.info(
                "Assuming CPU architecture '{}' from vmlinuz file."
                .format(arch)
            )

        elif arch not in Architecture.all:
            raise ValueError(
                "Can't build images for unknown architecture '{}'"
                .format(arch)
            )

        return Architecture(arch)

    @options.add
    @Argument("--format", nargs=1)
    def image_format(self, format_=None):
        """Kernel image format to use."""

        # Default to architecture-specific formats.
        if format_ is None:
            if self.arch in Architecture.arm:
                format_ = "fit"
            elif self.arch in Architecture.x86:
                format_ = "zimage"
            self.logger.info("Assuming image format '{}'.".format(format_))

        if format_ not in ("fit", "zimage"):
            raise ValueError(
                "Can't build images for unknown image format '{}'"
                .format(format_)
            )

        return format_

    @options.add
    @Argument("--kernel-start", nargs=1)
    def kernel_start(self, addr=None):
        """Start of depthcharge kernel buffer in memory."""
        if addr is not None:
            return parse_bytesize(addr)

        if self.arch in Architecture.x86:
            return 0x100000

    @Group
    def fit_options(self):
        """FIT image options"""

    @fit_options.add
    @Argument("-C", "--compress", nargs=1)
    def compress(self, type_=None):
        """Compress vmlinuz file before packing."""

        # We need to pass "-C none" to mkimage or it assumes gzip.
        if type_ is None and self.image_format == "fit":
            type_ = "none"

        if type_ not in (None, "none", "lz4", "lzma"):
            raise ValueError(
                "Compression type '{}' is not supported."
                .format(type_)
            )

        return type_

    @fit_options.add
    @Argument("-n", "--name", nargs=1)
    def name(self, desc=None):
        """Description of vmlinuz to put in the FIT."""

        # If we don't pass "-n <name>" to mkimage, the kernel image
        # description is left blank. Other images get "unavailable"
        # as their description, so it looks better if we match that.
        if desc is None and self.image_format == "fit":
            desc = "unavailable"

        return desc

    @fit_options.add
    @Argument("--ramdisk-load-address", nargs=1)
    def ramdisk_load_address(self, addr=None):
        """Add load address to FIT ramdisk image section."""
        if addr is not None:
            return parse_bytesize(addr)

        return None

    @fit_options.add
    @Argument(
        "--no-patch-dtbs", patch_dtbs=False,
        help=argparse.SUPPRESS,
    )
    @Argument(
        "--patch-dtbs", patch_dtbs=True,
        help="Add linux,initrd properties to device-tree binary files.",
    )
    def patch_dtbs(self, patch_dtbs=False):
        """Add linux,initrd properties to device-tree binary files."""
        if (
            patch_dtbs
            and self.kernel_start is None
            and self.ramdisk_load_address is None
        ):
            raise ValueError(
                "The kernel buffer start address or a ramdisk load address "
                "is required to patch DTB files for initramfs support."
            )

        return bool(patch_dtbs)

    @Group
    def zimage_options(self):
        """zImage format options"""

    @zimage_options.add
    @Argument(
        "--no-pad-vmlinuz", pad=False,
        help="Don't pad the vmlinuz file for safe decompression",
    )
    @Argument(
        "--pad-vmlinuz", pad=True,
        help=argparse.SUPPRESS,
    )
    def pad_vmlinuz(self, pad=None):
        """Pad vmlinuz for safe decompression"""
        if pad is None:
            return (
                self.image_format == "zimage"
                and self.initramfs is not None
            )

        return bool(pad)

    @Group
    def vboot_options(self):
        """Depthcharge image options"""

        keydirs = []
        if self.keydir is not None:
            keydirs += [self.keydir]

        # If any of the arguments are given, search nearby for others
        if self.keyblock is not None:
            keydirs += [self.keyblock.parent]
        if self.signprivate is not None:
            keydirs += [self.signprivate.parent]
        if self.signpubkey is not None:
            keydirs += [self.signpubkey.parent]

        if None in (self.keyblock, self.signprivate, self.signpubkey):
            for d in sorted(set(keydirs), key=keydirs.index):
                self.logger.info(
                    "Searching '{}' for vboot keys."
                    .format(d)
                )

            # Defaults to distro-specific paths for necessary files.
            keydir, keyblock, signprivate, signpubkey = vboot_keys(*keydirs)

            if keydir:
                self.logger.info(
                    "Defaulting to keys from '{}' for missing arguments."
                    .format(keydir)
                )

            if self.keyblock is None:
                self.keyblock = keyblock
            if self.signprivate is None:
                self.signprivate = signprivate
            if self.signpubkey is None:
                self.signpubkey = signpubkey

        # We might still not have the vboot keys after all that.
        if self.keyblock is None:
            raise ValueError(
                "Couldn't find a usable keyblock file."
            )
        elif not self.keyblock.is_file():
            raise ValueError(
                "Keyblock file '{}' does not exist."
                .format(self.keyblock)
            )
        else:
            self.logger.info(
                "Using keyblock file '{}'."
                .format(self.keyblock)
            )

        if self.signprivate is None:
            raise ValueError(
                "Couldn't find a usable signprivate file."
            )
        elif not self.signprivate.is_file():
            raise ValueError(
                "Signprivate file '{}' does not exist."
                .format(self.signprivate)
            )
        else:
            self.logger.info(
                "Using signprivate file '{}'."
                .format(self.signprivate)
            )

        if self.signpubkey is None:
            self.logger.warning(
                "Couldn't find a usable signpubkey file."
            )
        elif not self.signpubkey.is_file():
            self.logger.warning(
                "Signpubkey file '{}' does not exist."
                .format(self.keyblock)
            )
            self.signpubkey = None
        else:
            self.logger.info(
                "Using signpubkey file '{}'."
                .format(self.signpubkey)
            )

    @vboot_options.add
    @Argument("-c", "--cmdline", append=True, nargs="+")
    def cmdline(self, *cmd):
        """Command-line parameters for the kernel."""

        # If the cmdline is empty vbutil_kernel returns an error. We can use
        # "--" instead of putting a newline or a space into the cmdline.
        if len(cmd) == 0:
            cmdline = "--"
        elif len(cmd) == 1 and isinstance(cmd[0], str):
            cmdline = cmd[0]
        elif isinstance(cmd, (list, tuple)):
            cmdline = " ".join(cmd)

        # The firmware replaces any '%U' in the kernel cmdline with the
        # PARTUUID of the partition it booted from. Chrome OS uses
        # kern_guid=%U in their cmdline and it's useful information, so
        # prepend it to cmdline.
        if (self.kern_guid is None) or self.kern_guid:
            cmdline = " ".join(("kern_guid=%U", cmdline))

        return cmdline

    @vboot_options.add
    @Argument(
        "--no-kern-guid", kern_guid=False,
        help="Don't prepend kern_guid=%%U to the cmdline."
    )
    @Argument(
        "--kern-guid", kern_guid=True,
        help=argparse.SUPPRESS,
    )
    def kern_guid(self, kern_guid=True):
        """Prepend kern_guid=%%U to the cmdline."""
        return kern_guid

    @vboot_options.add
    @Argument("--bootloader", nargs=1)
    def bootloader(self, file_=None):
        """Bootloader stub binary to use."""
        if file_ is not None:
            file_ = Path(file_).resolve()

        if (
            self.image_format == "zimage"
            and self.initramfs is not None
            and file_ is not None
        ):
            raise ValueError(
                "Can't build images with both initramfs and "
                "bootloader stub for zimage format."
            )

        return file_

    @vboot_options.add
    @Argument("--keydir")
    def keydir(self, dir_):
        """Directory containing vboot keys to use."""
        if dir_ is not None:
            dir_ = Path(dir_).resolve()

        return dir_

    @vboot_options.add
    @Argument("--keyblock")
    def keyblock(self, file_):
        """The key block file (.keyblock)."""
        if file_ is not None:
            file_ = Path(file_).resolve()

        return file_

    @vboot_options.add
    @Argument("--signprivate")
    def signprivate(self, file_):
        """Private key (.vbprivk) to sign the image."""
        if file_ is not None:
            file_ = Path(file_).resolve()

        return file_

    @vboot_options.add
    @Argument("--signpubkey")
    def signpubkey(self, file_):
        """Public key (.vbpubk) to verify the image."""
        if file_ is not None:
            file_ = Path(file_).resolve()

        return file_

    def __call__(self):
        vmlinuz = self.vmlinuz
        initramfs = self.initramfs
        bootloader = self.bootloader
        dtbs = self.dtbs
        tmpdir = self.tmpdir

        if bootloader is not None:
            bootloader = copy(bootloader, tmpdir)

        # Depthcharge on arm64 with FIT supports these two compressions.
        if self.compress == "lz4":
            self.logger.info("Compressing kernel with lz4.")
            vmlinuz = lz4.compress(vmlinuz, self._tempfile("vmlinuz.lz4"))
        elif self.compress == "lzma":
            self.logger.info("Compressing kernel with lzma.")
            vmlinuz = lzma.compress(vmlinuz, self._tempfile("vmlinuz.lzma"))
        elif self.compress not in (None, "none"):
            fmt = "Compression type '{}' is not supported."
            msg = fmt.format(compress)
            raise ValueError(msg)

        # vbutil_kernel --config argument wants cmdline as a file.
        cmdline_file = self._tempfile("kernel.args")
        cmdline_file.write_text(self.cmdline)

        # vbutil_kernel --bootloader argument is mandatory, but it's
        # unused in depthcharge except as a multiboot ramdisk. Prepare
        # this empty file as its replacement where necessary.
        empty = self._tempfile("empty.bin")
        empty.write_bytes(bytes(512))

        if self.image_format == "fit":
            fit_image = self._tempfile("depthcharge.fit")

            initramfs_args = []
            if initramfs is not None:
                initramfs_args += ["-i", initramfs]

            dtb_args = []
            for dtb in dtbs:
                dtb_args += ["-b", dtb]

            # The subimage nodes can be <type>@1 or <type>-1.
            def subimage_by_type(fit_image, subimage_type):
                for subimage in fdtget.subnodes(fit_image, "/images"):
                    node = "/images/{}".format(subimage)
                    try:
                        if fdtget.get(fit_image, node, "type") == subimage_type:
                            return node
                    except:
                        continue

            # The later 32-bit ARM Chromebooks use Depthcharge, but
            # their stock versions don't have the code to support FIT
            # ramdisks. But since we know the fixed KERNEL_START we can
            # deduce where the initramfs will be, and inject its address
            # into the DTBs the way Linux expects bootloaders to do.
            if initramfs is not None and self.patch_dtbs:

                # We'll probably never need this, as only old U-Boot builds
                # need a ramdisk load address and those can handle the
                # initrd properties fine.
                if self.ramdisk_load_address:
                    initrd_start = self.ramdisk_load_address
                    initrd_end = initrd_start + initramfs.stat().st_size

                else:
                    # Allocate space for the properties we want to set,
                    # adding them later would shift things around.
                    self.logger.info("Preparing dtb files for initramfs support.")
                    for dtb in dtbs:
                        fdtput.put(dtb, "/chosen", "linux,initrd-start", 0)
                        fdtput.put(dtb, "/chosen", "linux,initrd-end", 0)

                    # Make a temporary image and search for the initramfs
                    # inside it, because I don't want to risk a wrong
                    # estimate and don't want to mess with pylibfdt.
                    self.logger.info("Packing files as temp FIT image:")
                    tmp_image = self._tempfile("depthcharge.fit.tmp")
                    proc = mkimage(
                        "-f", "auto",
                        "-A", self.arch.mkimage,
                        "-T", "kernel",
                        "-O", "linux",
                        "-C", self.compress,
                        "-n", self.name,
                        *initramfs_args,
                        *dtb_args,
                        "-d", vmlinuz,
                        tmp_image,
                    )
                    self.logger.info(proc.stdout)

                    # Mkimage breaks the config node key with -T kernel_noload.
                    # Apparently this shifts things around as well, so...
                    self.logger.info("Patching temp FIT for kernel_noload type.")
                    fdtput.put(
                        tmp_image, subimage_by_type(tmp_image, "kernel"),
                        "type", "kernel_noload",
                    )

                    with tmp_image.open("r+b") as f, mmap(f.fileno(), 0) as data:
                        initrd_offset = data.find(initramfs.read_bytes())
                        self.logger.info(
                            "Initramfs is at offset {:#x} in FIT image."
                            .format(initrd_offset)
                        )
                    initrd_start = initrd_offset + self.kernel_start
                    initrd_end = initrd_start + initramfs.stat().st_size

                self.logger.info(
                    "Initramfs should be at address {:#x} - {:#x} in memory."
                    .format(initrd_start, initrd_end)
                )

                self.logger.info("Patching dtb files for initramfs support.")
                for dtb in dtbs:
                    fdtput.put(dtb, "/chosen", "linux,initrd-start", initrd_start)
                    fdtput.put(dtb, "/chosen", "linux,initrd-end", initrd_end)

            self.logger.info("Packing files as FIT image:")
            proc = mkimage(
                "-f", "auto",
                "-A", self.arch.mkimage,
                "-T", "kernel",
                "-O", "linux",
                "-C", self.compress,
                "-n", self.name,
                *initramfs_args,
                *dtb_args,
                "-d", vmlinuz,
                fit_image,
            )
            self.logger.info(proc.stdout)

            # Earlier 32-bit ARM Chromebooks use U-Boot, which needs a
            # usable load address for the FIT ramdisk image section.
            if initramfs is not None and self.ramdisk_load_address:
                self.logger.info("Patching FIT for ramdisk load address.")
                fdtput.put(
                    fit_image, subimage_by_type(fit_image, "ramdisk"),
                    "load", self.ramdisk_load_address,
                )

            # Mkimage breaks the config node key with -T kernel_noload.
            self.logger.info("Patching FIT for kernel_noload type.")
            fdtput.put(
                fit_image, subimage_by_type(fit_image, "kernel"),
                "type", "kernel_noload",
            )

            if (
                initramfs is not None
                and self.patch_dtbs
                and self.ramdisk_load_address is None
            ):
                with fit_image.open("r+b") as f, mmap(f.fileno(), 0) as data:
                    if initrd_offset != data.find(initramfs.read_bytes()):
                        raise RuntimeError(
                            "Initramfs FIT offset changed after rebuild."
                        )

            self.logger.info("Packing files as depthcharge image.")
            proc = vbutil_kernel(
                "--version", "1",
                "--arch", self.arch.vboot,
                "--vmlinuz", fit_image,
                "--config", cmdline_file,
                "--bootloader", bootloader or empty,
                "--keyblock", self.keyblock,
                "--signprivate", self.signprivate,
                "--pack", self.output,
            )
            self.logger.info(proc.stdout)

        elif self.image_format == "zimage" and initramfs is None:
            self.logger.info("Packing files as depthcharge image.")
            proc = vbutil_kernel(
                "--version", "1",
                "--arch", self.arch.vboot,
                "--vmlinuz", vmlinuz,
                "--config", cmdline_file,
                "--bootloader", bootloader or empty,
                "--keyblock", self.keyblock,
                "--signprivate", self.signprivate,
                "--pack", self.output,
            )
            self.logger.info(proc.stdout)

        elif self.image_format == "zimage":
            # The bzImage overwrites parts of the buffer we control
            # while decompressing itself. We need to make sure we don't
            # place initramfs in that range. For that, we need to know
            # how offsets in file correspond to addresses in memory.

            def addr_to_offs(addr, load_addr=self.kernel_start):
                return addr - load_addr + 0x10000

            def offs_to_addr(offs, load_addr=self.kernel_start):
                return offs + load_addr - 0x10000

            def align_up(size, align=0x1000):
                return ((size + align - 1) // align) * align

            # bzImage header has the address the kernel will decompress
            # to, and the amount of memory it needs there to work.
            # See Documentation/x86/boot.rst in Linux tree for offsets.
            with vmlinuz.open("r+b") as f, mmap(f.fileno(), 0) as data:
                if data[0x202:0x206] != b"HdrS":
                    raise ValueError(
                        "Vmlinuz file is not a Linux kernel bzImage."
                    )

                pref_address, init_size = struct.unpack(
                    "<QI", data[0x258:0x264]
                )
                pad_to = align_up(addr_to_offs(pref_address + init_size))

                # Custom kernels may have PHYSICAL_START high enough
                # that we can squeeze the initramfs between vmlinuz and
                # the preferred_address. Very unlikely, but saves space.
                # Didn't work unless I padded the vmlinuz by 0x15000.
                low_pad = 0x18000
                low_usable = pref_address - align_up(data.size()) - low_pad
                if initramfs.stat().st_size < low_usable:
                    pad_to = align_up(data.size()) + low_pad

                if self.pad_vmlinuz and pad_to > data.size():
                    self.logger.info(
                        "Padding vmlinuz to size {:#x}"
                        .format(pad_to)
                    )
                    data.resize(pad_to)

            # vbutil_kernel picks apart the vmlinuz in ways I don't
            # really want to reimplement right now, so just call it.
            self.logger.info("Packing files as temporary image.")
            temp_img = self._tempfile("temp.img")
            proc = vbutil_kernel(
                "--version", "1",
                "--arch", self.arch.vboot,
                "--vmlinuz", vmlinuz,
                "--config", cmdline_file,
                "--bootloader", initramfs,
                "--keyblock", self.keyblock,
                "--signprivate", self.signprivate,
                "--pack", temp_img,
            )
            self.logger.info(proc.stdout)

            # Do binary editing for now, until I get time to write
            # parsers for vboot_reference structs and kernel headers.
            with temp_img.open("r+b") as f, mmap(f.fileno(), 0) as data:
                if data[:8] != b"CHROMEOS":
                    raise RuntimeError(
                        "Unexpected output format from vbutil_kernel, "
                        "expected 'CHROMEOS' magic at start of file."
                    )

                # File starts with a keyblock and a kernel preamble
                # immediately afterwards, and padding up to 0x10000.
                keyblock_size = struct.unpack(
                    "<I", data[0x10:0x14]
                )[0]
                p = preamble_offset = keyblock_size

                # Preamble has the "memory address" of the "bootloader"
                # but it assumes the body is loaded at 0x100000.
                bootloader_addr = struct.unpack(
                    "<I", data[p+0x38:p+0x3c]
                )[0]
                bootloader_offset = addr_to_offs(bootloader_addr, 0x100000)

                # Assume vbutil_kernel correctly put it as "bootloader"
                initramfs_offset = bootloader_offset
                initramfs_addr = offs_to_addr(initramfs_offset)
                initramfs_size = initramfs.stat().st_size

                self.logger.info(
                    "Initramfs is at offset {:#x}, address {:#x}, size {:#x}."
                    .format(initramfs_offset, initramfs_addr, initramfs_size)
                )

                # Params is immediately before bootloader with size 0x1000
                p = params_offset = bootloader_offset - 0x1000
                if data[p+0x202:p+0x206] != b"HdrS":
                    raise RuntimeError(
                        "Unexpected output format from vbutil_kernel, "
                        "expected 'HdrS' magic in boot params."
                    )

                # These get passed to the kernel unmodified by depthcharge.
                # "initrdmem=addr,size" in the cmdline would work, but
                # this looks like how bootloaders are supposed to do it.
                data[p+0x218:p+0x21c] = struct.pack("<I", initramfs_addr)
                data[p+0x21c:p+0x220] = struct.pack("<I", initramfs_size)

            self.logger.info("Re-signing edited temporary image.")
            proc = vbutil_kernel(
                "--keyblock", self.keyblock,
                "--signprivate", self.signprivate,
                "--oldblob", temp_img,
                "--repack", self.output,
            )
            self.logger.info(proc.stdout)

        self.logger.info("Verifying built depthcharge image:")
        signpubkey_args = []
        if self.signpubkey is not None:
            signpubkey_args += ["--signpubkey", self.signpubkey]

        proc = vbutil_kernel(
            "--verify", self.output,
            *signpubkey_args,
        )
        self.logger.info(proc.stdout)

        return self.output


if __name__ == "__main__":
    mkdepthcharge.main()
