#!/bin/sh
# Boot from a local disk installation

BOARD_NAME=${CONFIG_BOARD_NAME:-${CONFIG_BOARD}} 
MAIN_MENU_TITLE="${BOARD_NAME} | Heads Boot Menu"
export BG_COLOR_WARNING="${CONFIG_WARNING_BG_COLOR:-"--background-gradient 0 0 0 150 125 0"}"
export BG_COLOR_ERROR="${CONFIG_ERROR_BG_COLOR:-"--background-gradient 0 0 0 150 0 0"}"
export BG_COLOR_MAIN_MENU=""

. /etc/functions
. /tmp/config

last_half=X

mount_boot()
{
  
  # Mount local disk if it is not already mounted
  while ! grep -q /boot /proc/mounts ; do
    # try to mount if CONFIG_BOOT_DEV exists
    if [ -e "$CONFIG_BOOT_DEV" ]; then
      mount -o ro $CONFIG_BOOT_DEV /boot 
      [[ $? -eq 0 ]] && continue
    fi

    # CONFIG_BOOT_DEV doesn't exist or couldn't be mounted, so give user options
    BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
    whiptail $BG_COLOR_ERROR --clear --title "ERROR: No Bootable OS Found!" \
        --menu "    No bootable OS was found on the default boot device $CONFIG_BOOT_DEV.
    How would you like to proceed?" 30 90 4 \
        'b' ' Select a new boot device' \
        'u' ' Boot from USB' \
        'm' ' Continue to the main menu' \
        'x' ' Exit to recovery shell' \
        2>/tmp/whiptail || recovery "GUI menu failed"

    option=$(cat /tmp/whiptail)
    case "$option" in 
      b )
        config-gui.sh boot_device_select
        if [ $? -eq 0 ]; then
          # update CONFIG_BOOT_DEV
          . /tmp/config
          BG_COLOR_MAIN_MENU=""
        fi
        ;;
      u )
        exec /bin/usb-init
        ;;
      m )
        skip_to_menu="true"
        break
        ;;
      * )
        recovery "User requested recovery shell"
        ;;
    esac
  done
}

verify_global_hashes()
{
  # Check the hashes of all the files, ignoring signatures for now
  check_config /boot force
  TMP_HASH_FILE="/tmp/kexec/kexec_hashes.txt"
  TMP_PACKAGE_TRIGGER_PRE="/tmp/kexec/kexec_package_trigger_pre.txt"
  TMP_PACKAGE_TRIGGER_POST="/tmp/kexec/kexec_package_trigger_post.txt"

  if ( cd /boot && sha256sum -c "$TMP_HASH_FILE" > /tmp/hash_output ) then
    return 0
  elif [ ! -f $TMP_HASH_FILE ]; then
    if (whiptail $BG_COLOR_ERROR --clear --title 'ERROR: Missing Hash File!' \
      --yesno "The file containing hashes for /boot is missing!\n\nIf you are setting this system up for the first time, select Yes to update\nyour list of checksums.\n\nOtherwise this could indicate a compromise and you should select No to\nreturn to the main menu.\n\nWould you like to update your checksums now?" 30 90) then
      update_checksums
    fi
    BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
    return 1
  else
    CHANGED_FILES=$(grep -v 'OK$' /tmp/hash_output | cut -f1 -d ':')

    # if files changed before package manager started, show stern warning
    if [ -f "$TMP_PACKAGE_TRIGGER_PRE" ]; then 
      PRE_CHANGED_FILES=$(grep '^CHANGED_FILES' $TMP_PACKAGE_TRIGGER_POST | cut -f 2 -d '=' | tr -d '"')
      TEXT="The following files failed the verification process BEFORE package updates ran:\n${PRE_CHANGED_FILES}\n\nCompare against the files Heads has detected have changed:\n${CHANGED_FILES}\n\nThis could indicate a compromise!\n\nWould you like to update your checksums anyway?"

    # if files changed after package manager started, probably caused by package manager
    elif [ -f "$TMP_PACKAGE_TRIGGER_POST" ]; then 
      LAST_PACKAGE_LIST=$(grep -E "^(Install|Remove|Upgrade|Reinstall):" $TMP_PACKAGE_TRIGGER_POST)
      UPDATE_INITRAMFS_PACKAGE=$(grep '^UPDATE_INITRAMFS_PACKAGE' $TMP_PACKAGE_TRIGGER_POST | cut -f 2 -d '=' | tr -d '"')

      if [ "$UPDATE_INITRAMFS_PACKAGE" != "" ]; then
        TEXT="The following files failed the verification process AFTER package updates ran:\n${CHANGED_FILES}\n\nThis is likely due to package triggers in$UPDATE_INITRAMFS_PACKAGE.\n\nYou will need to update your checksums for all files in /boot.\n\nWould you like to update your checksums now?"
      else
        TEXT="The following files failed the verification process AFTER package updates ran:\n${CHANGED_FILES}\n\nThis might be due to the following package updates:\n$LAST_PACKAGE_LIST.\n\nYou will need to update your checksums for all files in /boot.\n\nWould you like to update your checksums now?"
      fi

    else
      TEXT="The following files failed the verification process:\n\n${CHANGED_FILES}\n\nThis could indicate a compromise!\n\nWould you like to update your checksums now?"
    fi

    if (whiptail $BG_COLOR_ERROR --clear --title 'ERROR: Boot Hash Mismatch' --yesno "$TEXT" 30 90) then
      update_checksums
    fi
    BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
    return 1
  fi
}

prompt_update_checksums()
{
  if (whiptail $BG_COLOR_WARNING --title 'Update Checksums and sign all files in /boot' \
      --yesno "You have chosen to update the checksums and sign all of the files in /boot.\n\nThis means that you trust that these files have not been tampered with.\n\nYou will need your GPG key available, and this change will modify your disk.\n\nDo you want to continue?" 16 90) then
    update_checksums
  else
    echo "Returning to the main menu"
  fi
}

generate_totp_htop()
{
  echo "Scan the QR code to add the new TOTP secret"
  /bin/seal-totp
  if [ -x /bin/hotp_verification ]; then
    echo "Once you have scanned the QR code, hit Enter to configure your HOTP USB Security Dongle (e.g. Librem Key or Nitrokey)"
    read
    /bin/seal-hotpkey
  else
    echo "Once you have scanned the QR code, hit Enter to continue"
    read
  fi
}

update_totp()
{
  # update the TOTP code every thirty seconds
  date=`date "+%Y-%m-%d %H:%M:%S"`
  seconds=`date "+%s"`
  half=`expr \( $seconds % 60 \) / 30`
  if [ "$CONFIG_TPM" = n ]; then
    TOTP="NO TPM"
  elif [ "$half" != "$last_half" ]; then
    last_half=$half;
    TOTP=`unseal-totp`
    if [ $? -ne 0 ]; then
      BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
      whiptail $BG_COLOR_ERROR --clear --title "ERROR: TOTP Generation Failed!" \
        --menu "    ERROR: Heads couldn't generate the TOTP code.\n
  If you have just completed a Factory Reset, or just reflashed
  your BIOS, you should generate a new HOTP/TOTP secret.\n
  If this is the first time the system has booted, you should
  reset the TPM and set your own password.\n
  If you have not just reflashed your BIOS, THIS COULD INDICATE TAMPERING!\n
  How would you like to proceed?" 30 90 4 \
        'g' ' Generate new HOTP/TOTP secret' \
        'i' ' Ignore error and continue to main menu' \
        'p' ' Reset the TPM' \
        'x' ' Exit to recovery shell' \
        2>/tmp/whiptail || recovery "GUI menu failed"

      option=$(cat /tmp/whiptail)
      case "$option" in 
        g )
          if (whiptail $BG_COLOR_WARNING --title 'Generate new TOTP/HOTP secret' \
              --yesno "This will erase your old secret and replace it with a new one!\n\nDo you want to proceed?" 16 90) then
            generate_totp_htop && BG_COLOR_MAIN_MENU=""
          fi
          ;;
        i )
          return 1
          ;;
        p )
          reset_tpm && BG_COLOR_MAIN_MENU=""
          ;;
        x )
          recovery "User requested recovery shell"
          ;;
      esac
    fi
  fi
}

update_hotp()
{
  if [ -x /bin/hotp_verification ]; then
    HOTP=`unseal-hotp`
    enable_usb
    if ! hotp_verification info ; then
      whiptail $BG_COLOR_WARNING --clear \
--title "WARNING: Please Insert Your $HOTPKEY_BRANDING" \
--msgbox "Your $HOTPKEY_BRANDING was not detected.\n\nPlease insert your $HOTPKEY_BRANDING" 30 90
      fi
    # Don't output HOTP codes to screen, so as to make replay attacks harder
    hotp_verification check $HOTP
    case "$?" in
      0 )
        HOTP="Success"
        BG_COLOR_MAIN_MENU=""
      ;;
      4 )
        HOTP="Invalid code"
        BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
      ;;
      * )
        HOTP="Error checking code, Insert $HOTPKEY_BRANDING and retry"
        BG_COLOR_MAIN_MENU=$BG_COLOR_WARNING
      ;;
    esac
  else
    HOTP='N/A'
  fi
}

clean_boot_check()
{
  # assume /boot mounted
  if ! grep -q /boot /proc/mounts ; then
    return
  fi

  # check for any kexec files in /boot
  kexec_files=`find /boot -name kexec*.txt`
  [ ! -z "$kexec_files" ] && return
  
  #check for GPG key in keyring
  GPG_KEY_COUNT=`gpg -k 2>/dev/null | wc -l`
  [ $GPG_KEY_COUNT -ne 0 ] && return

   # check for USB security token
  if [ "$CONFIG_HOTPKEY" = "y" ]; then
    enable_usb
    if ! gpg --card-status > /dev/null ; then
      return
    fi
  fi

  # OS is installed, no kexec files present, no GPG keys in keyring, security token present
  # prompt user to run OEM factory reset
  oem-factory-reset \
    "Clean Boot Detected - Perform OEM Factory Reset?" "$BG_COLOR_WARNING"
}

check_gpg_key()
{
  GPG_KEY_COUNT=`gpg -k 2>/dev/null | wc -l`
  if [ $GPG_KEY_COUNT -eq 0 ]; then
    BG_COLOR_MAIN_MENU=$BG_COLOR_ERROR
    whiptail $BG_COLOR_ERROR --clear --title "ERROR: GPG keyring empty!" \
      --menu "ERROR: Heads couldn't find any GPG keys in your keyring.\n\nIf this is the first time the system has booted,\nyou should add a public GPG key to the BIOS now.\n\nIf you just reflashed a new BIOS, you'll need to add at least one\npublic key to the keyring.\n\nIf you have not just reflashed your BIOS, THIS COULD INDICATE TAMPERING!\n\nHow would you like to proceed?" 30 90 4 \
      'g' ' Add a GPG key to the running BIOS' \
      'i' ' Ignore error and continue to main menu' \
      'x' ' Exit to recovery shell' \
      2>/tmp/whiptail || recovery "GUI menu failed"

    option=$(cat /tmp/whiptail)
    case "$option" in 
      g )
        gpg-gui.sh && BG_COLOR_MAIN_MENU=""
        ;;
      i )
        return 1
        ;;
      x )
        recovery "User requested recovery shell"
        ;;
    esac
  fi
}

prompt_auto_default_boot()
{
  # save IFS before changing, restore after read
  IFS_DEF=$IFS
  IFS=''
  first_pass=false
  echo -e "\nHOTP verification success\n\n"
  read -t $CONFIG_AUTO_BOOT_TIMEOUT -s -n 1 -p "Automatic boot in $CONFIG_AUTO_BOOT_TIMEOUT seconds unless interrupted by keypress... "
  if [[ $? -ne 0 ]]; then
    IFS=$IFS_DEF
    echo -e "\n\nAttempting default boot...\n\n"
    attempt_default_boot
  fi
  IFS=$IFS_DEF
}

show_main_menu()
{
  date=`date "+%Y-%m-%d %H:%M"`
  whiptail $BG_COLOR_MAIN_MENU --clear --title "$MAIN_MENU_TITLE" \
    --menu "$date\nTOTP: $TOTP | HOTP: $HOTP" 20 90 10 \
    'd' ' Default boot' \
    'r' ' Refresh TOTP/HOTP' \
    'o' ' Options -->' \
    's' ' System Info' \
    'p' ' Power Off' \
    2>/tmp/whiptail || recovery "GUI menu failed"

  option=$(cat /tmp/whiptail)
  case "$option" in 
    d )
      attempt_default_boot
      ;;
    r )
      update_totp && update_hotp
      ;;
    o )
      show_options_menu
      ;;
    s )
      show_system_info
      ;;
    p )
      poweroff
      ;;
  esac
}

show_options_menu()
{
  whiptail $BG_COLOR_MAIN_MENU --clear --title "HEADS Options" \
    --menu "" 20 90 10 \
    'b' ' Boot Options -->' \
    't' ' TPM/TOTP/HOTP Options -->' \
    'u' ' Update checksums and sign all files in /boot' \
    'c' ' Change configuration settings -->' \
    'f' ' Flash/Update the BIOS -->' \
    'g' ' GPG Options -->' \
    'F' ' OEM Factory Reset -->' \
    'x' ' Exit to recovery shell' \
    'r' ' <-- Return to main menu' \
    2>/tmp/whiptail || recovery "GUI menu failed"

  option=$(cat /tmp/whiptail)
  case "$option" in 
    b )
      show_boot_options_menu
      ;;
    t )
      show_tpm_totp_hotp_options_menu
      ;;
    u )
      prompt_update_checksums
      ;;
    c )
      config-gui.sh
      ;;
    f )
      flash-gui.sh
      ;;
    g )
      gpg-gui.sh
      ;;
    F )
      oem-factory-reset
      ;;
    x )
      recovery "User requested recovery shell"
      ;;
    r )
      ;;
  esac
}

show_boot_options_menu()
{
  whiptail $BG_COLOR_MAIN_MENU --clear --title "Boot Options" \
    --menu "Select A Boot Option" 20 90 10 \
    'm' ' Show OS boot menu' \
    'u' ' USB boot' \
    'i' ' Ignore tampering and force a boot (Unsafe!)' \
    'r' ' <-- Return to main menu' \
    2>/tmp/whiptail || recovery "GUI menu failed"

  option=$(cat /tmp/whiptail)
  case "$option" in 
    m )
      # select a kernel from the menu
      select_os_boot_option
      ;;
    u )
      exec /bin/usb-init
      ;;
    i )
      force_unsafe_boot
      ;;
    r )
      ;;
  esac
}

show_tpm_totp_hotp_options_menu()
{
  whiptail $BG_COLOR_MAIN_MENU --clear --title "TPM/TOTP/HOTP Options" \
    --menu "Select An Option" 20 90 10 \
    'g' ' Generate new TOTP/HOTP secret' \
    'r' ' Reset the TPM' \
    't' ' TOTP/HOTP does not match after refresh, troubleshoot' \
    'm' ' <-- Return to main menu' \
    2>/tmp/whiptail || recovery "GUI menu failed"

  option=$(cat /tmp/whiptail)
  case "$option" in 
    g )
      generate_totp_htop
      ;;
    r )
      reset_tpm
      ;;
    t )
      prompt_totp_mismatch
      ;;
    m )
      ;;
  esac
}

prompt_totp_mismatch()
{
  if (whiptail $BG_COLOR_WARNING --title "TOTP/HOTP code mismatched" \
    --yesno "TOTP/HOTP code mismatches could indicate either TPM tampering or clock drift:\n\nTo correct clock drift: 'date -s HH:MM:SS'\nand save it to the RTC: 'hwclock -w'\nthen reboot and try again.\n\nWould you like to exit to a recovery console?" 30 90) then
    echo ""
    echo "To correct clock drift: 'date -s HH:MM:SS'"
    echo "and save it to the RTC: 'hwclock -w'"
    echo "then reboot and try again"
    echo ""
    recovery "TOTP/HOTP mismatch"
  fi
}

reset_tpm()
{
  if [ "$CONFIG_TPM" = "y" ]; then
    if (whiptail $BG_COLOR_WARNING --title 'Reset the TPM' \
        --yesno "This will clear the TPM and TPM password, replace them with new ones!\n\nDo you want to proceed?" 16 90) then
      /bin/tpm-reset

      # now that the TPM is reset, remove invalid TPM counter files
      mount_boot
      mount -o rw,remount /boot
      rm -f /boot/kexec_rollback.txt

      # create Heads TPM counter before any others
      check_tpm_counter /boot/kexec_rollback.txt \
      || die "Unable to find/create tpm counter"
      counter="$TPM_COUNTER"

      increment_tpm_counter $counter \
      || die "Unable to increment tpm counter"

      sha256sum /tmp/counter-$counter > /boot/kexec_rollback.txt \
      || die "Unable to create rollback file"
      mount -o ro,remount /boot

      generate_totp_htop
    else
      echo "Returning to the main menu"
    fi
  else
    whiptail $BG_COLOR_ERROR --clear --title 'ERROR: No TPM Detected' --msgbox "This device does not have a TPM.\n\nPress OK to return to the Main Menu" 30 90
  fi
}

show_system_info()
{
  memtotal=$(cat /proc/meminfo | grep 'MemTotal' | tr -s ' ' | cut -f2 -d ' ')
  memtotal=$((${memtotal} / 1024 / 1024 + 1))
  cpustr=$(cat /proc/cpuinfo | grep 'model name' | uniq | sed -r 's/\(R\)//;s/\(TM\)//;s/CPU //;s/model name.*: //')
  kernel=$(uname -s -r)
  whiptail $BG_COLOR_MAIN_MENU --title 'System Info' \
    --msgbox "${BOARD_NAME}\n\nFW_VER: ${FW_VER}\nKernel: ${kernel}\n\nCPU: ${cpustr}\nRAM: ${memtotal} GB\n\n$(fdisk -l | grep -e '/dev/sd.:' -e '/dev/nvme.*:' | sed 's/B,.*/B/')" 16 60
}

select_os_boot_option()
{
  mount_boot
  if verify_global_hashes ; then
    kexec-select-boot -m -b /boot -c "grub.cfg" -g
  fi
}

attempt_default_boot()
{
  mount_boot
    
  if ! verify_global_hashes; then
    return
  fi
  DEFAULT_FILE=`find /boot/kexec_default.*.txt 2>/dev/null | head -1`
  if [ -r "$DEFAULT_FILE" ]; then
    kexec-select-boot -b /boot -c "grub.cfg" -g \
    || recovery "Failed default boot"
  elif (whiptail $BG_COLOR_WARNING --title 'No Default Boot Option Configured' \
        --yesno "There is no default boot option configured yet.\nWould you like to load a menu of boot options?\nOtherwise you will return to the main menu." 16 90) then
    kexec-select-boot -m -b /boot -c "grub.cfg" -g
  fi
}

force_unsafe_boot()
{
  # Run the menu selection in "force" mode, bypassing hash checks
  if (whiptail $BG_COLOR_WARNING --title 'Unsafe Forced Boot Selected!' \
      --yesno "WARNING: You have chosen to skip all tamper checks and boot anyway.\n\nThis is an unsafe option!\n\nDo you want to proceed?" 16 90) then
    mount_boot && kexec-select-boot -m -b /boot -c "grub.cfg" -g -f
  fi
}

# gui-init start

# Use stored HOTP key branding
if [ -r /boot/kexec_hotp_key ]; then
	HOTPKEY_BRANDING="$(cat /boot/kexec_hotp_key)"
else
	HOTPKEY_BRANDING="HOTP USB Security Dongle"
fi

if detect_boot_device ; then
  # /boot device with installed OS found
  clean_boot_check
else
  # can't determine /boot device or no OS installed, 
  # so fall back to interactive selection
  mount_boot
fi

# detect whether any GPG keys exist in the keyring, if not, initialize that first
[[ "$skip_to_menu" != "true" ]] && check_gpg_key && update_totp && update_hotp

if [[ "$HOTP" = "Success" && $CONFIG_AUTO_BOOT_TIMEOUT ]]; then
  prompt_auto_default_boot
fi

while true; do
  show_main_menu
done

recovery "Something failed during boot"