#!/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"

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.
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

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

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

[[ "$EUID" -eq 0 ]] || common::die "this script must be run as root (for example, with sudo)"
[[ -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)"

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"

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