#!/bin/bash

# Configure grub. Note that the various conditionals here are to handle
# different distributions gracefully.

if [ ${DIB_DEBUG_TRACE:-1} -gt 0 ]; then
    set -x
fi
set -eu
set -o pipefail

if [ ${DIB_EXTLINUX:-0} != "0" ]; then
    echo "DIB_EXTLINUX no longer supported"
    exit 1
fi

# Some distros have pre-installed grub in some other way, and want to
# skip this.
if [[ -f "/tmp/grub/install" ]]; then
    exit 0
fi

BOOT_DEV=$IMAGE_BLOCK_DEVICE
# All available devices, handy for some bootloaders...
declare -A DEVICES
eval DEVICES=( $IMAGE_BLOCK_DEVICES )

DIB_BLOCK_DEVICE=${DIB_BLOCK_DEVICE:-}

# Right now we can't use pkg-map to branch by arch, so tag an
# architecture specific virtual package so we can install the
# rigth thing based on distribution.
if [[ "$ARCH" =~ "ppc" ]]; then
    GRUB_PACKAGES="grub-ppc64"
elif [[ "${DIB_BLOCK_DEVICE}" == "mbr" ||
            "${DIB_BLOCK_DEVICE}" == "gpt" ]]; then
    GRUB_PACKAGES="grub-pc"
elif [[ "${DIB_BLOCK_DEVICE}" == "efi" ]]; then
    GRUB_PACKAGES="grub-efi grub-efi-$ARCH"
else
    GRUB_PACKAGES="grub-pc grub-efi grub-efi-$ARCH"
fi

if [[ "False" == "${DIB_SKIP_GRUB_PACKAGE_INSTALL:-False}" ]]; then
    install-packages -m bootloader $GRUB_PACKAGES
fi

GRUBNAME=$(type -p grub-install) || echo "trying grub2-install"
if [ -z "$GRUBNAME" ]; then
    GRUBNAME=$(type -p grub2-install)
fi

if type grub2-mkconfig >/dev/null; then
    GRUB_MKCONFIG="grub2-mkconfig"
else
    GRUB_MKCONFIG="grub-mkconfig"
fi

if [[ ! $($GRUBNAME --version) =~ ' 2.' ]]; then
    echo "Failure: not grub2"
    exit 1
fi

# Ensure BLS entries are updated with default args
if [[ $($GRUB_MKCONFIG --help) =~ '--update-bls-cmdline' ]]; then
    GRUB_MKCONFIG="$GRUB_MKCONFIG --update-bls-cmdline"
fi

# Some distros keep things in /boot/grub2, others in /boot/grub
if [ -d /boot/grub2 ]; then
    GRUB_CFG=/boot/grub2/grub.cfg
    GRUBENV=/boot/grub2/grubenv
else
    # NOTE(ianw) This used to be behind a "-d /boot/grub" but this
    # directory doesn't seem to exist for gentoo at this point;
    # something creates it later.  So we just fallback to this
    # unconditionally.
    GRUB_CFG=/boot/grub/grub.cfg
    GRUBENV=/boot/grub/grubenv
    mkdir -p /boot/grub
fi

# When using EFI image-based builds, particularly rhel element
# based on RHEL>=8.2 .qcow2, we might have /boot/grub2/grubenv
# as a dangling symlink to /boot/efi because we have extracted
# it from the root fs, but we didn't populate the separate EFI
# boot partition from the image.  grub2-install calls rename()
# on this file, so if it's a dangling symlink it errors.  Just
# remove it if it exists.
if [[ -L $GRUBENV ]]; then
    rm -f $GRUBENV
fi

function append_or_replace {
    local file=$1
    local key=$2
    local value=$3

    if grep -q "^$key=" "$file"; then
        sed -i "s/^$key=.*/$key=$value/" "$file"
    else
        echo $key=$value >>"$file"
    fi
}


# NOTE(TheJulia): Explicitly strip out embedded console settings
# in the source's grub defaults. This prevents re-introduction later.
# Intentionally matching 'console=ttyS0,115200n8 console=tty0'
# Also globally matches so it can handle 'console=tty0 stuff console=tty1'
# ... and also matches
# 'console=ttyS0,112500n81 console=ttyS1,n8 console=tty0'
if [[ -f /etc/default/grub ]]; then
    sed -i 's/\<console=[[:alnum:],]*[[:space:]]*//g' /etc/default/grub
fi

append_or_replace /etc/default/grub GRUB_DEVICE "LABEL=${DIB_ROOT_LABEL}"
append_or_replace /etc/default/grub GRUB_DISABLE_LINUX_UUID true
append_or_replace /etc/default/grub GRUB_TIMEOUT "${DIB_GRUB_TIMEOUT:-1}"
append_or_replace /etc/default/grub GRUB_TIMEOUT_STYLE "${DIB_GRUB_TIMEOUT_STYLE:-hidden}"
if [[ "True" == "${DIB_BOOTLOADER_USE_SERIAL_CONSOLE:-True}" ]]; then
    append_or_replace /etc/default/grub GRUB_TERMINAL \""serial console"\"
else
    append_or_replace /etc/default/grub GRUB_TERMINAL \"console\"
fi
append_or_replace /etc/default/grub GRUB_GFXPAYLOAD_LINUX auto

# NOTE(TheJulia): We need to remove any boot entry from the /etc/default/grub
# file that may already exist, such as what was added by fips being setup on
# either in the source image or by by an element, as we repack the image.
# with new filesystems.
# Matches any element which looks like " boot=" and the associated value
# in order for us to have a clean starting point to put a value in place,
# if applicable.
# Removes entry trailing with a space, or any entry where boot is set as
# the last argument on the line.
sed -i 's/\ boot=[0-9A-Za-z/=\-]\+//' /etc/default/grub
# NOTE(TheJulia): When using FIPS, dracut wants to evaluate
# the hmac files for the kernel checksum. However, if /boot is
# located on a separate filesystem from the root filesystem,
# than this fails. As a result, we need to identify IF /boot
# is a separate filesystem, and convey this fact as a boot
# argument so dracut does not halt the system on boot.

BOOT_FS=""

if [[ -n "${DIB_BOOT_LABEL}" ]]; then
    BOOT_FS=" boot=LABEL=${DIB_BOOT_LABEL}"
fi

# NOTE(TheJulia): While on the subject of FIPS, if there is not an
# explicit /boot partition, then the fips setup command will return
# a successful result, but then also tell you to update your grub
# configuration. This happens specifically with Rocky linux.
# as such, we check/reconcile the flag into place for the kernel
# as the utility will return a result code of 1 if the state is
# inconsistent, i.e. policy in place, but not kernel command line
# argument.

BOOT_FIPS=""

if [[ -x /bin/fips-mode-setup ]]; then
    set +e
    fips-mode-setup --is-enabled
    is_fips_enabled=$?
    set -e
    if [ $is_fips_enabled -eq 1 ]; then
        BOOT_FIPS=" fips=1"
    fi
fi


if [[ "True" == "${DIB_BOOTLOADER_USE_SERIAL_CONSOLE:-True}" ]]; then
    if [[ -n "${DIB_BOOTLOADER_SERIAL_CONSOLE}" ]]; then
        SERIAL_CONSOLE="console=${DIB_BOOTLOADER_SERIAL_CONSOLE}"
    elif [[ "powerpc ppc64 ppc64le" =~ "$ARCH" ]]; then
        # Serial console on Power is hvc0
        SERIAL_CONSOLE="console=hvc0"
    elif [[ "arm64" =~ "$ARCH" ]]; then
        SERIAL_CONSOLE="console=ttyAMA0,115200"
    else
        SERIAL_CONSOLE="console=ttyS0,115200"
    fi
else
    SERIAL_CONSOLE=""
fi

if [[ -n "${DIB_BOOTLOADER_VIRTUAL_TERMINAL}" ]]; then
    VIRTUAL_TERMINAL="console=${DIB_BOOTLOADER_VIRTUAL_TERMINAL}"
else
    VIRTUAL_TERMINAL=""
fi

if [[ "True" == "${DIB_NO_TIMER_CHECK:-True}" ]]; then
    NO_TIMER_CHECK="no_timer_check"
else
    NO_TIMER_CHECK=""
fi

GRUB_CMDLINE_LINUX_DEFAULT="${VIRTUAL_TERMINAL} ${SERIAL_CONSOLE} ${NO_TIMER_CHECK}"
echo "GRUB_CMDLINE_LINUX_DEFAULT=\"${GRUB_CMDLINE_LINUX_DEFAULT} ${DIB_BOOTLOADER_DEFAULT_CMDLINE}${BOOT_FS}${BOOT_FIPS}\"" >>/etc/default/grub


if [[ "True" == "${DIB_BOOTLOADER_USE_SERIAL_CONSOLE:-True}" ]]; then
    echo 'GRUB_SERIAL_COMMAND="serial --speed=115200 --unit=0 --word=8 --parity=no --stop=1"' >>/etc/default/grub
fi

# os-prober leaks /dev/sda into config file in dual-boot host
# Disable grub-os-prober to avoid the issue  while running
# grub-mkconfig
# Setting a flag to track whether the entry is already there in grub config
PROBER_DISABLED=
if ! grep -qe "^\s*GRUB_DISABLE_OS_PROBER=true" /etc/default/grub; then
    PROBER_DISABLED=true
    echo 'GRUB_DISABLE_OS_PROBER=true' >> /etc/default/grub
fi

# GRUB_MKCONFIG call needs to happen after we configure
# /etc/default/grub above. Without this we can set inappropriate
# root device labels and then images don't boot.
#
# This produces a legacy config which both bios and uefi can boot
# Later we copy the final config to an efi specific location to
# support uefi specific functionality like secure boot.
$GRUB_MKCONFIG -o $GRUB_CFG

# If we are using BLS, for debugging purposes dump out the kernel
if [[ -e /boot/loader/entries ]]; then
    grubby --info=ALL
fi

# Remove the fix to disable os_prober
if [ -n "$PROBER_DISABLED" ]; then
    sed -i '$d' /etc/default/grub
fi

# Fix efi specific instructions in grub config file
if [ -d /sys/firmware/efi ]; then
    sed -i 's%\(initrd\|linux\)efi /boot%\1 /boot%g' $GRUB_CFG
fi

# when using efi, and having linux16/initrd16, it needs to be replaced
# by linuxefi/initrdefi. When building images on a non-efi system,
# the 16 suffix is added to linux/initrd entries, but we need it to be
# linuxefi/initrdefi for the image to boot under efi
if [[ ${DIB_BLOCK_DEVICE} == "efi" ]]; then
    sed -i 's%\(linux\|initrd\)16 /boot%\1efi /boot%g' $GRUB_CFG
fi

# Finally copy the grub.cfg and grubenv to the EFI specific dir
# to support functionality like secure boot. We make a copy because
# /boot and /boot/efi may be different partitions and uefi looks
# for a specific partition UUID preventing symlinks from working.
if [ -n "${EFI_BOOT_DIR:-}" ] && [ -d /boot/efi/$EFI_BOOT_DIR ] ; then
    cp $GRUB_CFG /boot/efi/$EFI_BOOT_DIR/grub.cfg
    if [ -a $GRUBENV ]; then
        cp $GRUBENV /boot/efi/$EFI_BOOT_DIR/grubenv
    fi
fi

# Ensure paths in BLS entries account for /boot being a partition or part of the
# root partition and perform any required entry removal.
if [[ -e /boot/loader/entries ]]; then
    pushd /boot/loader/entries
    set +e
    mountpoint /boot
    bootmount=$?

    for entry in *; do
        if [[ $bootmount -eq 0 ]]; then
            sed -i "s| /boot/vmlinuz| /vmlinuz|" $entry
            sed -i "s| /boot/initramfs| /initramfs|" $entry
        else
            sed -i "s| /vmlinuz| /boot/vmlinuz|" $entry
            sed -i "s| /initramfs| /boot/initramfs|" $entry
        fi
    done
    set -e
    popd

    # Since we are already aware we're using BLS, we can go ahead and
    # do any cleanup to disable embedded serial console settings if they
    # already exist.
    if [[ "True" != "${DIB_BOOTLOADER_USE_SERIAL_CONSOLE:-True}" ]] && [[ "${VIRTUAL_TERMINAL}" == "" ]]; then
        # NOTE(TheJulia): This removes any console arguments and allows grub
        # to use it's default. It will also remove any embedded consoles in
        # source images, which may be highly desirable if your running
        # interrupt sensitive workloads, such as NFV workloads.
        grubby --update-kernel ALL --remove-args=console
    fi
    # Print resulting grubby output for debug purposes
    grubby --info=ALL
fi

if [[ ! "$ARCH" =~ "ppc" ]] && [[ -z  "${DIB_BLOCK_DEVICE}" ]]; then
    echo "WARNING: No bootloader installation will occur."
    echo "To install a bootloader ensure you have included a block-device-* element"
    exit 0
fi

echo "Installing GRUB2..."

# We need --force so grub does not fail due to being installed on the
# root partition of a block device.
GRUB_OPTS="--force "

if [[ "$ARCH" =~ "ppc" ]] ; then
    # For PPC (64-Bit regardless of Endian-ness), we use the "boot"
    # partition as the one to point grub-install to, not the loopback
    # device.  ppc has a dedicated PReP boot partition.
    # For grub2 < 2.02~beta3 this needs to be a /dev/mapper/... node after
    # that a dev/loopXpN node will work fine.
    $GRUBNAME --modules="part_msdos" $GRUB_OPTS ${DEVICES[boot]} --no-nvram
else
    # This set of modules is sufficient for all installs (mbr/gpt/efi)
    modules="part_msdos part_gpt lvm"

    if [[ ${DIB_BLOCK_DEVICE} == "mbr" || ${DIB_BLOCK_DEVICE} == "gpt" ]]; then
        if [[ ! "x86_64 amd64" =~ ${ARCH} ]]; then
            echo "*** ${ARCH} is not supported by mbr/gpt"
        fi
        $GRUBNAME --modules="$modules biosdisk" --target=i386-pc \
                $GRUB_OPTS $BOOT_DEV
    elif [[ ${DIB_BLOCK_DEVICE} == "efi" ]]; then
        # We need to manually set the target if it's different to
        # the host.  Setup for EFI
        case $ARCH in
            "x86_64"|"amd64")
                # This call installs grub for BIOS compatability
                # which makes portable EFI/BIOS images.
                $GRUBNAME --modules="$modules" --target=i386-pc $BOOT_DEV
                # Set the x86_64 specific efi target for the generic
                # installation below.
                GRUB_OPTS="--target=x86_64-efi"
                ;;
            # At this point, we don't need to override the target
            # for any other architectures.
        esac
        # If we don't have a distro specific dir with presigned efi targets
        # we install a generic one.
        if [ ! -d /boot/efi/$EFI_BOOT_DIR ]; then
            echo "WARNING: /boot/efi/$EFI_BOOT_DIR does not exist, UEFI secure boot not supported"
            # This tells the EFI install to put the EFI binaries into
            # the generic /BOOT directory and avoids trying to update
            # nvram settings.
            extra_options="--removable"
            $GRUBNAME --modules="$modules" $extra_options $GRUB_OPTS $BOOT_DEV
        fi
    fi
fi