#!/bin/bash
###################Wipe (optional)
DEVICE=${1}
wipedelay=20

# Required packages
REQUIRED_PACKAGES=("hdparm" "dialog" "dc3dd" "util-linux")

# Check for missing packages
check_missing_packages()
{
	for package in "${REQUIRED_PACKAGES[@]}"; do
		if ! dpkg -s "${package}" >/dev/null 2>&1; then
			echo "Wipe script requires the following packages:"
			for p in "${REQUIRED_PACKAGES[@]}"; do
				echo "  ${p}"
			done
			exit 1
		fi
	done
}

# Get device from the user if not specified or invalid
get_device()
{
	if [ -z "$DEVICE" ] || [ ! -b "$DEVICE" ]; then
		# Create a list of available devices
		W=()
		while read -r line; do
			dev=$(echo $line | cut -f1 -d" ")
			rest=$(echo $line | cut -f2- -d" " | tr -s " ")
			W+=("/dev/${dev}" "${rest}")
		done < <(lsblk -l -oname,size,model,type | grep -e disk)

		# Display device selection menu
		DEVICE=$(dialog --backtitle "CERTBw - SecureErase" --title "Available Devices" --menu "Which disk should be wiped?" 24 80 17 "${W[@]}" 3>&2 2>&1 1>&3)
	fi
}

# cleanup function to unset the ATA Password if execution  gets interrupted
cleanup()
{
	echo
	echo "==WIPE : Removing ATA password due to user interruption..."
	hdparm --user-master u --security-disable certbw "${DEVICE}"
	echo "==WIPE : ATA password removed."
	exit 1
}

# Display warning and countdown
display_warning()
{
	dialog --backtitle "CERTBw - SecureErase" --defaultno --cancel-label "Cancel" --colors --title "\Z1!WARNING!\Zn" --pause "\n\Z1The device ${DEVICE} will be completely erased!\Zn\n\nThe SecureErase process must not be interrupted, as this will lock the device, and it will need to be manually unlocked afterward.\n\n\nThe process will automatically continue after the countdown expires.\n\nTo cancel the DiskWipe, you can:\n        \Z4Select \"Cancel\"\n        Press \"ESC\"\n        Press \"CTRL + C\"\n        Turn off the computer\Zn" 24 80 ${wipedelay}
	if [ "$?" -gt 0 ]; then
		echo "==WIPE : Wipe was canceled by the user."
		sleep 1
		read -p "Press [ENTER] key for Shell..."
		exit 1
	fi
}

# Securely erase the device
secure_erase()
{
	if hdparm -I "${DEVICE}" | grep supported | grep -q erase; then
		echo "==WIPE : Secure Erase is supported by ${DEVICE}"
		if ! (hdparm -I "${DEVICE}" | grep not | grep -q frozen); then
			echo "==WIPE : The device ${DEVICE} is frozen"
			echo "==WIPE : The notebook will now be put to sleep for 10 seconds."
			echo "==WIPE : Do not turn off the notebook."
			sleep 5
			rtcwake -s 10 -m mem
			echo "==WIPE : The notebook has woken up. Checking the status of ${DEVICE}."
		fi
		if hdparm -I "${DEVICE}" | grep not | grep -q frozen; then
			echo "==WIPE : The device ${DEVICE} is 'not frozen'"
			echo
			echo "==WIPE : A temporary ATA password (certbw) must be set for SecureErase."
			echo "==WIPE : If the SecureErase process is interrupted, the disk will be unusable until manually unlocked."
			echo "==WIPE : Do not turn off the notebook."
			sleep 5
			# Set a trap to catch SIGINT and call the cleanup function
			trap 'cleanup' SIGINT
			# Set ATA password
			hdparm --user-master u --security-set-pass certbw "${DEVICE}"
			# Issue Secure Erase command
			hdparm --user-master u --security-erase certbw "${DEVICE}"
			# Remove the trap after the Secure Erase is completed
			trap - SIGINT
		else
			# Normal wipe because unfreeze didn't work
			echo "==WIPE : The device could not be unfrozen."
			echo "==WIPE : The device ${DEVICE} will be overwritten."
			/usr/bin/dc3dd wipe="${DEVICE}"
		fi
	else
		# Normal wipe because Secure Erase is not supported
		echo "==WIPE : Secure Erase is NOT supported."
		echo "==WIPE : The device ${DEVICE} will be overwritten."
		/usr/bin/dc3dd wipe="${DEVICE}"
	fi
}

check_missing_packages
get_device
if [ ! -b "${DEVICE}" ]; then
	echo "==WIPE : Kein gültiges BLOCK-Device ausgewählt."
	sleep 1
	read -p "Press [ENTER] key for Shell..."
	exit 1
fi
display_warning
secure_erase