#!/bin/bash

# Init script information
NAME=grub-bootconf-dsos
DESC=""

# The number of the EFI System Partition
ESP_PART_NUM=1

# The number of the EFI loader recovery partition
EFILDRRECO_PART_NUM=2

# If the machine has a menu boot entry (yes | no)
HAS_DSOS_MENU_BOOT_ENTRY=yes

# Maximum and counter file to avoid repeatedly changing the EFI boot order
EFI_BOOT_REORDER_MAX=3
EFI_BOOT_REORDER_CNT_FILE=/var/lib/grub-bootconf-dsos/efi-boot-reorder-cnt

# The directory where the secure boot initialization keys are stored
SBINITKEYSDIR=/usr/share/resources/default/pki/sb-init

# The location of the Secure Boot key database
SBKEYSTORE=/usr/share/resources/default/pki/sb-keydb

# The GUIDs of UEFI global variables and UEFI security database variables
UEFI_GLBL_GUID=8be4df61-93ca-11d2-aa0d-00e098032b8c
UEFI_SECDB_GUID=d719b2cb-3d3a-4596-a3bc-dad00e67656f

LOGGER="logger -p daemon.err -t $NAME -s"

# Load identifiers before starting
[ -f /etc/spinetix/identifiers ] && . /etc/spinetix/identifiers

is_nfsroot() {
    grep -q '\bnfsroot=' /proc/cmdline
}

create_efi_boot_entries() {
    local ret=0
    local bootinfo
    local bdev=/dev/"$SPX_MAIN_DISK_NAME"

    # NOTE: all entries except 'SpinetiX DSOS' are created with
    # --create-only, otherwise they would become the default boot option,
    # which is not what we want. For 'SpinetiX DSOS', which is our desired
    # default boot option, we create it with --create instead so that it
    # immediately becomes the default one.

    bootinfo="$(efibootmgr)"
    if [ $? -ne 0 ]; then
        $LOGGER "ERROR: could not query current EFI boot manager data"
        return 1
    fi

    if [ ! -b "$bdev" ]; then
        $LOGGER "ERROR: block device of main disk '$SPX_MAIN_DISK_NAME' missing"
        return 1
    fi

    # Verify that the hardware type is DSOS. We have SpinetiX DSOS menu entries present only on DSOS.
    if [ "$HAS_DSOS_MENU_BOOT_ENTRY" = "yes" ]; then
        if ! echo "$bootinfo" | grep -q '^Boot[a-fA-F0-9]\+\*\s\+SpinetiX DSOS menu (redundant boot)$'; then
            $LOGGER "WARN: Create SpinetiX DSOS menu (redundant boot) boot entry"
            efibootmgr -q --create-only --label 'SpinetiX DSOS menu (redundant boot)' --disk "$bdev" --part "$ESP_PART_NUM" --loader '\EFI\DSOS\menu\grubx64.efi' || ret=1
        fi
    fi

    # Create SpinetiX DSOS (redundant boot) entry if it is not present.
    if ! echo "$bootinfo" | grep -q '^Boot[a-fA-F0-9]\+\*\s\+SpinetiX DSOS (redundant boot)$'; then
        $LOGGER "WARN: Create SpinetiX DSOS (redundant boot) boot entry"
        efibootmgr -q --create-only --label 'SpinetiX DSOS (redundant boot)' --disk "$bdev" --part "$EFILDRRECO_PART_NUM" --loader '\EFI\DSOS\grubx64.efi' || ret=1
    fi

    # Create SpinetiX DSOS menu entry if it is not present.
    if [ "$HAS_DSOS_MENU_BOOT_ENTRY" = "yes" ]; then
        if ! echo "$bootinfo" | grep -q '^Boot[a-fA-F0-9]\+\*\s\+SpinetiX DSOS menu$'; then
            $LOGGER "WARN: Create SpinetiX DSOS menu boot entry"
            efibootmgr -q --create-only --label 'SpinetiX DSOS menu' --disk "$bdev" --part "$EFILDRRECO_PART_NUM" --loader '\EFI\DSOS\menu\grubx64.efi' || ret=1
        fi
    fi

    # Create SpinetiX DSOS entry if it is not present.
    if ! echo "$bootinfo" | grep -q '^Boot[a-fA-F0-9]\+\*\s\+SpinetiX DSOS$'; then
        $LOGGER "WARN: Create SpinetiX DSOS boot entry"
        efibootmgr -q --create --label 'SpinetiX DSOS' --disk "$bdev" --part "$ESP_PART_NUM" --loader '\EFI\DSOS\grubx64.efi' || ret=1
    fi

    return $ret
}

order_efi_boot_entries() {
    local menu_redundant_boot_order
    local menu_boot_order
    local bootinfo

    bootinfo="$(efibootmgr)"
    if [ $? -ne 0 ]; then
        $LOGGER "ERROR: could not query current EFI boot manager data"
        return 1
    fi

    if [ "$HAS_DSOS_MENU_BOOT_ENTRY" = "yes" ]; then
        # Get boot order for different menu entries
        menu_redundant_boot_order=$(echo "$bootinfo" | sed -n -E '/^Boot[a-fA-F0-9]+\*\s+SpinetiX DSOS menu \(redundant boot\)$/ s/^Boot([a-fA-F0-9]+)\*.*/\1/p')
        menu_boot_order=$(echo "$bootinfo" | sed -n -E '/^Boot[a-fA-F0-9]+\*\s+SpinetiX DSOS menu$/ s/^Boot([a-fA-F0-9]+)\*.*/\1/p')
    else
        menu_redundant_boot_order=-
        menu_boot_order=-
    fi

    # Get boot order for different entries
    local dsos_redundant_boot_order=$(echo "$bootinfo" | sed -n -E '/^Boot[a-fA-F0-9]+\*\s+SpinetiX DSOS \(redundant boot\)$/ s/^Boot([a-fA-F0-9]+)\*.*/\1/p')
    local dsos_boot_order=$(echo "$bootinfo" | sed -n -E '/^Boot[a-fA-F0-9]+\*\s+SpinetiX DSOS$/ s/^Boot([a-fA-F0-9]+)\*.*/\1/p')

    if [ -z "$menu_redundant_boot_order" -o -z "$menu_boot_order" -o -z "$dsos_redundant_boot_order" -o -z "$dsos_boot_order" ]; then
        $LOGGER "ERROR: failed to recover order of boot entries"
        return 1
    fi

    # NOTE: some BIOS (e.g., Intel Dawson Canyon NUC) only allow one boot
    # entry per partition and hide any extra entries by wrapping them in a
    # VenHw() descriptor and pushing them to the end of the BootOrder list,
    # keeping only the first boot entry of every partition. So, for third
    # party devices, we only check the first two boot entries to be in the
    # expected place in BootOrder, as they are on different partitions, and
    # for the other ones we only check their presence and relative order in
    # BootOrder.

    local good_boot_order_1 good_boot_order_2 good_boot_order_3 good_boot_order_4

    if [ "$HAS_DSOS_MENU_BOOT_ENTRY" = "yes" ]; then
        good_boot_order_1="$dsos_boot_order"
        good_boot_order_2="$menu_boot_order"
        good_boot_order_3="$dsos_redundant_boot_order"
        good_boot_order_4="$menu_redundant_boot_order"
    else
        good_boot_order_1="$dsos_boot_order"
        good_boot_order_2="$dsos_redundant_boot_order"
        good_boot_order_3=
        good_boot_order_4=
    fi

    local bad_boot_order=
    local current_boot_order=$(echo "$bootinfo" | sed -n -E '/^BootOrder:\s/ s/^[^\s]*:\s+//p')
    if ! [[ "$current_boot_order" =~ ^$good_boot_order_1,$good_boot_order_2(,|$) ]]; then
        bad_boot_order=1
    elif [ -n "$good_boot_order_3" ] && ! [[ "$current_boot_order" =~ ,$good_boot_order_3,$good_boot_order_4(,|$) ]]; then
        bad_boot_order=1
    fi

    if [ -z "$bad_boot_order" ]; then

        rm -f $EFI_BOOT_REORDER_CNT_FILE

    else

        # Construct the boor order so as to preserve existing boot entries in their order, but making sure the first two
        # DSOS entries are always first and in the expected order and the trailing two DSOS entries, if any, appear in
        # the expected relative order at the current position if present or just after the first two DSOS entries if
        # none. This way we add all the DSOS entries in the expected order, while still avoiding changes on every boot
        # when the BIOS hide and push the trailing DSOS entries to the end of the boot order (though it may take two
        # reboots for the list to stabilize when the BIOS hides and moves the trailing DSOS entries).

        # Some BIOS treat the boot entry that is automatically added when there are no valid boot entries, and that
        # points to the default bootloader location in the EFI System Partition, in a special way. These BIOS will
        # insist on making that entry the first entry in the boot order every time the BIOS boots, which conflicts with
        # the boot order we want, and use it to boot. That automatically added boot entry would then be the one that
        # booted this time and thus have its index in BootCurrent. So we skip and remove the BootCurrent entry if it is
        # not one the index of one of the DSOS boot entries. Note that an equivalent entry may be added back by the BIOS
        # later, but it will not be the automatic default and will not be traated specially by the BIOS.

        local cnt boot_current remove_boot_current bentry good_boot_order trailing_added

        [ -f $EFI_BOOT_REORDER_CNT_FILE ] && cnt=$(cat $EFI_BOOT_REORDER_CNT_FILE)
        [ -n "$cnt" ] || cnt=0
        if ! [ "$cnt" -lt $EFI_BOOT_REORDER_MAX ]; then
            $LOGGER "ERROR: unexpected EFI boot order found but alredy reset EFI order $cnt times, aborting"
            echo "$bootinfo" | $LOGGER
            return 1
        fi
        echo $((cnt + 1)) > $EFI_BOOT_REORDER_CNT_FILE

        # loop over each existing entries, keeping them as needed
        boot_current=$(echo "$bootinfo" | sed -n -E '/^BootCurrent:\s/ s/^[^\s]*:\s+//p')
        good_boot_order=
        trailing_added=
        remove_boot_current=
        for bentry in ${current_boot_order//,/ } ; do
            case "$bentry" in
                "$good_boot_order_1"|"$good_boot_order_2")
                    # leading entries done at end
                    ;;
                "$good_boot_order_3"|"$good_boot_order_4")
                    # add trailing entries only once at first match
                    if [ -z "$trailing_added" ]; then
                        good_boot_order="${good_boot_order:+$good_boot_order,}$good_boot_order_3,$good_boot_order_4"
                        trailing_added=1
                    fi
                    ;;
                "$boot_current")
                    # this is the entry that booted us, but is not DSOS, probably the default recovery entry added by the BIOS
                    remove_boot_current=1
                    ;;
                *)
                    # keep current relative position for other entries
                    good_boot_order="${good_boot_order:+$good_boot_order,}$bentry"
                    ;;
            esac
        done

        # prepend the trailing DSOS boot entries, if any and not already added
        if [ -n "$good_boot_order_3" -a -z "$trailing_added" ]; then
            good_boot_order="$good_boot_order_3,$good_boot_order_4${good_boot_order:+,$good_boot_order}"
        fi
        # prepend the first two DSOS boot entries
        good_boot_order="$good_boot_order_1,$good_boot_order_2${good_boot_order:+,$good_boot_order}"

        $LOGGER "WARN: unexpected EFI boot order ($current_boot_order), resetting to $good_boot_order"
        efibootmgr -o $good_boot_order | $LOGGER

        if [ -n "$remove_boot_current" ]; then
            $LOGGER "WARN: removing current boot entry $boot_current as it may conflict"
            efibootmgr -B -b "$boot_current" | $LOGGER
        fi
    fi
}

make_sb_vars_immutable() {
    local make_immutable="$1"
    local attribute

    if [ "$make_immutable" -eq 1 ]; then
        attribute="+i" # Immutable
    else
        attribute="-i" # Mutable
    fi

    if [ -f /sys/firmware/efi/efivars/PK-"$UEFI_GLBL_GUID" ]; then
        chattr $attribute /sys/firmware/efi/efivars/PK-"$UEFI_GLBL_GUID"
    fi
    if [ -f /sys/firmware/efi/efivars/KEK-"$UEFI_GLBL_GUID" ]; then
        chattr $attribute /sys/firmware/efi/efivars/KEK-"$UEFI_GLBL_GUID"
    fi
    if [ -f /sys/firmware/efi/efivars/db-"$UEFI_SECDB_GUID" ]; then
        chattr $attribute /sys/firmware/efi/efivars/db-"$UEFI_SECDB_GUID"
    fi
    if [ -f /sys/firmware/efi/efivars/dbx-"$UEFI_SECDB_GUID" ]; then
        chattr $attribute /sys/firmware/efi/efivars/dbx-"$UEFI_SECDB_GUID"
    fi
    return 0
}

# Writes an EFI variable
write_efi_variable() {
    local file="$1"
    local guid="$2"
    local name="$3"

    # The efivar utility cannot be used because it requires to separate the
    # attributes and format them in decimal, so just use dd and efivarfs directly.
    # Data must be written in a single write call, so use a large block size to
    # ensure that all fits in a single write call.
    dd bs=512k status=none if="$file" of=/sys/firmware/efi/efivars/"$name"-"$guid"
}

update_secure_boot_keys() {
    local initkeysdir="$SBINITKEYSDIR"
    local keystore="$SBKEYSTORE"
    local writepk=

    # NOTE: PK must be written last, even after running sbkeysync, or
    # otherwise there is the risk of exiting setup mode before adding all db
    # entries necessary to boot again with Secure Boot fully enabled.

    # SetupMode is an 8-bit unsigned value, convert it to an hex string.
    if [ "$(hexdump -e '5/1 "%02x"' /sys/firmware/efi/efivars/SetupMode-"$UEFI_GLBL_GUID" | cut -c 9-)" = "01" ]; then
        $LOGGER "WARN: Secure Boot in setup mode, installing Secure Boot keys"

        make_sb_vars_immutable 0

        write_efi_variable "$initkeysdir"/db.siglist.signed "$UEFI_SECDB_GUID" db
        write_efi_variable "$initkeysdir"/dbx.siglist.signed "$UEFI_SECDB_GUID" dbx
        write_efi_variable "$initkeysdir"/KEK.siglist.signed "$UEFI_GLBL_GUID" KEK
        writepk=1

        # NOTE: some of the files for these EFI variables may have not
        # existed and were thus created, but when created the immutable
        # attribute will be automatically set.

        $LOGGER "WARN: Secure Boot db, dbx and KEK keys installed"
    fi

    # NOTE: the keys from $initkeysdir are only a bare minimum set, so as to
    # create the Secure Boot variables, so that sbkeysync can then run (it
    # can only append to existing variables). The full set of keys is in
    # $keystore.

    if [ -f /sys/firmware/efi/efivars/SecureBoot-"$UEFI_GLBL_GUID" ]; then
        if [ "$(hexdump -e '5/1 "%02x"' /sys/firmware/efi/efivars/SecureBoot-"$UEFI_GLBL_GUID" | cut -c 9-)" != "01" ]; then
            $LOGGER "WARN: Secure Boot is disabled"
        fi

        echo -n "sync SB key store... "

        make_sb_vars_immutable 0

        if ! sbkeysync --keystore "$keystore"; then
            $LOGGER "WARN: Secure Boot key database update failed, booting may fail"
        fi
    else
        $LOGGER "WARN: Secure Boot is not supported"
    fi

    if [ -n "$writepk" ]; then
        $LOGGER "WARN: writing Secure Boot PK, will exit setup mode"

        write_efi_variable "$initkeysdir"/PK.siglist.signed "$UEFI_GLBL_GUID" PK

        $LOGGER "WARN: Secure Boot PK installed"
    fi

    make_sb_vars_immutable 1
}

do_checks() {
    if ! is_nfsroot ; then
        create_efi_boot_entries && order_efi_boot_entries
        update_secure_boot_keys
    fi
}

#
# Main
#

case "$1" in
    start)
        echo -n "Doing DSOS bootloader checks... "

        do_checks

        echo "done"

        ;;
    stop)
        ;;
esac
