#!/usr/bin/env bash
# vim:et:ts=2:sts=2:sw=2

# Get absolute repository root (especially when symlinked)
ROOT="$(realpath "$(dirname "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")")")"

# Load accompanying bash library
source "$ROOT/lib/lib.bash"

# Exit now if run as root user
assert-is-root

# Path to USB disk
DISK=""
IS_SECURE="${IS_SECURE:-false}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --secure)
      IS_SECURE=true
      shift
      ;;
    -h | --help)
      HELP=true
      shift
      ;;
    *)
      if [ -z "$DISK" ]; then
        DISK="$1"
      else
        err "Unknown argument: $1"
        exit 1
      fi
      shift
      ;;
  esac
done

if [ -z "$DISK" ] || [ -n "$HELP" ]; then
  say "
  USAGE:
    $(g "${BASH_SOURCE##*/}") [$(c --secure)] $(c /dev/sdx)

  DESCRIPTION:
    $(g "${BASH_SOURCE##*/}") installs and configures bootloaders for the
    TechLit installer USB on $(c /dev/sdx).
  "
  exit 1
fi

# Mountpoint to use
PREFIX="${PREFIX:-"$(mktemp -d)"}"

# partition infix
[[ "$DISK" =~ mmc ]] || [[ "$DISK" =~ nvm ]] && p=p || p=

say "Mounting system (and un-mounting on errors or exit)"
{
  cleanup() {
    say "Un-mounting system"
    run umount -qR "$PREFIX"
  }
  run mount "$DISK$p"2 "$PREFIX"
}

GRUB_MODULES="all_video boot btrfs cat chain configfile echo efifwsetup efinet \
ext2 fat font gettext gfxmenu gfxterm gfxterm_background gzio halt help hfsplus \
iso9660 jpeg keystatus loadenv loopback linux ls lsefi lsefimmap lsefisystab \
lssal memdisk minicmd normal ntfs part_apple part_msdos part_gpt password_pbkdf2 \
png probe reboot regexp search search_fs_uuid search_fs_file search_label sleep \
smbios squash4 test true video xfs zfs zfscrypt zfsinfo cpuid play tpm cryptodisk \
gcry_arcfour gcry_blowfish gcry_camellia gcry_cast5 gcry_crc gcry_des gcry_dsa \
gcry_idea gcry_md4 gcry_md5 gcry_rfc2268 gcry_rijndael gcry_rmd160 gcry_rsa \
gcry_seed gcry_serpent gcry_sha1 gcry_sha256 gcry_sha512 gcry_tiger gcry_twofish \
gcry_whirlpool luks lvm mdraid09 mdraid1x raid5rec raid6rec keylayouts usbms usb_keyboard udf ext2 reiserfs"

SBAT="$ROOT"/roles/global/sbat.csv

install-grub() {
  local arch="$1"
  local disk="$2"
  local secure_args=()

  if [[ "$arch" == "x86_64-efi" ]] && [[ "$IS_SECURE" == "true" ]]; then
    secure_args=(--disable-shim-lock --modules="$GRUB_MODULES" --sbat "$SBAT")
  fi

  say "Installing $(c "$arch") bootloader"
  run grub-install \
    --force --recheck --removable --skip-fs-probe \
    "${secure_args[@]}" \
    --target="$arch" \
    --efi-directory="$PREFIX" $NONVRAM \
    --boot-directory="$PREFIX" \
    "$disk"
}

for arch in i386-pc i386-efi x86_64-efi; do
  # If we are doing secure boot, we must re-install x86_64-efi to get SBAT
  if [[ "$IS_SECURE" == "true" ]] && [[ "$arch" == "x86_64-efi" ]]; then
    install-grub "$arch" "$DISK"
  elif [ ! -d "$PREFIX"/grub/"$arch" ]; then
    install-grub "$arch" "$DISK"
  fi
done

say "Configuring bootloaders"
{
  iso_uuid=$(lsblk "$DISK$p"3 -no UUID)

  run write "$PREFIX/grub/grub.cfg" "
insmod font
insmod keylayouts
insmod part_gpt

set default=0
set timeout=0

menuentry 'TechLit Installer'{
	search --no-floppy --fs-uuid --set=root $iso_uuid
	loopback loop $iso_uuid

  linux (loop)/boot/vmlinuz-linux
	initrd (loop)/boot/initramfs_x86_64.img
	configfile /boot/grub/grub.cfg
}
"
}

# TODO: use FEATURE_FLAGS
if [[ "$IS_SECURE" == "true" ]]; then
  say "Setting up Secure boot"
  {
    EFI="$PREFIX"/EFI/BOOT
    KEYS="$ROOT"/roles/global/secure-keys
    SHIM="$ROOT"/roles/global/shim

    # grub-install --removable overwrites BOOTX64.efi, so move it to grubx64.efi (shim target)
    if [[ -f "$EFI"/BOOTX64.efi ]]; then
      run mv -f "$EFI"/BOOTX64.efi "$EFI"/grubx64.efi
    fi
    run cp "$SHIM"/shimx64.efi "$EFI"/BOOTX64.efi
    run cp "$SHIM"/mmx64.efi "$EFI"/

    say "Signing GRUB..."
    if [[ -f "$EFI"/grubx64.efi ]]; then
      if ! is-signed "$EFI"/grubx64.efi; then
        run sbsign --key "$KEYS"/MOK.key --cert "$KEYS"/MOK.crt --output \
          "$EFI"/grubx64.efi "$EFI"/grubx64.efi
      fi
    fi

    run cp "$KEYS"/MOK.cer "$EFI/ENROLL_THIS.cer"
  }
fi

# Unmount system and quiet exit hook
# shellcheck disable=SC2218
cleanup
cleanup() { :; }
say "Done."
