#!/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
		case "${1:-}" in
			-h|--help)
				usage
				exit 0
				;;
			"")
				usage
				exit 1
				;;
			-*)
				die "Invalid option."
				;;
		esac

		usage
		exit 1
	fi
}

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 "$@"
