CONFIG_FILE="/home/nps/.config/wipe.conf" # Configuration file location. DO NOT CHANGE!
EXCLUDE_DEVICE="/dev/nvme0n1" # Boot drive, should not be changed. DO NOT CHANGE!

BUS_TYPE="" # NVMe wipe protocol by default.
MEDIA_TYPE="" # SSD by default.
DEVICE="" # Default device
LOG_FILE=""

wipe-confirmation() {
  local wipe_status="$1"
  local smart_status="$2"

  if [ "$wipe_status" -eq 0 ]; then
          wipe_status="WIPE:  [ SUCCESS ]"
  else
          wipe_status="WIPE:  [ FAILED  ]"
  fi

  if [ "$smart_status" -eq 0 ]; then
          smart_status="SMART: [ PASSED  ]"
  else
          smart_status="SMART: [ FAILED  ]"
  fi

  { echo "
===================================
          RESULTS SUMMARY          
===================================

        $wipe_status

        $smart_status

===================================
  "; cat "$LOG_FILE"; } > tmp.log && mv tmp.log $LOG_FILE


  echo "
===================================
          RESULTS SUMMARY          
===================================

        $wipe_status

        $smart_status

===================================
  "
}

smart-info() {
  smartctl -a "$DEVICE" >> "$LOG_FILE" 2>&1
  return ${PIPESTATUS[0]}
}

init-log() {
  mkdir -p "/home/nps/logs"

  SN="$(smartctl -a "$DEVICE" | awk '/Serial Number:/ {print $3; exit}')"
  LOG_FILE="/home/nps/logs/$SN.log"
  echo "" > "$LOG_FILE"
}

verify-wipe() {
  local wipe_status
  echo "Verifying wipe..." | tee -a "$LOG_FILE"
  sleep 2

  if dd if="$DEVICE" bs=1M count=1024 status=none | cmp -n 1073741824 -l - /dev/zero; then
    echo "Verification passed." | tee -a "$LOG_FILE"
    wipe_status=0
  else
    echo "Verification failed." | tee -a "$LOG_FILE"
    wipe_status=1
  fi

  return "$wipe_status"
}

# Returns 0 if ready, 1 if unsupported, 2 if frozen/aborted.
sata-ssd-security-state() {
  local status
  status=$(sudo hdparm -I "$DEVICE")

  # Check if supported
  if ! echo "$status" | grep -q "supported"; then
    echo "Error: Drive does not support ATA security features."
    return 1 # Unsupported
  fi

  # Check if frozen
  if echo "$status" | grep -q "frozen" | grep -v -q "not"; then
    echo "Error: Drive is in a 'frozen' state. Hot-plug required."
    read -p "Has drive been hot-plugged? (y/N)" -r -n 1
    echo
    
    if [[ "$REPLY" =~ ^[Yy]$ ]]; then
      sata-ssd-security-state
      return $?
    fi

    return 2 # Frozen/aborted by user.
  fi

  return 0 # Ready
}

wipe-sata-ssd() {
  local wipe_status

  echo "Checking security state..."
  sata-ssd-security-state
  local state=$?

  case $state in
    "0")
      echo "Executing ATA Secure Erase..." | tee -a "$LOG_FILE"
      hdparm --user-master u --security-set-pass p "$DEVICE" >> "$LOG_FILE" 2>&1
      hdparm --user-master u --security-erase p "$DEVICE" >> "$LOG_FILE" 2>&1
      wipe_status=$?
    ;;
    "1")
      echo "ATA Security not supported. Falling back to blkdiscard..." | tee -a "$LOG_FILE"
      blkdiscard --secure "$DEVICE" >> "$LOG_FILE" 2>&1
      wipe_status=$?
    ;;
    "2")
      echo "Aborted: Drive is frozen or user declined hot-plug."
      return 2
    ;;
    *)
      echo "Unknown state returned."
      return 1
    ;;
  esac

  if [ $wipe_status -eq 0 ]; then
    blockdev --rereadpt "$DEVICE"
    partprobe "$DEVICE"
    udevadm settle

    verify-wipe
    wipe_status=$?
  fi

  # Write SMART info to LOG and save result for confirmation display.
  smart-info
  wipe-confirmation "$wipe_status" "$?"
  echo -e "\n\nSATA SSD wipe complete and ready to remove."
}

wipe-hdd() {
  echo "Wiping HDD..." | tee -a "$LOG_FILE"
  shred -vz -n 1 "$DEVICE"
  local wipe_status=$?

  if [ $wipe_status -eq 0 ]; then
    verify-wipe
    wipe_status=$?
  fi

  # Write SMART info to LOG and save result for confirmation display.
  smart-info
  wipe-confirmation "$wipe_status" "$?"
  echo -e "\n\nHDD wipe complete and ready to remove."
}

wipe-emmc() {
  echo "Wiping eMMC..." | tee -a "$LOG_FILE"
  blkdiscard "$DEVICE" >> "$LOG_FILE" 2>&1
  local wipe_status=$?

  if [ $wipe_status -eq 0 ]; then
    verify-wipe
    wipe_status=$?
  fi

  # Write SMART info to LOG and save result for confirmation display.
  smart-info
  wipe-confirmation "$wipe_status" "$?"
  echo -e "\n\neMMC wipe complete and ready to remove."
}

wipe-usb() {
  echo "Wiping USB..." | tee -a "$LOG_FILE"
  shred -vz -n 1 "$DEVICE"
  local wipe_status=$?

  if [ $wipe_status -eq 0 ]; then
    verify-wipe
    wipe_status=$?
  fi

  # Write SMART info to LOG and save result for confirmation display.
  smart-info
  wipe-confirmation "$wipe_status" "$?"
  echo -e "\n\nUSB wipe complete and ready to remove."
}

wipe-nvme() {
  # Wipe NVMe via Secure Erase NVMe feature.
  echo "Wiping NVMe..."
  nvme format "$DEVICE" --ses=1 --force >> "$LOG_FILE" 2>&1
  local wipe_status=$?

  # Write SMART info to LOG and save result for confirmation display.
  smart-info
  local smart_status=$?

  # Dynamic disconnect of NVMe drive
  local dev_node="${DEVICE#/dev/}"
  local pci_path=$(readlink -f "/sys/block/$dev_node/device" | xargs basename)
  echo "NVMe disconnect will take 15 seconds..."
  echo 1 > "/sys/bus/pci/devices/$pci_path/remove"
  sleep 15

  wipe-confirmation "$wipe_status" "$smart_status"
  echo -e "\n\nNVMe disconnected and is safe to remove."
}

show-help() {
  echo '
Created by Josh Ashton
Last modified: 20-07-2026
Email: me@joshashton.dev
Secondary Email: joshua.ashton@npsstore.com

Drive wiping utility for various drive types.

Usage:

  - `wipe`
    Use default wipe method (NVMe SSD format)

  - `wipe -h` & `wipe --help`
    Show this help page.
    Cannot be combined with other arguments.

  - `wipe -b nvme` & `wipe --bus-type nvme`
    Set the bus type that a drive is connected to (ie. nvme, sata). nvme is
    the default bus type. And will automatically set the media type to SSD.
    Use without providing a bus type to list available bus types.

  - `wipe -m ssd` & `wipe --media-type ssd`
    Set the media type to wipe (ie. ssd, hdd, emmc, usb). ssd is the
    default media type. 
    Use without providing a media type to list available media types.

  - `wipe -d /dev/nvme0n1` & `wipe --device /dev/nvme0n1`
    Set the device to wipe. Defaults to the first device of the bus type.
    Use without providing a device to list available devices.

  - `wipe -c` & `wipe --configure`
    Perform guided configuration to set new defaults. Changes are persistent
    and can be directly edited at `$HOME/.config/wipe.conf`.
    Cannot be combined with other arguments.

  Arguments can be combined for more control:
  - `wipe -b sata -m hdd -d /dev/sdb`
  '

  exit
}

detect-device() {
  if [[ "$BUS_TYPE" == "nvme" ]]; then
    echo "Detecting in 15 seconds..."
    sleep 10
    echo 1 > /sys/bus/pci/rescan
    sleep 5

    if [ ! -e "$DEVICE" ]; then
      echo "Rebooting in 3 seconds. If still not detected, NVMe is DOA."
      sleep 3
      reboot
    else
      echo "NVMe detected!"
    fi
  elif [[ "$BUS_TYPE" == "sata" && "$MEDIA_TYPE" == "ssd" ]]; then
    udevadm trigger
    
    if [ ! -e "$DEVICE" ]; then
      echo "Drive is DOA."
      exit 1
    fi
  fi

  init-log
}

run-wipe() {
  # Safety check
  if [[ "$DEVICE" == "$EXCLUDE_DEVICE" ]] || grep -q "$DEVICE" /proc/mounts; then
    echo "CRITICAL ERROR: Attempted to wipe protected or mounted device: $DEVICE"
    exit 1
  fi

  case "$BUS_TYPE:$MEDIA_TYPE" in
    nvme:ssd)  wipe-nvme ;;
    sata:ssd)  wipe-sata-ssd ;;
    *:hdd)     wipe-hdd ;;
    *:emmc)    wipe-emmc ;;
    *:usb)     wipe-usb ;;
    *)         echo "Unsupported configuration: $BUS_TYPE/$MEDIA_TYPE"; exit 1 ;;
  esac
}

configure-wipe() {
  # 1. Select Bus Type
  BUS_TYPE=$(dialog --title "Bus Type" --menu "Select bus type:" 15 60 4 \
    "nvme" "NVMe Drive" \
    "sata" "SATA Drive" \
    "usb"  "USB Storage" 3>&1 1>&2 2>&3) || exit

  # 2. Select Media Type
  MEDIA_TYPE=$(dialog --title "Media Type" --menu "Select media type:" 15 60 4 \
    "ssd" "Solid State Drive" \
    "hdd" "Hard Disk Drive" \
    "emmc" "eMMC Flash" 3>&1 1>&2 2>&3) || exit

  # 3. Select Device (Dynamic list of block devices)
  local dev_list=()
  while read -r dev; do
    dev_list+=("$dev" "Available device")
  done < <(lsblk -dpno NAME | grep -v "$EXCLUDE_DEVICE")

  if [ "${#dev_list[@]}" -eq 0 ]; then
    dialog --msgbox "No available drives detected." 8 40
    clear
    exit 1
  fi

  DEVICE=$(dialog --title "Device Selection" --menu "Select device to wipe:" 15 60 10 \
    "${dev_list[@]}" 3>&1 1>&2 2>&3) || exit

  # 4. Save to Config
  mkdir -p "$(dirname "$CONFIG_FILE")"
  echo "
BUS_TYPE=$BUS_TYPE
MEDIA_TYPE=$MEDIA_TYPE
DEVICE=$DEVICE
  " > "$CONFIG_FILE"

  dialog --msgbox "Configuration saved to $CONFIG_FILE" 8 40
  clear
  exit 0
}

main() {
  local GETOPT_BIN="/run/current-system/sw/bin/getopt"
  [ -f "$CONFIG_FILE" ] && source "$CONFIG_FILE"

  local OPTS
  OPTS=$("$GETOPT_BIN" -o b:m:d:hc -l bus-type:,media-type:,device:,help,configure -- "$@")
  if [ $? -ne 0 ]; then help; exit 1; fi
  eval set -- "$OPTS"

  while true; do
      case "$1" in
          -b|--bus-type)   BUS_TYPE="$2"; shift 2 ;;
          -m|--media-type) MEDIA_TYPE="$2"; shift 2 ;;
          -d|--device)     DEVICE="$2"; shift 2 ;;
          -h|--help)       show-help ;;
          -c|--configure)  configure-wipe ;;
          --)              shift; break ;;
          *)               echo "Unexpected error"; help; exit 1 ;;
      esac
  done

  if [[ -z "$DEVICE" ]]; then
      echo "Error: No device specified. Use -d or run -c to configure."
      exit 1
  fi
  init-log
  detect-device
  run-wipe
}

main "$@"