#!/bin/bash

# Copyright (C) 2025 Matthias Luescher
#
# Authors:
#  Matthias Luescher
#
# This file is part of edi-boot-shim.
#
# edi-boot-shim is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# edi-boot-shim is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with edi-boot-shim.  If not, see <http://www.gnu.org/licenses/>.

# ----------------------------------------------------------------------------
# Small utility that helps updating the u-boot boot configuration
# without copying around initial ramdisks and kernels.
# ----------------------------------------------------------------------------

set -o errexit
set -o nounset
set -o pipefail

echo_usage()
{
    echo "Usage: edi-boot-shim KERNEL_PACKAGE_NAME"
    echo "Example: edi-boot-shim linux-image-4.18.0-2-arm64"
    echo "Please check the man page for more information."
}

cleanup()
{
    rm -f "${GENERATED_SCRIPT}"
    rm -f "${GENERATED_SCRIPT_ITS}"
    rm -f "${GENERATED_UIMAGE_ITS}"
}

cleanup_and_exit()
{
    cleanup
    >&2 echo "Error: Abnormal script termination."
    exit 1
}

copy_or_link()
{
    local SOURCE="${1}"
    local DEST="${2}"
    if [ "$(stat -c "%m" "${SOURCE%/*}")" == "$(stat -c "%m" "${DEST%/*}")" ]
    then
        ln -srf "${SOURCE}" "${DEST}"
    else
        cp "${SOURCE}" "${DEST}"
    fi
}

DEBIAN_ARCHITECTURE="$(dpkg --print-architecture)"

# default values
BOOTLOADER_TYPE="u-boot"
BOOTLOADER_DIRECTORY="/boot/firmware"
CHECK_MOUNT=true
BOOT_COMMAND_TEMPLATE="/etc/edi-boot-shim/boot.cmd.rpi.${DEBIAN_ARCHITECTURE}"
UIMAGE_FIT_TEMPLATE=""
SCRIPT_FIT_TEMPLATE="/etc/edi-boot-shim/boot.scr.its"

EDI_BOOT_SHIM_CFG="/etc/edi-boot-shim/edi-boot-shim.cfg"
if [ -f "${EDI_BOOT_SHIM_CFG}" ]
then
    # shellcheck source=../config/edi-boot-shim.cfg
    source "${EDI_BOOT_SHIM_CFG}"
fi

CONFIG_PARTS="${EDI_BOOT_SHIM_CFG}.d"
if [ -d "${CONFIG_PARTS}" ]
then
    for config_part in "${CONFIG_PARTS}"/*.cfg
    do
        [[ -e "${config_part}" ]] || break
        # shellcheck source=../config/edi-boot-shim.cfg
        source "${config_part}"
    done
fi

if [[ "$#" -eq 1 && ( "$1" == "-h" || "$1" == "--help" ) ]]
then
    echo_usage
    exit 0
fi

if [ "$#" -ne 1 ]
then
    >&2 echo "Error: Illegal number of parameters."
    echo_usage
    exit 1
fi

if [[ $EUID -ne 0 ]]
then
    >&2 echo "Error: This script must be run as root."
    exit 1
fi

KERNEL_PACKAGE_NAME="${1}"

KERNEL_PACKAGE_PREFIX="linux-image-"
if [[ ${KERNEL_PACKAGE_NAME} != "${KERNEL_PACKAGE_PREFIX}"* ]]
then
    >&2 echo "Error: ${KERNEL_PACKAGE_NAME} does not look like a kernel package name."
    exit 1
fi

if ! dpkg -l "${KERNEL_PACKAGE_NAME}" > /dev/null 2>&1
then
    >&2 echo "Error: ${KERNEL_PACKAGE_NAME} is not installed."
    exit 1
fi

if dpkg -s "${KERNEL_PACKAGE_NAME}" | grep "^Depends: linux-image" > /dev/null
then
    >&2 echo "Error: ${KERNEL_PACKAGE_NAME} is a meta package."
    exit 1
fi

KERNEL_PACKAGE_VERSION="${KERNEL_PACKAGE_NAME#"${KERNEL_PACKAGE_PREFIX}"}"
GENERATED_SCRIPT="/tmp/${KERNEL_PACKAGE_NAME}.boot.cmd"
GENERATED_SCRIPT_ITS="/tmp/${KERNEL_PACKAGE_NAME}.boot.scr.its"
GENERATED_UIMAGE_ITS="/tmp/${KERNEL_PACKAGE_NAME}.uimage.its"

trap cleanup_and_exit EXIT

cp "${BOOT_COMMAND_TEMPLATE}" "${GENERATED_SCRIPT}"

sed -i 's@__KERNEL_PACKAGE_VERSION__@'"${KERNEL_PACKAGE_VERSION}"'@' "${GENERATED_SCRIPT}"

# This script assumes that the bootloader partition is mounted to BOOTLOADER_DIRECTORY.

if [ "${CHECK_MOUNT}" = false ]
then
    # Mount point check disabled by configuration.
    true
elif ischroot
then
    # Chroot detected - skip mount point check.
    true
elif test -e /usr/bin/systemd-detect-virt && systemd-detect-virt -q
then
    # Virtualization detected - skip mount point check.
    true
elif ! mountpoint -q ${BOOTLOADER_DIRECTORY}
then
    >&2 echo "Error: Missing ${BOOTLOADER_DIRECTORY} - you should probably mount it."
    exit 1
fi

mkdir -p ${BOOTLOADER_DIRECTORY}

if [[ ${BOOTLOADER_TYPE} == "u-boot" ]]
then
    case "${DEBIAN_ARCHITECTURE}" in
        "armhf" | "armel")
            U_BOOT_ARCHITECTURE="arm"
      ;;
        *)
            U_BOOT_ARCHITECTURE="${DEBIAN_ARCHITECTURE}"
      ;;
    esac

    if [ -z "${UIMAGE_FIT_TEMPLATE}" ]
    then
        mkimage -A "${U_BOOT_ARCHITECTURE}" -O linux -T script -d "${GENERATED_SCRIPT}" ${BOOTLOADER_DIRECTORY}/boot.scr > /dev/null
    else
        cp "${UIMAGE_FIT_TEMPLATE}" "${GENERATED_UIMAGE_ITS}"
        sed -i 's@__KERNEL_PACKAGE_VERSION__@'"${KERNEL_PACKAGE_VERSION}"'@' "${GENERATED_UIMAGE_ITS}"
        mkimage -f "${GENERATED_UIMAGE_ITS}" "${BOOTLOADER_DIRECTORY}/uimage.${KERNEL_PACKAGE_VERSION}" > /dev/null

        cp "${SCRIPT_FIT_TEMPLATE}" "${GENERATED_SCRIPT_ITS}"
        sed -i 's@__BOOT_COMMAND_SCRIPT__@'"${GENERATED_SCRIPT}"'@' "${GENERATED_SCRIPT_ITS}"
        sed -i 's@__U_BOOT_ARCHITECTURE__@'"${U_BOOT_ARCHITECTURE}"'@' "${GENERATED_SCRIPT_ITS}"
        mkimage -f "${GENERATED_SCRIPT_ITS}" ${BOOTLOADER_DIRECTORY}/boot.scr > /dev/null
    fi
elif  [[ ${BOOTLOADER_TYPE} == "raspberry-pi" ]]
then
    kernel_image=""
    device_tree_binaries=""
    device_tree_binary_overlays=""
    # shellcheck source=../template/boot.spec.rpi5.arm64
    source "${GENERATED_SCRIPT}"
    copy_or_link "${kernel_image}" "${BOOTLOADER_DIRECTORY}/vmlinuz"
    if [ -n "${initramfs_image}" ]
    then
        if [ -f "${initramfs_image}" ]
        then
            copy_or_link "${initramfs_image}" "${BOOTLOADER_DIRECTORY}/initrd.img"
        else
            echo "Remark: Initramfs image '${initramfs_image}' does not exist!"
        fi
    fi
    for dtb in "${device_tree_binaries[@]}"
    do
        if [ -f "${dtb}" ]
        then
            copy_or_link "${dtb}" "${BOOTLOADER_DIRECTORY}/"
        else
            echo "Warning: Device tree binary '${dtb}' does not exist!"
        fi
    done
    mkdir -p ${BOOTLOADER_DIRECTORY}/overlays
    # A README file needs to be present: https://www.raspberrypi.com/documentation/computers/config_txt.html#os_prefix
    echo "Dear bootloader, please take the device tree overlays from here!" > "${BOOTLOADER_DIRECTORY}/overlays/README"
    for dtbo in "${device_tree_binary_overlays[@]}"
    do
        if [ -f "${dtbo}" ]
        then
            copy_or_link "${dtbo}" "${BOOTLOADER_DIRECTORY}/overlays/"
        else
            echo "Warning: Device tree binary overlay '${dtbo}' does not exist!"
        fi
    done
else
    >&2 echo "Error: Unknown bootloader type: ${BOOTLOADER_TYPE}."
    exit 1
fi

cleanup
trap - EXIT

echo "Success: ${KERNEL_PACKAGE_NAME} is configured for next boot!"

