#!/bin/bash

function usage {
  echo "$(basename "$0") [OPTIONS]"
  echo "  -h            shows usage"
  echo "  -g <version>  generate UKI image for specified kernel version"
  echo "  -a            generate UKI images for all available kernels"
}

function check_root {
  [ $EUID -eq 0 ] && return
  echo "dracut-ukify requires root privileges to work" >&2
  exit 1
}

if [[ ${#} -eq 0 ]]; then
   usage
fi

declare -a ukify_global_args=()

ESP_PATH=$(bootctl --print-esp-path)
if [ -z "$ESP_PATH" ]; then
  exit 1
else

  if [ ! -d "${ESP_PATH}/EFI/Linux" ]; then
    if [ -d "${ESP_PATH}/EFI/linux" ]; then
      BOOT_PATH="${ESP_PATH}/EFI/linux"
    fi
  else
    BOOT_PATH="${ESP_PATH}/EFI/Linux"
  fi

  if [ -z "$BOOT_PATH" ]; then
    mkdir -p "${ESP_PATH}/EFI/Linux"
    BOOT_PATH="${ESP_PATH}/EFI/linux"
  fi

fi

ukify_conf="/etc/kernel/uki.conf"
keys_count=0

function check_uki_conf_and_keys_and_gen_keys {
  if [ ! -f "${ukify_conf}" ]; then
    echo "Create ${ukify_conf}"
    cat >"${ukify_conf}" <<EOF
[UKI]
SecureBootPrivateKey=/etc/kernel/secure-boot.key.pem
SecureBootCertificate=/etc/kernel/secure-boot.cert.pem
SignKernel=yes
PCRBanks=sha384,sha512
SBAT="sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
      uki.author.myimage,1,UKI for System,uki.author.myimage,1,https://www.freedesktop.org/software/systemd/man/systemd-stub.html"

[PCRSignature:initrd]
PCRPrivateKey=/etc/kernel/pcr-initrd.key.pem
PCRPublicKey=/etc/kernel/pcr-initrd.pub.pem
Phases=enter-initrd

[PCRSignature:system]
PCRPrivateKey=/etc/kernel/pcr-system.key.pem
PCRPublicKey=/etc/kernel/pcr-system.pub.pem
Phases=enter-initrd:leave-initrd
       enter-initrd:leave-initrd:sysinit
       enter-initrd:leave-initrd:sysinit:ready
EOF
  fi

  echo "Check Secure Boot keys"
  declare -a keys=("/etc/kernel/secure-boot.key.pem" "/etc/kernel/secure-boot.cert.pem" "/etc/kernel/pcr-initrd.key.pem" "/etc/kernel/pcr-initrd.pub.pem" "/etc/kernel/pcr-system.key.pem" "/etc/kernel/pcr-system.pub.pem")

  for i in "${keys[@]}"
  do
    if [ -f "${i}" ]
    then
      echo "${i} exist"
      keys_count=$(($keys_count + 1))
    fi
  done

  if [ $keys_count -lt 6 ]
  then
    for i in "${keys[@]}"
    do
      if [ -f "${i}" ]
      then
        echo "${i} remove"
        rm "${i}"
        keys_count=$(($keys_count - 1))
      fi
    done
  fi
  echo "Keys = $keys_count"

  if [ -f "${ukify_conf}" ] && [ $keys_count == 0 ]
  then
    echo "Generate keys"
    /usr/lib/systemd/ukify genkey --config "${ukify_conf}"
  fi
}

declare -A kernels
update_all=0

while getopts ":hag:xyz" arg; do
  case ${arg} in
    g)
      found=0
      for line in $(pacman -Qql "$OPTARG"); do
        if [[ $line =~ ^/usr/lib/modules/([^/]+)/pkgbase$ ]]; then
          read -r kernel_name < "/${line}"
          kernels["${kernel_name}"]="${BASH_REMATCH[1]}"
          found=1
          break
        fi
      done
      if (( ! found )); then
        echo "Error occurred during '$OPTARG' package traversal" >&2
        exit 1
      fi
      ;;
    a)
      update_all=1
      ;;
    x)
      check_root
      # Trigger some IO on ESP path to be sure it's mounted by autofs if it's the case
      # Otherwise upgrading systemd may cause ESP partition not mounted at the time dracut attempt to write new image
      stat "$ESP_PATH" >/dev/null
      ;;
    y)
      check_root
      while read -r line; do
        if [[ "$line" == 'usr/lib/modules/'+([^/])'/pkgbase' ]]; then
          read -r kernel_name < "/${line}"
          path="$(grep -lE "^${kernel_name}\$" /usr/lib/modules/*/pkgbase)"
          kernel_version=$(basename "${path%/pkgbase}")

          efi_image="$BOOT_PATH/$kernel_name.efi"
          efi_image_fallback="$BOOT_PATH/$kernel_name-fallback.efi"
          efi_image_fallback_terminal="$BOOT_PATH/$kernel_name-fallback-terminal.efi"

          declare -a images=()
          images+=("$efi_image")
          images+=("$efi_image_fallback")
          images+=("$efi_image_fallback_terminal")

          for image in "${images[@]}";
          do
            if [ -f "$image" ]; then
              echo "==> Removing $image..."
              rm -f "$image"
            fi
          done
        fi
      done
      exit 0
      ;;
    z)
      check_root
      while read -r line; do
        if [[ $line =~ ^usr/lib/modules/([^/]+)/pkgbase$ ]]; then
          read -r kernel_name < "/${line}"
          kernels["${kernel_name}"]="${BASH_REMATCH[1]}"
        else
          update_all=1
          break
        fi
      done
      ;;
    h)
      usage
      ;;
    *)
      usage
      ;;
  esac
done

if (( update_all )); then
  for kernel_path in /usr/lib/modules/*; do
    [ -f "$kernel_path/pkgbase" ] || continue

    kernel_version=$(basename "$kernel_path")
    kernel_dir="/usr/lib/modules/$kernel_version"
    pkgbase="$kernel_dir/pkgbase"
    kernel_name=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' "$pkgbase")

    if [ -n "$kernel_name" ]; then
      kernels["$kernel_name"]="$kernel_version"
    fi
  done
fi



function gen_image() {
  check_root
  check_uki_conf_and_keys_and_gen_keys

  kernel_name="$1"
  kernel_version="$2"

  kernel_dir="/usr/lib/modules/$kernel_version"

  if [ -f "/etc/kernel/cmdline" ]; then
    cmdline=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' "/etc/kernel/cmdline")

  elif [ -f "/boot/refind_linux.conf" ]; then
      cmdline=$(head -n 1 < "/boot/refind_linux.conf" | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  elif [ -f "/efi/refind_linux.conf" ]; then
      cmdline=$(head -n 1 < "/efi/refind_linux.conf" | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  elif [ -f "/boot/efi/refind_linux.conf" ]; then
      cmdline=$(head -n 1 < "/boot/efi/refind_linux.conf" | sed -e 's|"Boot with standard options"||g' | awk '{$1=$1;print}' | sed -e 's/"//g')

  else
    cmdline=$(sed -e 's/^[[:space:]]//g' -e 's/[[:space:]]$//g' -e 's/initrd.*$//g' "/proc/cmdline")
  fi

  cmdline_terminal="$cmdline systemd.unit=multi-user.target"

  vmlinuz="$kernel_dir/vmlinuz"

  initrd_image="$(mktemp)"
  initrd_image_fallback="$(mktemp)"

  efi_image="$BOOT_PATH/$kernel_name.efi"
  efi_image_fallback="$BOOT_PATH/$kernel_name-fallback.efi"
  efi_image_fallback_terminal="$BOOT_PATH/$kernel_name-fallback-terminal.efi"

  echo "==> Building initrd image $kernel_name $initrd_image"
#   dracut --force --hostonly --no-hostonly-cmdline --kver "$kernel_version" "$initrd_image"
  dracut --force --hostonly --no-hostonly-i18n --early-microcode --kernel-cmdline="$cmdline" --kver "$kernel_version" "$initrd_image"

  declare -a ukify_args=("${ukify_global_args[@]}")
  ukify_args+=(--linux "$vmlinuz")
  if [ -f "/boot/amd-ucode.img" ];then
    ukify_args+=(--initrd "/boot/amd-ucode.img")
  fi
  if [ -f "/boot/intel-ucode.img" ];then
    ukify_args+=(--initrd "/boot/intel-ucode.img")
  fi
  ukify_args+=(--initrd "$initrd_image")
  ukify_args+=(--cmdline "$cmdline")
  ukify_args+=(--os-release "/etc/os-release")
  ukify_args+=(--uname "$kernel_version")
  ukify_args+=(--output "$efi_image")
  ukify_args+=(--measure)

  echo "==> Ukify image $kernel_name $efi_image"
  if [ -f "${ukify_conf}" ]; then
#     ukify_args+=(--sbat='sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
#       uki.author.myimage,1,UKI for System,uki.author.myimage,1,https://www.freedesktop.org/software/systemd/man/systemd-stub.html')
    /usr/lib/systemd/ukify --config "${ukify_conf}" build "${ukify_args[@]}"
  else
    /usr/lib/systemd/ukify build "${ukify_args[@]}"
  fi

  echo "==> Building initrd image $kernel_name $initrd_image_fallback"
#   dracut --force --kver "$kernel_version" "$initrd_image_fallback"
  dracut --force --no-hostonly --no-hostonly-i18n --early-microcode --kernel-cmdline="$cmdline" --kver "$kernel_version" "$initrd_image_fallback"

  declare -a ukify_args=("${ukify_global_args[@]}")
  ukify_args+=(--linux "$vmlinuz")
  if [ -f "/boot/amd-ucode.img" ];then
    ukify_args+=(--initrd "/boot/amd-ucode.img")
  fi
  if [ -f "/boot/intel-ucode.img" ];then
    ukify_args+=(--initrd "/boot/intel-ucode.img")
  fi
  ukify_args+=(--initrd "$initrd_image_fallback")
  ukify_args+=(--cmdline "$cmdline")
  ukify_args+=(--os-release "/etc/os-release")
  ukify_args+=(--uname "$kernel_version")
  ukify_args+=(--output "$efi_image_fallback")
  ukify_args+=(--measure)

  echo "==> Ukify image $kernel_name $efi_image_fallback"
  if [ -f "${ukify_conf}" ]; then
#     ukify_args+=(--sbat='sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
#       uki.author.myimage,1,UKI for System,uki.author.myimage,1,https://www.freedesktop.org/software/systemd/man/systemd-stub.html')
    /usr/lib/systemd/ukify --config "${ukify_conf}" build "${ukify_args[@]}"
  else
    /usr/lib/systemd/ukify build "${ukify_args[@]}"
  fi

  echo "==> Building initrd image $kernel_name $initrd_image_fallback TERMINAL"
#   dracut --force --kver "$kernel_version" "$initrd_image_fallback"
  dracut --force --no-hostonly --no-hostonly-i18n --early-microcode --kernel-cmdline="$cmdline_terminal" --kver "$kernel_version" "$initrd_image_fallback"

  declare -a ukify_args=("${ukify_global_args[@]}")
  ukify_args+=(--linux "$vmlinuz")
  if [ -f "/boot/amd-ucode.img" ];then
    ukify_args+=(--initrd "/boot/amd-ucode.img")
  fi
  if [ -f "/boot/intel-ucode.img" ];then
    ukify_args+=(--initrd "/boot/intel-ucode.img")
  fi
  ukify_args+=(--initrd "$initrd_image_fallback")
  ukify_args+=(--cmdline "$cmdline_terminal")
  ukify_args+=(--os-release "/etc/os-release")
  ukify_args+=(--uname "$kernel_version")
  ukify_args+=(--output "$efi_image_fallback_terminal")
  ukify_args+=(--measure)

  echo "==> Ukify image $kernel_name $efi_image_fallback_terminal TERMINAL"
  if [ -f "${ukify_conf}" ]; then
#     ukify_args+=(--sbat='sbat,1,SBAT Version,sbat,1,https://github.com/rhboot/shim/blob/main/SBAT.md
#       uki.author.myimage,1,UKI for System,uki.author.myimage,1,https://www.freedesktop.org/software/systemd/man/systemd-stub.html')
    /usr/lib/systemd/ukify --config "${ukify_conf}" build "${ukify_args[@]}"
  else
    /usr/lib/systemd/ukify build "${ukify_args[@]}"
  fi

  rm -f "$initrd_image"
  rm -f "$initrd_image_fallback"
}

for kernel_name in "${!kernels[@]}"; do
  kernel_version="${kernels[$kernel_name]}"
  gen_image "$kernel_name" "$kernel_version"
done