#!/bin/sh

set -eu

ask_disk() {
  disks=$(awk '$4 ~ /^((md|mmcblk)[0-9]+|nvme[0-9]+n[0-9]+|[sv]d[a-z]|xvd[a-z])$/ {
                 print "/dev/"$4
               }' /proc/partitions)
  [ -n "$disks" ] || return 1
  printf %s "$disks" | fzf -1 --height=5 --layout=reverse-list --color=bg+:#FFFFFF,fg+:#000000
}

prepare_for_erase() {
  ### There might be leftover nvram dumps that could prevent
  ### a successful bootloader installation later on. Clean them up.
  rm -f /sys/firmware/efi/efivars/dump-*
  ### Disable all swap partitions, unmount Puavo LVM.
  echo 'Disabling all swap partitions..'
  swapoff -a
  echo 'Unmounting all LVM volumes named "puavo"'
  lvchange -an puavo
}

if [ "$(id -u)" -ne 0 ]; then
  echo 'This tool must be run as root!' >&2
  exit 1
fi

bootserver_address=$(cat /run/puavo/nbd-server 2>/dev/null) || true

usage() {
  cat <<EOF
$(basename $0) [--force] [--normal] [device]

  --force               Use the Force
  --normal              We default to doing an Enhanced Secure Erase (ATA) or a Cryptographic Erase (NVMe), provided the device supports such action.
                        With this, one can opt to do a "normal" Secure Erase instead.
  device                Specify the device name (e.g. sda, sdb). If not specified here, it will be asked.
EOF
  exit 1
}

useforce="false"
enhanced="true"
erase_mode=""

if ! args=$(getopt -n "$0" -o + -l 'force,normal' -- "$@"); then
  usage
fi

eval "set -- $args"
while [ $# -ne 0 ]; do
  case "$1" in
    --force)    useforce="true" ;  shift ;;
    --normal)   enhanced="false";  shift ;;
    --) shift; break ;;
    *)   ;;
  esac
done

if [ "$@" ]; then
  target_disk="/dev/$@"
else
  echo 'Choose a disk to Secure Erase:'
  if ! target_disk=$(ask_disk); then
    echo 'No disk selected, aborting!' >&2
    exit 1
  fi
fi

echo "Chosen << ${target_disk} >> for erase"
case "$target_disk" in
  *nvme*)
    pde_error_code=0
    blkd_error_code=0

    echo 'Preparing system for secure erase. This should not take long.'
    prepare_for_erase && \
    echo 'Preparing finished, moving on..'
    echo 'Selected disk is an NVMe disk. Using "nvme-cli" utility for Secure Erase.'
    echo 'Going to sleep to as a work around for some NVMe disk firmware bugs that may prevent Secure Erase.'
    sleep 1

    rtcwake -m mem -s 5
    sleep 1

    if ("$enhanced"); then
      if nvme id-ctrl "$target_disk" --vendor-specific -H | grep "Crypto Erase Supported as part of Secure Erase" >&2; then
        echo 'Disk supports Cryptographic Erase and override was not given. Running Cryptographic Erase.'
        erase_mode="-s2"
      fi
    else
      echo 'Disk either does not support Cryptographic Erase or override was given. Running Secure Erase.'
      erase_mode="-s1"
    fi

    if ("$useforce"); then
      command time --format "Time elapsed: %E" nvme format --force $erase_mode "$target_disk" || pde_error_code=$?
    else
      command time --format "Time elapsed: %E" nvme format $erase_mode "$target_disk" || pde_error_code=$?
    fi

    if [ "${pde_error_code}" -ne 0 ]; then
      ### This means we have ran into some issue during Secure Erase
      echo "Secure Erase did not work. Fallback to blkdiscard (Secure)"
      command time --format "Time elapsed: %E" /usr/sbin/blkdiscard --secure --force "$target_disk" || blkd_error_code=$?
      if [ "${blkd_error_code}" -ne 0 ]; then
        printf "%s\n%s\n" "blkdiscard (Secure) failed, error code ${blkd_error_code}." "Fallback to standard blkdiscard"
        command time --format "Time elapsed: %E" /usr/sbin/blkdiscard --force "$target_disk"
      fi
    fi
    ;;
  *)
    pde_error_code=0
    blkd_error_code=0

    echo 'Preparing system for secure erase. This should not take long.'
    prepare_for_erase && echo 'Preparing finished, moving on..'

    if [ -n "$bootserver_address" ]; then
      if ! hdparm -I "$target_disk" | grep "SECURITY ERASE UNIT" >&2; then
        echo 'Selected disk does not support Secure Erase!' >&2
        exit 1
      elif ("$enhanced") && hdparm -I "$target_disk" | grep "ENHANCED SECURITY ERASE UNIT" >&2; then
        echo 'Disk supports Enhanced Secure Erase and override was not given. Running Enhanced Secure Erase.'
        erase_mode="--security-erase-enhanced"
      else
       echo 'Disk either does not support Enhanced Secure Erase or override was given. Running Secure Erase.'
        erase_mode="--security-erase"
      fi

      if ! hdparm -I "$target_disk" | grep -P "not\tfrozen" >&2; then
        if ("$useforce"); then
          echo 'Flag "--force" was given.'
          echo 'Disk is in frozen state! Going to sleep now.'
        else
          echo 'Disk is in frozen state, cannot continue!' >&2
          echo 'Please close the lid of the device for a few seconds, reopen it and try again.\n' >&2
          echo 'If we are netbooted, closing the lid will not put the device to sleep.'
          echo 'We can however put it to sleep and wake it up with the usage of rtcwake.\n'

          read -p "Shall I run the command for you now? (y/n)[y] " gosleep
          case "$gosleep" in
            [yY]|"")
              echo 'Going to sleep now.'
              sleep 1
              ;;
            [nN])
              echo 'Not going to sleep. You must close and reopen the lid, or run the command by hand. Then retry.'
              sleep 1
              exit 1
              ;;
            *)
              echo 'Invalid answer, exiting.' >&2
              sleep 1
              exit 1
              ;;
          esac
        fi

        rtcwake -m mem -s 5
        sleep 1

        if ! hdparm -I "$target_disk" | grep -P "not\tfrozen" >&2; then
          echo 'Disk is still in frozen state. Aborting Secure Erase.' >&2
          echo 'You could either, run Secure Erase from device BIOS (if supported natively),\n' >&2
          echo 'or use blkdiscard and then fill the drive with junk.' >&2
          pde_error_code=3
        elif [ "${pde_error_code}" -ne 0 ]; then
          ### This means we have ran into some issue during Secure Erase
          echo "Secure Erase did not work. Fallback to blkdiscard (Secure)"
          command time --format "Time elapsed: %E" /usr/sbin/blkdiscard --secure --force "$target_disk" || blkd_error_code=$?
          if [ "${blkd_error_code}" -ne 0 ]; then
            printf "%s\n%s\n" "blkdiscard (Secure) failed, error code ${blkd_error_code}." "Fallback to standard blkdiscard"
            command time --format "Time elapsed: %E" /usr/sbin/blkdiscard --force "$target_disk"
          fi
        else
          echo 'Disk is not in frozen state, we can continue.'

          hdparm --user-master user --security-set-pass Pass "$target_disk"
          if ! hdparm -I "$target_disk" | grep "Security level high" >&2; then
            echo 'Failed setting user and password, exiting.' >&2
            exit 2
          fi

          echo 'User and password set, proceeding to Secure Erase.'
          command time --format "Time elapsed: %E" hdparm --user-master user "$erase_mode" Pass "$target_disk"
          echo 'Command completed. You may check how long the erase took above.'

          if ! hdparm -I "$target_disk" | grep -P "not\tlocked"; then
            echo 'Disk seems to be locked after erase. Unlocking now.'
            hdparm --user-master user --security-unlock Pass "$target_disk"
          fi
        fi
      fi
    fi
    ;;
esac

partprobe

if [ -e "/sys/firmware/efi" ]; then
  efitable="$(efibootmgr | awk '/debian/ {print substr($0,5,4)} /Windows/ {print substr($0,5,4)} /Ubuntu/ {print substr($0,5,4)}')"

  if [ -n "$efitable" ]; then
    echo 'Removing leftover EFI boot entries.'
    for entry in $efitable;
    do
      efibootmgr -q -b "$entry" -B
    done
    echo 'Leftover EFI boot entries have been removed.'
  else
    echo 'No leftover EFI boot entries to remove.'
  fi
fi

echo 'You may check how long the erase took above.'
echo 'If you see no errors, you may re-plug the disk or reboot the device.\n'
sleep 2
