#!/usr/bin/env bash
#
# erase-storage.sh — overwrite a partition or securely erase a whole disk.
# Full documentation: see erase-storage.md in this same folder.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib/common.sh
source "${SCRIPT_DIR}/../lib/common.sh"

if [[ "$EUID" -ne 0 ]]; then
    common::log "WARNING: This script must be run as root (for example, with sudo)."
    common::die "stopping before processing arguments or checking dependencies"
fi

SCRIPT_HELP=$(cat <<'HELP'
Usage: sudo ./HDD/erase-storage.sh DEVICE [zero|random]

For a partition, writes one full pass of zeros (default) or pseudorandom data.
For a whole disk, detects and runs a supported ATA Secure Erase or NVMe Sanitize
operation. Whole-disk erasure requires two separate confirmations.
Lists available block devices and mount points before continuing, highlighting
the selected device in red when run in a terminal.
At startup, asks you to confirm that you can keep the computer and drive
powered and connected for the hours the operation may take.

Arguments:
  DEVICE             Partition (for example /dev/sda1) or whole disk
  zero               Write zeros to a partition (default)
  random             Write pseudorandom data to a partition

Options:
  -h, --help         Show this help message and exit

Examples:
  sudo ./HDD/erase-storage.sh /dev/sda1
  sudo ./HDD/erase-storage.sh /dev/nvme0n1p2 random
  sudo ./HDD/erase-storage.sh /dev/sda
  sudo ./HDD/erase-storage.sh /dev/nvme0n1
HELP
)

common::show_help_if_requested "$@"

[[ $# -ge 1 && $# -le 2 ]] || common::die "expected DEVICE [zero|random]; run with --help for usage"

DEVICE_INPUT="$1"
if [[ "$DEVICE_INPUT" == dev/* ]]; then
    DEVICE_INPUT="/$DEVICE_INPUT"
fi
[[ "$DEVICE_INPUT" == /dev/* ]] || common::die "DEVICE must be a path under /dev"

FILL_MODE="${2:-zero}"
case "$FILL_MODE" in
    zero|random)
        ;;
    *)
        common::die "partition fill mode must be 'zero' or 'random'"
        ;;
esac

for required_command in awk findmnt lsblk readlink swapon date dd sleep; do
    common::require_cmd "$required_command"
done

[[ -b "$DEVICE_INPUT" ]] || common::die "$DEVICE_INPUT does not exist or is not a block device"

DEVICE="$(readlink -f -- "$DEVICE_INPUT")" || common::die "cannot resolve device path: $DEVICE_INPUT"
DEVICE_TYPE="$(lsblk -dnro TYPE -- "$DEVICE")" || common::die "cannot determine device type"
[[ "$DEVICE_TYPE" == "part" || "$DEVICE_TYPE" == "disk" ]] \
    || common::die "only a partition or a whole disk is accepted (detected: $DEVICE_TYPE)"
common::show_block_devices "$DEVICE" || common::die "cannot list available block devices"

ROOT_SOURCE="$(findmnt -n -o SOURCE /)" || common::die "cannot determine the root filesystem device"
ROOT_SOURCE="${ROOT_SOURCE%%\[*}"
[[ -b "$ROOT_SOURCE" ]] || common::die "cannot safely resolve the root filesystem device ($ROOT_SOURCE)"
ROOT_DEVICE="$(readlink -f -- "$ROOT_SOURCE")" || common::die "cannot resolve the root filesystem device"
ROOT_TREE="$(lsblk --inverse --list --noheadings --paths --output PATH -- "$ROOT_DEVICE")" \
    || common::die "cannot determine the devices backing the root filesystem"

while IFS= read -r root_device; do
    [[ -n "$root_device" ]] || continue
    if [[ "$(readlink -f -- "$root_device")" == "$DEVICE" ]]; then
        common::die "$DEVICE backs the root filesystem; refusing to erase it"
    fi
done <<< "$ROOT_TREE"

DEVICE_TREE="$(lsblk --list --noheadings --paths --output PATH -- "$DEVICE")" \
    || common::die "cannot inspect devices below $DEVICE"
ACTIVE_SWAP_DEVICES="$(swapon --show --noheadings --raw --output NAME)" \
    || common::die "cannot inspect active swap devices"

while IFS= read -r child_device; do
    [[ -n "$child_device" ]] || continue
    if findmnt -rn -S "$child_device" >/dev/null 2>&1; then
        common::die "$child_device is mounted; unmount it before running this script"
    fi

    while IFS= read -r swap_device; do
        [[ -n "$swap_device" ]] || continue
        if [[ -b "$swap_device" ]]; then
            swap_source="$swap_device"
        else
            swap_mount_source="$(findmnt -n -o SOURCE -T "$swap_device")" \
                || common::die "cannot determine the filesystem backing active swap file $swap_device"
            swap_mount_source="${swap_mount_source%%\[*}"
            [[ -b "$swap_mount_source" ]] \
                || common::die "cannot safely resolve the block device backing active swap file $swap_device"
            swap_source="$swap_mount_source"
        fi

        swap_source_tree="$(lsblk --inverse --list --noheadings --paths --output PATH -- "$swap_source")" \
            || common::die "cannot determine the devices backing active swap $swap_device"
        while IFS= read -r swap_backing_device; do
            [[ -n "$swap_backing_device" ]] || continue
            if [[ "$(readlink -f -- "$swap_backing_device")" == "$(readlink -f -- "$child_device")" ]]; then
                common::die "$child_device backs active swap ($swap_device); disable it before running this script"
            fi
        done <<< "$swap_source_tree"
    done <<< "$ACTIVE_SWAP_DEVICES"
done <<< "$DEVICE_TREE"

common::log "WARNING: Erasing data may take hours. Do not power off, reboot, or suspend the computer, or disconnect power or the drive."
common::confirm "Can the computer and drive remain powered and connected until the operation finishes?" || exit 0

if [[ "$DEVICE_TYPE" == "part" ]]; then
    [[ $# -eq 1 || "$FILL_MODE" == "zero" || "$FILL_MODE" == "random" ]] \
        || common::die "invalid fill mode; use 'zero' or 'random'"
    DEVICE_SIZE="$(lsblk -dnro SIZE --bytes -- "$DEVICE")" || common::die "cannot determine partition size"

    common::log "WARNING: this will irreversibly overwrite all data on partition $DEVICE ($DEVICE_SIZE bytes)."
    read -r -p "Type the exact device path '$DEVICE' to continue: " CONFIRM
    [[ "$CONFIRM" == "$DEVICE" ]] || common::die "operation cancelled; no data was changed"

    common::log "Writing one full $FILL_MODE pass to $DEVICE."
    if [[ "$FILL_MODE" == "zero" ]]; then
        dd if=/dev/zero of="$DEVICE" bs=4M count="$DEVICE_SIZE" iflag=count_bytes conv=fsync status=progress
    else
        dd if=/dev/urandom of="$DEVICE" bs=4M count="$DEVICE_SIZE" iflag=count_bytes conv=fsync status=progress
    fi
    common::log "Partition overwrite completed successfully."
    exit 0
fi

[[ $# -eq 1 ]] || common::die "a fill mode is only valid for a partition"
DEVICE_SIZE="$(lsblk -dnro SIZE --bytes -- "$DEVICE")" || common::die "cannot determine disk size"
DEVICE_NAME="${DEVICE##*/}"
ERASE_METHOD=""
ERASE_TOOL=""
NVME_ACTION=""
NVME_ACTION_NAME=""
ATA_ERASE_OPTION=""

if [[ "$DEVICE_NAME" =~ ^(nvme[0-9]+)n[0-9]+$ ]]; then
    common::require_cmd nvme
    common::require_cmd jq
    NVME_CONTROLLER="/dev/${BASH_REMATCH[1]}"
    NVME_IDENTIFY="$(nvme id-ctrl "$NVME_CONTROLLER" -o json)" \
        || common::die "could not read NVMe controller capabilities"
    SANICAP="$(jq -er '.sanicap' <<< "$NVME_IDENTIFY")" \
        || common::die "NVMe controller did not report SANICAP"
    [[ "$SANICAP" =~ ^[0-9]+$ ]] || common::die "invalid NVMe SANICAP value: $SANICAP"
    SANICAP_VALUE=$((SANICAP))

    if (( SANICAP_VALUE & 1 )); then
        NVME_ACTION=4
        NVME_ACTION_NAME="Crypto Erase"
    elif (( SANICAP_VALUE & 2 )); then
        NVME_ACTION=2
        NVME_ACTION_NAME="Block Erase"
    elif (( SANICAP_VALUE & 4 )); then
        NVME_ACTION=3
        NVME_ACTION_NAME="Overwrite"
    else
        common::die "NVMe controller does not advertise a supported Sanitize action"
    fi
    ERASE_METHOD="NVMe Sanitize: $NVME_ACTION_NAME"
    ERASE_TOOL="nvme"
else
    common::require_cmd hdparm
    common::require_cmd awk
    common::require_cmd grep
    ATA_IDENTIFY="$(hdparm -I "$DEVICE" 2>&1)" || common::die "could not read ATA security capabilities"
    ATA_SECURITY="$(awk '
        /^[[:space:]]*Security:/ { in_security = 1; print; next }
        in_security && /^[^[:space:]]/ { exit }
        in_security { print }
    ' <<< "$ATA_IDENTIFY")"
    [[ -n "$ATA_SECURITY" ]] || common::die "device did not report an ATA Security section"
    grep -Eiq '^[[:space:]]+supported[[:space:]]*$' <<< "$ATA_SECURITY" \
        || common::die "device does not report ATA Security Erase support"
    grep -Eiq '^[[:space:]]+not frozen[[:space:]]*$' <<< "$ATA_SECURITY" \
        || common::die "ATA security is frozen or its state could not be determined; refusing to continue"
    grep -Eiq '^[[:space:]]+not enabled[[:space:]]*$' <<< "$ATA_SECURITY" \
        || common::die "ATA security is already enabled; refusing to risk leaving the disk locked"

    if grep -Eiq 'supported:[[:space:]]+enhanced erase' <<< "$ATA_SECURITY"; then
        ATA_ERASE_OPTION="--security-erase-enhanced"
        ERASE_METHOD="ATA Enhanced Security Erase"
    else
        ATA_ERASE_OPTION="--security-erase"
        ERASE_METHOD="ATA Security Erase"
    fi
    ERASE_TOOL="hdparm"
fi

common::log "WARNING: whole-disk operation on $DEVICE ($DEVICE_SIZE bytes)."
common::log "Detected supported method: $ERASE_METHOD. ALL DATA on this disk will be destroyed."
read -r -p "First confirmation: type ERASE to continue: " FIRST_CONFIRM
[[ "$FIRST_CONFIRM" == "ERASE" ]] || common::die "operation cancelled; no data was changed"
common::log "Final warning: $ERASE_METHOD will erase the entire disk $DEVICE, including every partition."
read -r -p "Final confirmation: type the exact device name '$DEVICE_NAME': " FINAL_CONFIRM
[[ "$FINAL_CONFIRM" == "$DEVICE_NAME" ]] || common::die "operation cancelled; no data was changed"

if [[ "$ERASE_TOOL" == "nvme" ]]; then
    common::log "Starting $ERASE_METHOD on $DEVICE. The device may continue sanitizing after the command returns."
    nvme sanitize "$DEVICE" --sanact="$NVME_ACTION"

    while true; do
        SANITIZE_LOG="$(nvme sanitize-log "$DEVICE" -o json)" \
            || common::die "could not read NVMe Sanitize status; the operation may still be running"
        SANITIZE_STATUS="$(jq -er '.sstat' <<< "$SANITIZE_LOG")" \
            || common::die "NVMe Sanitize status was missing or invalid"
        [[ "$SANITIZE_STATUS" =~ ^[0-9]+$ ]] || common::die "invalid NVMe Sanitize status: $SANITIZE_STATUS"

        case $((SANITIZE_STATUS & 7)) in
            1|4)
                common::log "NVMe Sanitize completed successfully."
                break
                ;;
            2)
                common::log "NVMe Sanitize is still in progress; checking again in 30 seconds."
                sleep 30
                ;;
            3)
                common::die "NVMe Sanitize reported failure; consult the device sanitize log"
                ;;
            *)
                common::die "NVMe Sanitize is not running or reported an unknown status ($SANITIZE_STATUS)"
                ;;
        esac
    done
else
    common::log "Starting $ERASE_METHOD on $DEVICE. Do not interrupt power or disconnect the drive."
    ATA_PASSWORD="strix-erase"
    hdparm --user-master u --security-set-pass "$ATA_PASSWORD" "$DEVICE"
    if ! hdparm --user-master u "$ATA_ERASE_OPTION" "$ATA_PASSWORD" "$DEVICE"; then
        common::die "ATA erase failed; the disk may remain security-locked with the temporary password '$ATA_PASSWORD'"
    fi
    common::log "$ERASE_METHOD completed successfully."
fi