#!/bin/bash

set -eu
set -o pipefail

g_exitval=1
g_no_efi_is_ok=false
g_cmd=
g_esp_mount_path=
g_esp_devpath=
g_win_c_partition_devpath=
g_win_disk=
PID_FILE=/run/puavo-manage-efi.pid # Locks this file to ensure only one puavo-manage-efi is running
ESP_WINDOWS_DIR_SUFFIX='/EFI/Microsoft'

log() {
  logger -t puavo-manage-efi -p "user.${1}" "$2" || true
}

on_exit()
{
  set +e

  if [ -n "${g_esp_mount_path}" ]; then
    umount "${g_esp_mount_path}" || log err "failed to umount '${g_esp_mount_path}'"
  fi

  rm -f "${PID_FILE}"

  if [ ${g_exitval} -ne 0 ]; then
    log err 'failed!'
  fi

  exit ${g_exitval}
}

get_windows_boot_entries_to_change() {
  efibootmgr | awk -v flag="$1" '
    /Windows/ {
      is_active = (substr($0, 9, 1) == "*")
      if ((flag == "-A" && is_active) || flag == "-a" && !is_active) {
        print substr($0, 5, 4)
      }
    }
  '
}

get_esp_devpaths()
{
  lsblk -n -l -o PATH,PARTTYPE -- "$@" | awk '$2 == "c12a7328-f81f-11d2-ba4b-00a0c93ec93b" {print $1}'
}

get_one_esp_devpath()
{
  local esp_devpath esp_devpath_count line

  esp_devpath_count=0

  while read -r line; do
    if [ -z "${line}" ]; then
      continue
    fi
    esp_devpath="${line}"
    log notice "Found ESP at '${esp_devpath}'"
    esp_devpath_count=$((esp_devpath_count + 1))
  done < <(get_esp_devpaths "$@") || {
    log err 'failed to find ESP'
    return 1
  }

  if [ ${esp_devpath_count} -eq 0 ]; then
    log err 'failed to find ESP'
    return 1
  elif [ ${esp_devpath_count} -gt 1 ]; then
    log err 'found multiple ESPs'
    return 1
  fi

  echo -n "${esp_devpath}"
}

get_partnum()
{
  local devpath partnum

  devpath="$1"
  shift

  partnum=$(stat --format %Lr "${devpath}") || {
    log err "failed to get the partition number of '${devpath}'"
    return 1
  }

  if [ "$partnum" -lt 1 ]; then
    log err "file '${devpath}' is not a block device"
    return 1
  fi

  echo -n "$partnum"
}

get_disk()
{
  local part_devpath disk_devpath

  part_devpath=$1
  shift

  disk_devpath=$(lsblk -p -J -s "${part_devpath}" | jq -r '.blockdevices[0].children[0] | select(.type == "disk").name') || {
    log err "failed to find disk for '${part_devpath}'"
    return 1
  }

  if [ -z "${disk_devpath}" ]; then
    log err 'failed to find disk (jq empty output)'
    return 1
  fi

  echo -n "${disk_devpath}"
}

efibootmgr_windows_create() {
  local esp_partnum esp_partuuid

  log notice "creating EFI boot entry for Windows Boot Manager in ESP '${g_esp_devpath}'"

  esp_partnum=$(get_partnum "${g_esp_devpath}") || return 1

  esp_partuuid=$(lsblk -n -oPARTUUID "${g_esp_devpath}") || {
    log err "failed to query GPT partuuid of ESP '${g_esp_devpath}' (lsblk failed)"
    return 1
  }

  if ! [[ "${esp_partuuid}" =~ ^[A-Fa-f0-9]{8}(-[A-Fa-f0-9]{4}){3}-[A-Fa-f0-9]{12}$ ]]; then
    log err "failed to query GPT partuuid of ESP '${g_esp_devpath}' (lsblk printed invalid UUID: '${esp_partuuid}')"
    return 1
  fi

  if efibootmgr -v | sed -r -n 's|^Boot[0-9]+[* ] [^\t]+\tHD\([0-9]+,GPT,([^,]+),[^)]+\)/File\(\\EFI\\Microsoft\\Boot\\bootmgfw\.efi\).*$|\1|p' | grep -i -q -x "${esp_partuuid}"; then
    log notice "not creating EFI boot entry for Windows Boot Manager in ESP '${g_esp_devpath}', because at least one already exists"
    return 0
  fi

  efibootmgr -q -c -L 'Windows Boot Manager' -l '\EFI\Microsoft\Boot\bootmgfw.efi' -d "${g_win_disk}" -p "${esp_partnum}" || {
    log err "failed to create EFI boot entry for Windows Boot Manager in ESP '${g_esp_devpath}'"
    return 1
  }

  log notice "created EFI boot entry for Windows Boot Manager in ESP '${g_esp_devpath}'"
}

efibootmgr_windows() {
  local boot_entry efibootmgr_flag enable_or_disable_msg

  efibootmgr_flag=$1
  shift
  case "$efibootmgr_flag" in
    -A) enable_or_disable_msg='disabling' ;;
    -a) enable_or_disable_msg='enabling'  ;;
     *) return 1;;
  esac


  for boot_entry in $(get_windows_boot_entries_to_change "$efibootmgr_flag"); do
    log notice "${enable_or_disable_msg} Windows boot entry ${boot_entry}"
    efibootmgr -q "$efibootmgr_flag" -b "$boot_entry" || return 1
  done

  return 0
}

esp_mount() {
  local esp_devpath esp_or_boot_efi_mount

  esp_devpath="$1"
  shift

  # Umount if already mounted to ensure we control mount options.
  while IFS= read -r -d $'\0' esp_or_boot_efi_mount; do
    log warning "'${esp_or_boot_efi_mount}' was already mounted, umounting it"
    umount "${esp_or_boot_efi_mount}" || return 1
  done < <(awk "-vdev=${esp_devpath}" '$1 == dev && $2 == "/boot/efi" { printf "%s", $2 }' /proc/mounts) || return 1

  mkdir -p /boot/efi || return 1
  mount -onodev,nosuid "${esp_devpath}" /boot/efi 1>/dev/null || return 1
  log info "mounted ESP '${esp_devpath}' to /boot/efi"
  echo -n /boot/efi
}

esp_break_windows() {
  local defused_windows_file esp_windows_dir tmp_defused_windows_file

  esp_windows_dir="${g_esp_mount_path}${ESP_WINDOWS_DIR_SUFFIX}"
  defused_windows_file="${g_esp_mount_path}/EFI/Puavo-defused-Windows.tar.gz"
  tmp_defused_windows_file="${defused_windows_file}.tmp"
  rm -f "$tmp_defused_windows_file"

  if [ -e "$defused_windows_file" ]; then
    # nothing to do
    if [ -d "${esp_windows_dir}" ]; then
      log info "cleaning up ${esp_windows_dir} that should not exist"
      rm -rf "${esp_windows_dir}" || true
    fi
    return 0
  fi

  if [ ! -d "${esp_windows_dir}" ]; then
    # not doing/logging anything, perhaps we do not have Windows installed
    # and that should be normal
    return 0
  fi

  log info "defusing Windows in ESP '${g_esp_devpath}' ..."
  tar -C "${g_esp_mount_path}/EFI" -z -c -f "$tmp_defused_windows_file" Microsoft || {
    log err 'failed to defuse Windows (tar error)'
    rm -f "$tmp_defused_windows_file" || true
    return 1
  }
  sync || true
  if ! mv "$tmp_defused_windows_file" "$defused_windows_file"; then
    log err 'failed to defuse Windows (mv error)'
    return 1
  fi
  if ! rm -rf "${esp_windows_dir}"; then
    log warning "error cleaning up ${esp_windows_dir}"
    return 1
  fi

  log notice 'Windows defused'

  return 0
}

esp_install_windows() {
  local defused_windows_file esp_windows_dir tmp_windows_dir

  esp_windows_dir="${g_esp_mount_path}${ESP_WINDOWS_DIR_SUFFIX}"
  defused_windows_file="${g_esp_mount_path}/EFI/Puavo-defused-Windows.tar.gz"
  tmp_windows_dir="${esp_windows_dir}.tmp"
  rm -rf "${tmp_windows_dir}"

  if [ -e "${defused_windows_file}" ]; then
    log err 'ESP already contains a defused Windows'
    return 1
  fi

  if [ -e "${esp_windows_dir}" ]; then
    log err 'ESP already contains Windows'
    return 1
  fi

  log info 'installing Windows bootloader...'
  mkdir -p "${tmp_windows_dir}" || {
    log err 'failed to install Windows bootloader (mkdir)'
    return 1
  }
  puavo-reset-windows --mode install-boot-dir "${g_win_c_partition_devpath}" "${tmp_windows_dir}" || {
    log err 'failed to install Windows bootloader (puavo-reset-windows)'
    rm -rf "${tmp_windows_dir}" || true
    return 1
  }
  mv -n -T "${tmp_windows_dir}" "${esp_windows_dir}" || {
    log err 'failed to install Windows bootloader (mv)'
    rm -rf "${tmp_windows_dir}" || true
    return 1
  }
  if [ -d "${tmp_windows_dir}" ]; then
    # mv failed because the source dir still exists
    log err 'failed to install Windows bootloader (mv)'
    rm -rf "${esp_windows_dir}" || true
    rm -rf "${tmp_windows_dir}" || true
    return 1
  fi

  sync || true

  log info 'Windows bootloader installed'
}

esp_fix_windows() {
  local defused_windows_file tmp_defused_windows_file

  esp_windows_dir="${g_esp_mount_path}${ESP_WINDOWS_DIR_SUFFIX}"
  defused_windows_file="${g_esp_mount_path}/EFI/Puavo-defused-Windows.tar.gz"
  tmp_defused_windows_file="${defused_windows_file}.tmp"
  rm -f "$tmp_defused_windows_file"

  if [ ! -e "$defused_windows_file" ]; then
    # nothing to do
    return 0
  fi

  log info "repriming Windows in ESP '${g_esp_devpath}' ..."
  tar -C "${g_esp_mount_path}/EFI" -z -x -f "$defused_windows_file" || {
    log err 'error repriming Windows (tar error)'
    rm -rf "${esp_windows_dir}" || true
    return 1
  }
  sync || true
  rm -f "$defused_windows_file" || return 1

  log info 'Windows reprimed'
}

efi_disable_windows() {
  if [ "$(puavo-conf puavo.grub.windows.defuse_efi_boot_when_disabled)" != 'true' ]; then
    return 0
  fi

  efibootmgr_windows -A || return 1
  esp_break_windows
}

efi_enable_windows() {
  esp_fix_windows || return 1
  efibootmgr_windows -a
}

efi_install_windows() {
  esp_install_windows || return 1
  efibootmgr_windows_create
}

has_efi() {
  grep -q '^efivarfs ' /proc/mounts
}

usage() {
  echo "Usage: $0 [OPTIONS] disable-windows WIN_DISK_DEVPATH"
  echo "       $0 [OPTIONS] enabled-windows WIN_DISK_DEVPATH"
  echo "       $0 [OPTIONS] install-windows WIN_C_PARTITION_DEVPATH"
  echo "       $0 --help"
}

usage_error() {
  local msg

  msg=$1
  shift

  echo "ERROR: ${msg}" >&2
  usage >&2

  exit 1
}

while [ $# -gt 0 ]; do
  case $1 in
    -h|--help)
      shift
      {
        usage
        echo
        echo "Manage EFI boot. Modifies EFI boot order and ESP."
        echo
        echo "Commands:"
        echo "    disable-windows              disable Windows installed on WIN_DISK_DEVPATH"
        echo "    enable-windows               enable Windows installed on WIN_DISK_DEVPATH"
        echo "    install-windows              install Windows boot manager to ESP, lookup Windows version from WIN_C_PARTITION_DEVPATH"
        echo
        echo "Options:"
        echo "    --no-efi-is-ok               exit with status 0 if the system does not use EFI"
        echo "    -h, --help                   print help and exit"
        echo
      } >&2
      exit 0
      ;;
    --no-efi-is-ok)
      shift
      g_no_efi_is_ok=true
      ;;
    --)
      shift
      break
      ;;
    -*)
      usage_error "invalid argument '$1'"
      ;;
    *)
      break
      ;;
  esac
done

if [ $# -lt 1 ]; then
  usage_error 'COMMAND is missing'
fi

g_cmd=$1
shift

case "${g_cmd}" in
  'install-windows')
    if [ $# -ne 1 ]; then
      usage_error 'WIN_C_PARTITION_DEVPATH is missing'
    fi
    g_win_c_partition_devpath=$1
    shift
    ;;
  'enable-windows'|'disable-windows')
    if [ $# -ne 1 ]; then
      usage_error 'WIN_DISK_DEVPATH is missing'
    fi
    g_win_disk=$1
    shift
    ;;
  *)
    usage_error "invalid command '${g_cmd}'"
    ;;
esac

if [ $# -ne 0 ]; then
  usage_error 'too many arguments'
fi

if ! has_efi; then
  if ${g_no_efi_is_ok}; then
    log warning "this system does not have EFI"
    exit 0
  else
    log err "this system does not have EFI"
    exit 1
  fi
fi

if [ "$(id -u)" -ne 0 ]; then
  echo 'You need to be root!' >&2
  exit 1
fi

exec {lockfd}<> "${PID_FILE}"
flock -x -n "${lockfd}" || {
  log err "puavo-manage-efi is already running!"
  exit 1
}

trap on_exit EXIT

echo "$$" >"${PID_FILE}"

case "${g_cmd}" in
  enable-windows)
    g_esp_devpath=$(get_one_esp_devpath "${g_win_disk}")
    g_esp_mount_path=$(esp_mount "${g_esp_devpath}")
    efi_enable_windows
    ;;
  disable-windows)
    g_esp_devpath=$(get_one_esp_devpath "${g_win_disk}")
    g_esp_mount_path=$(esp_mount "${g_esp_devpath}")
    efi_disable_windows
    ;;
  install-windows)
    g_win_disk=$(get_disk "${g_win_c_partition_devpath}")
    g_esp_devpath=$(get_one_esp_devpath "${g_win_disk}")
    g_esp_mount_path=$(esp_mount "${g_esp_devpath}")
    efi_install_windows
    ;;
  *)
    usage_error "invalid command '${g_cmd}'"
    ;;
esac

g_exitval=0
