#!/bin/bash

set -eu
set -o pipefail

g_exitval=1
g_no_efi_is_ok=false
g_cmd=
g_umount_on_exit=
PID_FILE=/run/puavo-manage-efi.pid # Locks this file to ensure only one puavo-manage-efi is running

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

on_exit()
{
  set +e

  if [ -n "${g_umount_on_exit}" ]; then
    umount "${g_umount_on_exit}" || log err "failed to umount '${g_umount_on_exit}'"
  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)
      }
    }
  '
}

efibootmgr_windows() {
  local boot_entry efibootmgr_flag enable_or_disable_msg

  efibootmgr_flag=$1
  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_boot_efi_mounts

  esp_devpath=$(lsblk -n -l -o PATH,PARTTYPE | awk '$2 == "c12a7328-f81f-11d2-ba4b-00a0c93ec93b" {print $1}') || return 1
  if [ -z "${esp_devpath}" ]; then
    log err 'EFI system partition not found'
    return 1
  fi

  esp_boot_efi_mounts=$(awk "-vdev=${esp_devpath}" '$1 == dev && $2 == "/boot/efi" {print $2}' /proc/mounts) || return 1

  # Umount if already mounted to /boot/efi to ensure we control mount options.
  if [ -n "${esp_boot_efi_mounts}" ]; then
    log warning '/boot/efi was already mounted, umounting it'
    umount /boot/efi || return 1
  fi

  mkdir -p /boot/efi || return 1
  mount -onodev,nosuid "${esp_devpath}" /boot/efi || return 1
  echo /boot/efi
}

esp_break_windows() {
  local defused_windows_file tmp_defused_windows_file

  defused_windows_file='/boot/efi/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 '/boot/efi/EFI/Microsoft' ]; then
      log info "cleaning up /boot/efi/EFI/Microsoft that should not exist"
      rm -rf '/boot/efi/EFI/Microsoft' || true
    fi
    return 0
  fi

  if [ ! -d '/boot/efi/EFI/Microsoft' ]; then
    # not doing/logging anything, perhaps we do not have Windows installed
    # and that should be normal
    return 0
  fi

  log info 'defusing Windows...'
  tar -C '/boot/efi/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 '/boot/efi/EFI/Microsoft'; then
    log warning "error cleaning up /boot/efi/EFI/Microsoft"
    return 1
  fi

  log notice 'Windows defused'

  return 0
}

esp_fix_windows() {
  local defused_windows_file tmp_defused_windows_file

  defused_windows_file='/boot/efi/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...'
  tar -C '/boot/efi/EFI' -z -x -f "$defused_windows_file" || {
    log err 'error repriming Windows (tar error)'
    rm -rf '/boot/efi/EFI/Microsoft' || 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
}

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

usage() {
  echo "Usage: $0 [OPTIONS] COMMAND"
  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"
        echo "    enable-windows               enable Windows"
        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 [ $# -ne 1 ]; then
  usage_error "invalid number of arguments ($#), expected 1"
fi

g_cmd=$1
shift

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}"

g_umount_on_exit=$(esp_mount)

case "${g_cmd}" in
  enable-windows)
    efi_enable_windows
    ;;
  disable-windows)
    efi_disable_windows
    ;;
  *)
    usage_error "invalid command '${g_cmd}'"
    ;;
esac

g_exitval=0
