#!/bin/bash

set -e

DIR="/etc/arch-secure-boot"
KEYSDIR="$DIR/keys"
SBKEYSDIR="/etc/secureboot/keys"
mkdir -p "$DIR" "$KEYSDIR" "$SBKEYSDIR"

[ -f "$DIR/config" ] && source "$DIR/config"

autoesp="$(findmnt -t vfat -o TARGET|grep -iw efi)"
autoesp=${autoesp:-efi}
ESP="${ESP:-$autoesp}"
EFI="${EFI:-/EFI/opensuse-uki}"
KERNEL="${KERNEL:-$(uname -r)}"
NAME="${NAME:-secure-boot-$KERNEL}"
SHELL_NAME="${SHELL_NAME:-UEFI shell}"
SUBVOLUME_ROOT="${SUBVOLUME_ROOT:-root}"
SUBVOLUME_SNAPSHOT="${SUBVOLUME_SNAPSHOT:-/@/.snapshots/%1/snapshot}" # %1 is replaced with snapshot ID

CMDLINE=/etc/kernel/cmdline
[ -f "$CMDLINE" ] || CMDLINE=/proc/cmdline

cmd="$0 $@"
print_config() {
    cat >&2 << EOF

== Command ==
$cmd

== Config ==
ESP=$ESP
EFI=$EFI
KERNEL=$KERNEL
NAME=$NAME
CMDLINE=$CMDLINE
EOF
}
trap 'print_config' ERR

error() {
    echo >&2 "$@"
    exit 1
}

add_entry() {
    local entry="$EFI/$1.efi"
    [ -f "$ESP/$entry" ] || error "Error: EFI images are not generated yet."
    local mount="$(findmnt -n -o SOURCE -T "$ESP")"
    local partition="${mount##*[!0-9]}"
    local loader="${entry//\//\\}"

    if efibootmgr -v | grep -Fq "File($loader)" ; then
        echo "$entry already in boot order"
        return
    fi

    efibootmgr -d "$mount" -p "$partition" -c -l "${entry//\//\\}" -L "$1"
}

case "$1" in
    initial-setup)
        "$0" generate-keys
        "$0" generate-efi
        "$0" add-efi
        "$0" enroll-keys || true
        ;;

    generate-snapshots)
        snapper --no-dbus -t 0 -c root list --disable-used-space --columns number,date,description > "$ESP/snapshots.txt"
        ;;

    generate-efi)
        echo "Generating EFI images..."

        find "$KEYSDIR" -mindepth 1 | read || error "Error: Secure Boot keys are not generated yet."

        tmp="$(mktemp -d)"
        trap 'rm -rf $tmp' EXIT
        cd "$tmp"

        grep -m1 -v "^#" "$CMDLINE" > cmdline
        # clear any current subvols
        sed -i "s|subvol=[^ ]*||g" cmdline
        # systemd-fstab-generator will merge multiple rootflags=
        sed "s|$| rootflags=subvol=$SUBVOLUME_SNAPSHOT|" cmdline >cmdline-subvol
        sed "s|%%NAME%%|$NAME|g; s|%%CMDLINE%%|$(cat cmdline-subvol)|g" "$DIR/recovery.nsh" > recovery.nsh

        cp /usr/share/ovmf/Shell.efi "efi-shell-unsigned.efi"

        objcopy \
            --add-section .osrel=/etc/os-release --change-section-vma .osrel=0x20000 \
            --add-section .cmdline=cmdline --change-section-vma .cmdline=0x30000 \
            --add-section .linux="/usr/lib/modules/$KERNEL/vmlinuz" --change-section-vma .linux=0x40000 \
            --add-section .initrd="/boot/initrd-$KERNEL"  --change-section-vma .initrd=0x3000000 \
            /usr/lib/systemd/boot/efi/linuxx64.efi.stub "$NAME-unsigned.efi"

        objcopy \
            --add-section .osrel=/etc/os-release --change-section-vma .osrel=0x20000 \
            --add-section .linux="/usr/lib/modules/$KERNEL/vmlinuz" --change-section-vma .linux=0x40000 \
            --add-section .initrd="/boot/initrd-$KERNEL" --change-section-vma .initrd=0x3000000 \
            /usr/lib/systemd/boot/efi/linuxx64.efi.stub "$NAME-recovery-unsigned.efi"


        sbsign --key "$KEYSDIR/db.key" --cert "$KEYSDIR/db.crt" --output "efi-shell.efi" "efi-shell-unsigned.efi"
        for flavor in '' '-recovery'; do
            sbsign --key "$KEYSDIR/db.key" --cert "$KEYSDIR/db.crt" --output "${NAME}${flavor}.efi" "${NAME}${flavor}-unsigned.efi"
        done

        mkdir -p "$ESP/$EFI"
        cp recovery.nsh "$ESP"

        cp "efi-shell.efi" "$ESP/$EFI"
        for flavor in '' '-recovery'; do
            cp "${NAME}${flavor}.efi" "$ESP/$EFI"
        done
        ;;

    add-efi)
        echo "Adding boot entries for EFI images..."

        add_entry "efi-shell"
        add_entry "$NAME"
        ;;

    enroll-keys)
        echo "Enrolling Secure Boot keys..."

        find "$SBKEYSDIR" -mindepth 2 | read || error "Error: keys are not generated yet."

        sbkeysync --verbose --pk
        ;;

    generate-keys)
        echo "Generating Secure Boot keys..."

        cd "$KEYSDIR"

        uuidgen > uuid
        read uuid < uuid

        for pair in PK=PK KEK=PK db=KEK; do
            key="${pair%=*}"
            from="${pair#*=}"

            openssl req -new -x509 -newkey rsa:4096 -subj "/CN=SecureBoot $key/" -keyout "$key.key" -out "$key.crt" -days 3650 -nodes -sha256
            openssl x509 -in "$key.crt" -inform PEM -out "$key.der" -outform DER

            sbsiglist --owner "$uuid" --type x509 --output "$key.esl" "$key.der"
            sbvarsign --key "$from.key" --cert "$from.crt" --output "$key.auth" "$key" "$key.esl"

            chmod 0400 "$key".{key,auth}

            mkdir -p "$SBKEYSDIR/$key/"
            cp "$key.auth" "$SBKEYSDIR/$key/"
        done
        ;;

    *)
        error "Usage: $0 <initial-setup|generate-snapshots|generate-efi|add-efi|generate-keys|enroll-keys>"
        ;;
esac
