org.osbuild.grub2

#!/usr/bin/python3
import os
import shutil
import string
import sys

import osbuild.api

# The main grub2 configuration file template. Used for UEFI and legacy
# boot. The parameters are currently:
#   - $search: to specify the search criteria of how to locate grub's
#     "root device", i.e. the device where the "OS images" are stored.
#   - $ignition: configuration for ignition, if support for ignition
#     is enabled
GRUB_CFG_TEMPLATE = """
set timeout=${timeout}

# load the grubenv file
load_env

# selection of the next boot entry via variables 'next_entry' and
# `saved_entry` present in the 'grubenv' file. Both variables are
# set by grub tools, like grub2-reboot, grub2-set-default

if [ "${next_entry}" ] ; then
   set default="${next_entry}"
   set next_entry=
   save_env next_entry
   set boot_once=true
else
   set default="${saved_entry}"
fi

search --no-floppy --set=root $search
set boot=$${root}
function load_video {
  insmod all_video
}
${features}${serial}${terminal_input}${terminal_output}
blscfg
"""


# The grub2 redirect configuration template. This is used in case of
# hybrid (uefi + legacy) boot. In this case this configuration, which
# is located in the EFI directory, will redirect to the main grub.cfg
# (GRUB_CFG_TEMPLATE).
# The parameters are:
#   - $root: specifies the path to the grub2 directory relative to
#     to the file-system where the directory is located on
GRUB_REDIRECT_TEMPLATE = """
search --no-floppy --set prefix --file ${root}grub2/grub.cfg
set prefix=($$prefix)${root}grub2
configfile $$prefix/grub.cfg
"""


# Template for ignition support in the grub.cfg
#
# it was taken verbatim from Fedora CoreOS assembler's grub.cfg
# See https://github.com/coreos/coreos-assembler/
#
# The parameters are:
#   - $root: specifies the path to the grub2 directory relative
#     to the file-system where the directory is located on

IGNITION_TEMPLATE = """
# Ignition support
set ignition_firstboot=""
if [ -f "${root}ignition.firstboot" ]; then
    # Default networking parameters to be used with ignition.
    set ignition_network_kcmdline=''

    # Source in the `ignition.firstboot` file which could override the
    # above $ignition_network_kcmdline with static networking config.
    # This override feature is primarily used by coreos-installer to
    # persist static networking config provided during install to the
    # first boot of the machine.
    source "${root}ignition.firstboot"

    set ignition_firstboot="ignition.firstboot $${ignition_network_kcmdline}"
fi
"""


GREENBOOT = """
# greenboot support, aka boot counter and boot success reporting
insmod increment
# Check if boot_counter exists and boot_success=0 to activate this behaviour.
if [ -n "${boot_counter}" -a "${boot_success}" = "0" ]; then
  # if countdown has ended, choose to boot rollback deployment,
  # i.e. default=1 on OSTree-based systems.
  if  [ "${boot_counter}" = "0" -o "${boot_counter}" = "-1" ]; then
    set default=1
    set boot_counter=-1
  # otherwise decrement boot_counter
  else
    decrement boot_counter
  fi
  save_env boot_counter
fi

# Reset boot_success for current boot
set boot_success=0
save_env boot_success
"""


def fs_spec_decode(spec):
    for key in ["uuid", "label"]:
        val = spec.get(key)
        if val:
            return key.upper(), val
    raise ValueError("unknown filesystem type")


def copy_modules(tree, platform):
    """Copy all modules from the build image to /boot"""
    target = f"{tree}/boot/grub2/{platform}"
    source = f"/usr/lib/grub/{platform}"
    os.makedirs(target, exist_ok=True)
    for dirent in os.scandir(source):
        (_, ext) = os.path.splitext(dirent.name)
        if ext not in ('.mod', '.lst'):
            continue
        if dirent.name == "fdt.lst":
            continue
        shutil.copy2(f"/{source}/{dirent.name}", target)


def copy_font(tree):
    """Copy a unicode font into /boot"""
    os.makedirs(f"{tree}/boot/grub2/fonts", exist_ok=True)
    shutil.copy2("/usr/share/grub/unicode.pf2", f"{tree}/boot/grub2/fonts/")


def copy_efi_data(tree, efi_src_dir, vendor):
    """Copy the EFI binaries & data into /boot/efi"""
    for d in ['BOOT', vendor]:
        source = f"{efi_src_dir}/{d}"
        target = f"{tree}/boot/efi/EFI/{d}/"
        shutil.copytree(source, target,
                        symlinks=False)


# pylint: disable=too-many-instance-attributes
class GrubConfig:
    def __init__(self, rootfs, bootfs):
        self.rootfs = rootfs
        self.bootfs = bootfs
        self.path = "boot/grub2/grub.cfg"
        self.default_entry = None
        self.ignition = False
        self.greenboot = False
        self.kernel_opts = ""
        self.disable_recovery = None
        self.disable_submenu = None
        self.distributor = None
        self.serial = ""
        self.terminal = None
        self.terminal_input = None
        self.terminal_output = None
        self.timeout = 0
        self.timeout_style = None

    @property
    def grubfs(self):
        """The filesystem containing the grub files,

        This is either a separate partition (self.bootfs if set) or
        the root file system (self.rootfs)
        """
        return self.bootfs or self.rootfs

    @property
    def separate_boot(self):
        return self.bootfs is not None

    @property
    def grub_home(self):
        return "/" if self.separate_boot else "/boot/"

    def make_terminal_config(self, terminal):
        config = getattr(self, terminal)
        if not config:
            return {}

        val = (
            "\n" +
            terminal +
            " " +
            " ".join(config)
        )
        return {terminal: val}

    def write(self, tree):
        """Write the grub config to `tree` at `self.path`"""
        path = os.path.join(tree, self.path)

        fs_type, fs_id = fs_spec_decode(self.grubfs)
        type2opt = {
            "UUID": "--fs-uuid",
            "LABEL": "--label"
        }

        ignition = ""
        if self.ignition:
            tplt = string.Template(IGNITION_TEMPLATE)
            subs = {"root": self.grub_home}
            ignition = tplt.safe_substitute(subs)

        greenboot = ""
        if self.greenboot:
            greenboot = GREENBOOT

        features = "\n".join(filter(bool, [ignition, greenboot]))

        # configuration options for the main template
        config = {
            "timeout": self.timeout,
            "search": type2opt[fs_type] + " " + fs_id,
            "features": features,
        }

        if self.serial:
            config["serial"] = "\n" + self.serial

        config.update(self.make_terminal_config("terminal_input"))
        config.update(self.make_terminal_config("terminal_output"))

        tplt = string.Template(GRUB_CFG_TEMPLATE)
        data = tplt.safe_substitute(config)

        with open(path, "w", encoding="utf8") as cfg:
            cfg.write(data)

    def write_redirect(self, tree, path):
        """Write a grub config pointing to the other cfg"""
        print("hybrid boot or unified grub config enabled. "
              "Writing alias grub config")

        # configuration options for the template
        config = {
            "root": self.grub_home
        }

        tplt = string.Template(GRUB_REDIRECT_TEMPLATE)
        data = tplt.safe_substitute(config)

        with open(os.path.join(tree, path), "w", encoding="utf8") as cfg:
            cfg.write(data)

    def defaults(self):
        # NB: The "GRUB_CMDLINE_LINUX" variable contains the kernel command
        # line but without the `root=` part, thus we just use `kernel_opts`.
        data = (
            f'GRUB_CMDLINE_LINUX="{self.kernel_opts}"\n'
            f"GRUB_TIMEOUT={self.timeout}\n"
            "GRUB_ENABLE_BLSCFG=true\n"
        )

        if self.disable_recovery is not None:
            data += f"GRUB_DISABLE_RECOVERY={str(self.disable_recovery).lower()}\n"

        if self.disable_submenu is not None:
            data += f"GRUB_DISABLE_SUBMENU={str(self.disable_submenu).lower()}\n"

        if self.distributor:
            data += f'GRUB_DISTRIBUTOR="{self.distributor}"\n'

        if self.serial:
            data += f'GRUB_SERIAL_COMMAND="{self.serial}"\n'

        if self.terminal:
            val = " ".join(self.terminal)
            data += f'GRUB_TERMINAL="{val}"\n'

        if self.terminal_input:
            val = " ".join(self.terminal_input)
            data += f'GRUB_TERMINAL_INPUT="{val}"\n'

        if self.terminal_output:
            val = " ".join(self.terminal_output)
            data += f'GRUB_TERMINAL_OUTPUT="{val}"\n'

        if self.timeout_style:
            data += f'GRUB_TIMEOUT_STYLE={self.timeout_style}\n'

        if self.default_entry is not None:
            data += f'GRUB_DEFAULT={self.default_entry}\n'

        return data


# pylint: disable=too-many-statements,too-many-branches
def main(tree, options):
    root_fs = options.get("rootfs")
    boot_fs = options.get("bootfs")
    kernel_opts = options.get("kernel_opts", "")
    legacy = options.get("legacy", None)
    uefi = options.get("uefi", None)
    write_cmdline = options.get("write_cmdline", True)
    write_defaults = options.get("write_defaults", True)
    ignition = options.get("ignition", False)
    saved_entry = options.get("saved_entry")
    cfg = options.get("config", {})

    # backwards compatibility
    if not root_fs:
        root_fs = {"uuid": options["root_fs_uuid"]}

    if not boot_fs and "boot_fs_uuid" in options:
        boot_fs = {"uuid": options["boot_fs_uuid"]}

    # legacy boolean means the
    if isinstance(legacy, bool) and legacy:
        legacy = "i386-pc"

    # Check if hybrid boot support is requested, i.e. the resulting image
    # should support booting via legacy and also UEFI. In that case the
    # canonical grub.cfg and the grubenv will be in /boot/grub2. The ESP
    # will only contain a small config file redirecting to the one in
    # /boot/grub2 and will not have a grubenv itself.
    hybrid = uefi and legacy

    # Prepare the actual grub configuration file, will be written further down
    config = GrubConfig(root_fs, boot_fs)
    config.ignition = ignition
    config.greenboot = options.get("greenboot", False)
    config.kernel_opts = kernel_opts
    config.disable_recovery = cfg.get("disable_recovery")
    config.disable_submenu = cfg.get("disable_submenu")
    config.distributor = cfg.get("distributor")
    config.serial = cfg.get("serial")
    config.terminal = cfg.get("terminal")
    config.terminal_input = cfg.get("terminal_input")
    config.terminal_output = cfg.get("terminal_output")
    config.timeout = cfg.get("timeout", 0)
    config.timeout_style = cfg.get("timeout_style")
    config.default_entry = cfg.get("default")

    # Create the configuration file that determines how grub.cfg is generated.
    if write_defaults:
        os.makedirs(f"{tree}/etc/default", exist_ok=True)
        with open(f"{tree}/etc/default/grub", "w", encoding="utf8") as default:
            default.write(config.defaults())

    os.makedirs(f"{tree}/boot/grub2", exist_ok=True)
    grubenv = f"{tree}/boot/grub2/grubenv"

    if hybrid:
        # The rpm grub2-efi package will have installed a symlink from
        # /boot/grub2/grubenv to the ESP. In the case of hybrid boot we
        # want a single grubenv in /boot/grub2; therefore remove the link
        try:
            os.unlink(grubenv)
        except FileNotFoundError:
            pass

    with open(grubenv, "w", encoding="utf8") as env:
        fs_type, fs_id = fs_spec_decode(root_fs)
        data = "# GRUB Environment Block\n"

        if write_cmdline:
            data += f"kernelopts=root={fs_type}={fs_id} {kernel_opts}\n"

        if saved_entry:
            data += f"saved_entry={saved_entry}\n"

        # The 'grubenv' file is, according to the documentation,
        # a 'preallocated 1024-byte file'. The empty space is
        # needs to be filled with '#' as padding
        data += '#' * (1024 - len(data))
        assert len(data) == 1024

        print(data)
        env.write(data)

    if uefi is not None:
        # UEFI support:
        # The following config files are needed for UEFI support:
        # /boot/efi/EFI/<vendor>/
        # - grubenv: in the case of non-hybrid boot it should have
        #     been written to via the link from /boot/grub2/grubenv
        #     created by grub2-efi-{x64, ia32}.rpm
        # - grub.cfg: needs to be generated, either the canonical one
        #     or a shim one that redirects to the canonical one in
        #     /boot/grub2 in case of hybrid boot (see above)
        efi_src_dir = uefi.get("efi_src_dir", "/boot/efi/EFI")
        vendor = uefi["vendor"]
        unified = uefi.get("unified", False)

        # EFI binaries and accompanying data can be installed from
        # the build root instead of using an rpm package
        if uefi.get('install', False):
            copy_efi_data(tree, efi_src_dir, vendor)

        grubcfg = f"boot/efi/EFI/{vendor}/grub.cfg"
        if hybrid or unified:
            config.write_redirect(tree, grubcfg)
        else:
            config.path = grubcfg

    # Now actually write the main grub.cfg file
    config.write(tree)

    if legacy:
        copy_modules(tree, legacy)
        copy_font(tree)

    return 0


if __name__ == '__main__':
    args = osbuild.api.arguments()
    r = main(args["tree"], args["options"])
    sys.exit(r)