#!/usr/bin/env bash

. lib.bash || exit

usage() {
	echo "Usage: $0 [-E] [-F] <dev>"
	echo ""
	echo_opt "-E" "use ATA ENHANCED SECURITY ERASE (for SSDs)"
	echo_opt "-F" "override safety check for non-USB transport"
}

force_tran=die
enhanced=''
noconfirm=0

if [[ $SUDO_UID ]]; then
	notify-send() {
		sudo -u "#$SUDO_UID" \
			XDG_RUNTIME_DIR="/run/user/$SUDO_UID" \
			notify-send "$@"
	}
fi

while getopts ":EFy" OPT; do
	case $OPT in
	F) force_tran=warn;;
	E) enhanced='-enhanced';;
	y) noconfirm=1;;
	*) lib:die_getopts;;
	esac
done; shift $((OPTIND-1))

dev=$1
[[ -b $dev ]] || die "path '$dev' is not a block device"
[[ -w $dev ]] || die "no write permission for '$dev'"

info "ATA information:"
hdparm_data=$(hdparm -I "$dev")
echo "$hdparm_data" | awk 'x&&/^[^\t]/{x=0} /^(Security:|ATA device,)/{x=1} x{print}'

if ! echo "$hdparm_data" | grep -qs '^Security:'; then
	die "Disk '$dev' does not support ATA SECURITY features"
fi

if [[ $enhanced ]] &&
   ! echo "$hdparm_data" | grep -qs 'ENHANCED SECURITY ERASE UNIT'; then
	die "Disk '$dev' does not support enhanced erase mode"
fi

minutes=$(echo "$hdparm_data" | if [[ $enhanced ]]; then
	sed -E -n 's/.* ([0-9]+)min for ENHANCED SECURITY ERASE UNIT.*/\1m/p'
else
	sed -E -n 's/^	([0-9]+)min for SECURITY ERASE UNIT.*/\1m/p'
fi)
if [[ $minutes ]]; then
	seconds=$(( ${minutes%m} * 60 ))
	info "Expected erase duration: $(interval $seconds)"
else
	warn "Device did not report expected erase duration."
fi

info "Device identification:"
lsblk -S "$dev"
tran=$(lsblk -Sno TRAN "$dev")
if [[ $tran != 'usb' ]]; then
	$force_tran "Disk '$dev' is not USB-attached (found: $tran)"
fi

if [[ $(wipefs -p "$dev") ]]; then
	info "Device contents:"
	lsblk -o +FSTYPE,LABEL "$dev"
	warn "Disk $dev is not empty"
fi

((noconfirm)) || confirm "Erase '$dev'?" || exit

# Remove all partition device nodes.
for part in {1..15}; do
	if [[ -e $dev$part ]]; then
		info "Clearing partition $part"
		if findmnt -S "$dev$part" >/dev/null; then
			do: umount "$dev$part"
		fi &&
		do: wipefs -a "$dev$part" &&
		sleep 0.2 &&
		do: partx --delete --nr "$part" "$dev" ||
		die "Could not unmount $dev$part"
	fi
done

# Wipe partition table manually, so that if erase is interrupted, the next time
# we won't need to do this at all.
info "Clearing partition table"
wipefs -a "$dev"

info "Dumping SMART data before erase"
attrib="/tmp/${dev//'/'/'_'}"
smartctl -A "$dev" | tee "$attrib-old.txt" | awk '
	/Reallocated_Sector_Ct|Current_Pending_Sector|Offline_Uncorrectable/ {
		if ($10 != "0") {
			$0 = "\033[1;93m" $0 "\033[m"
		} else {
			$0 = "\033[1;92m" $0 "\033[m"
		}
	}
	{ print }
'

info "Starting ATA secure erase of $dev (expected duration: $minutes)"
password="foo"
hdparm --security-set-pass "$password" "$dev"

# For some reason, the 1st ERASE sometimes immediately fails. 
for tries in {1..3}; do
	# Give some old disks a moment.
	# (On Seagate 160GB disks, immediately started erase seems to fail instantly.)
	sleep 2

	Tstart=$(date +%s)
	info "Started at $(date -d "@$Tstart" +%c)"

	flock -F "$dev" hdparm --security-erase$enhanced "$password" "$dev" & hdparm_pid=$!
	countdown "$minutes" & countdown_pid=$!
	trap 'kill $hdparm_pid $countdown_pid; exit' INT EXIT

	wait -n -p pid; r=$?
	if [[ $pid == $hdparm_pid ]]; then
		kill $countdown_pid
		echo "hdparm finished with $r"
	else
		echo "hdparm still running..."
		wait -n; r=$?
		echo "hdparm now finished with $r"
	fi
	trap - INT EXIT

	Tfinish=$(date +%s)
	info "Finished at $(date -d "@$Tfinish" +%c), in $(interval $[Tfinish-Tstart])"
	if (( Tfinish-Tstart > 1 )); then
		notify-send "Security erase of $dev done in $(interval $[Tfinish-Tstart])"
		break
	fi
	break
done

if ! < $dev; then
	die "Device lost."
fi

# Verify that security has been disabled after erase.
info "ATA security information:"
hdparm -I "$dev" | awk '/^Security:/{x=1} /^[^S\t]/{x=0} x{print}'

info "SMART attributes after erase:"
smartctl -A "$dev" > "$attrib-new.txt"
! colordiff -U999 "$attrib-old.txt" "$attrib-new.txt" || cat "$attrib-new.txt"

info "Initializing a new GPT on $dev"
sgdisk -o "$dev"

# Give udev and udisks some time to probe the new GPT before spindown.
sleep 2

info "Putting $dev to standby"
scsi_stop "$dev"
