#!/bin/bash
set -euo pipefail
IFS=$'\n\t'

VERSION="0.0.4"
STARTDATE=0
STARTDATESTRING=""

usage() {
	cat <<EOF
Usage: mtm-ddwipe DEVICE

Wipe a block device.

Warnings:
  - Destructive and irreversible.
  - Target must be a real block device, not mounted, and not in use.
  - Optional tools: blkdiscard, ddrescue, dd_rescue, nvme, hdparm.
  - dd fallback may take a long time.

Version: ${VERSION}
EOF
}

log() {
	echo "$*" >&2
}

die() {
	echo "Error: $*" >&2
	exit 1
}

check_args() {
	if [ $# -ne 1 ]; then
		usage
		exit 1
	fi

	case "$1" in
		-h|--help)
			usage
			exit 0
			;;
		-*)
			die "Invalid option."
			;;
	esac
}

check_device() {
	local dev="$1"

	[ -e "$dev" ] || die "Missing device: $dev"
	[ -b "$dev" ] || die "Not a block device: $dev"
}

check_device_not_in_use() {
	local dev="$1"

	if findmnt -rn --target "$dev" >/dev/null 2>&1; then
		die "Device is mounted: $dev"
	fi

	if lsblk -nrpo NAME,MOUNTPOINT "$dev" | awk '$2 != "" { found=1 } END { exit !found }'; then
		die "Device or child is mounted: $dev"
	fi
}

is_nvme_device() {
	local dev="$1"
	local sysdev

	sysdev="/sys$(lsblk -ndo PATH "$dev" 2>/dev/null | head -n1)"
	[ -n "$sysdev" ] || return 1
	[ -e "$sysdev" ] || return 1
	[ -d "$sysdev/device" ] || return 1
	[ -e "$sysdev/device/uevent" ] || return 1
	grep -q '^NVME=1$' "$sysdev/device/uevent"
}

is_ata_device() {
	local dev="$1"
	local tran

	tran="$(lsblk -ndo TRAN "$dev" 2>/dev/null | head -n1)"
	[ "$tran" = "sata" ] || [ "$tran" = "ata" ]
}

confirm_wipe() {
	local dev="$1"
	local choice=""

	echo "Device:"
	lsblk -o NAME,SIZE,TYPE,FSTYPE,MOUNTPOINT,MODEL,SERIAL "$dev"
	echo ""
	echo "Type exactly: WIPE $dev"
	read -r -p "Confirm: " choice
	[ "$choice" = "WIPE $dev" ] || die "Canceled"
	echo ""
}

confirm_root() {
	if [ "${EUID:-$(id -u)}" -ne 0 ]; then
		die "This tool must be run as root."
	fi
}

format_duration() {
	local total="$1"
	local hours minutes seconds

	hours=$((total / 3600))
	minutes=$(((total % 3600) / 60))
	seconds=$((total % 60))

	printf '%02d:%02d:%02d\n' "$hours" "$minutes" "$seconds"
}

print_time() {
	local enddate calctime

	echo ""
	log "Start date :"
	log "$STARTDATESTRING"

	enddate=$(date +%s)
	calctime=$((enddate - STARTDATE))

	echo ""
	log "End date :"
	date >&2

	echo ""
	log "Total time :"
	format_duration "$calctime" >&2
}

wipe_with_blkdiscard_secure() {
	local dev="$1"

	log "blkdiscard secure"
	blkdiscard -f -p 500M -s -v "$dev"
}

wipe_with_blkdiscard_zero() {
	local dev="$1"

	log "blkdiscard zero"
	blkdiscard -f -p 500M -z -v "$dev"
}

wipe_with_dd() {
	local dev="$1"

	log "dd zero"
	dd if=/dev/zero of="$dev" bs=1M status=progress conv=fsync
}

wipe_with_ddrescue() {
	local dev="$1"

	command -v ddrescue >/dev/null 2>&1 || return 1
	log "ddrescue zero"
	ddrescue -f -n /dev/zero "$dev"
}

wipe_with_dd_rescue() {
	local dev="$1"

	command -v dd_rescue >/dev/null 2>&1 || return 1
	log "dd_rescue zero"
	dd_rescue -f /dev/zero "$dev"
}

wipe_with_nvme() {
	local dev="$1"

	command -v nvme >/dev/null 2>&1 || return 1
	is_nvme_device "$dev" || return 1
	log "nvme format"
	nvme format "$dev" -s 1 >/dev/null
}

wipe_with_hdparm() {
	local dev="$1"

	command -v hdparm >/dev/null 2>&1 || return 1
	is_ata_device "$dev" || return 1
	log "hdparm secure erase"
	hdparm --security-erase NULL "$dev"
}

wipe_dev() {
	local dev="$1"

	STARTDATE=$(date +%s)
	STARTDATESTRING="$(date)"

	log "Begin wiping: $dev"
	echo ""
	log "Start date :"
	log "$STARTDATESTRING"
	echo ""

	if wipe_with_blkdiscard_secure "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_blkdiscard_zero "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_dd "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_ddrescue "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_dd_rescue "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_nvme "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	echo ""
	if wipe_with_hdparm "$dev"; then
		echo ""
		log "Device $dev wiped."
		return
	fi

	die "Wipe failed. The device may not be fully overwritten."
}

main() {
	check_args "$@"
	confirm_root
	check_device "$1"
	check_device_not_in_use "$1"
	confirm_wipe "$1"
	wipe_dev "$1"
	print_time
}

main "$@"
