#!/usr/bin/env bash

NAME="bdpurge"
CODENAME="bdpurge"
COPYRIGHT="Copyright (C) 2016 Nathan Shearer"
LICENSE="GNU General Public License 2.0"
VERSION="2.3.0.0"

# \brief Analyze a block device and display relevant information
# \param $1 The device
function bdpurge_analyze
{
	local DEVICE="$1"
	
	if [ ! -b "$DEVICE" ]; then
		printf "%s\n" "error: \"$DEVICE\" is not a block device"
		return 1
	fi
	if ! $FORCE && mount | grep "$DEVICE" >/dev/null; then
		printf "%s\n" "error: \"$DEVICE\" is currently mounted"
		return 2
	fi
	
	local DEVICE_MODEL=$(hdparm -I "$DEVICE" 2>/dev/null | grep --text 'Model Number' | cut -d ':' -f 2 | sed -r -e 's/ *(.*?) */\1/')
	local DEVICE_SERIAL=$(hdparm -I "$DEVICE" 2>/dev/null | grep --text 'Serial Number' | cut -d ':' -f 2 | sed -r -e 's/ *(.*?) */\1/')
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	
	printf "%s\n" "Device: $DEVICE"
	if [ "$DEVICE_MODEL" != "" ]; then
		printf "%s\n" "  Model Number: $DEVICE_MODEL"
	fi
	if [ "$DEVICE_SERIAL" != "" ]; then
		printf "%s\n" "  Serial Number: $DEVICE_SERIAL"
	fi
	printf "%s\n" "  Size: $DEVICE_SIZE Bytes"
}

# \brief Ensures dependencies are present
# \param $@ The dependencies to check for
function bdpurge_check_dependencies
{
	for TOOL in "$@"; do
		if ! type "$TOOL" >/dev/null 2>/dev/null; then
			echo "$CODENAME: \"$TOOL\" is required for this application to work correctly." >&2
			exit
		fi
	done
}

# \brief Sends an email
# \param $1 The e-mail address
# \param $2 The subject
# \param $3 The file containing the message
function bdpurge_email
{
	if [ $# -eq 3 -a "$1" != "" ]; then
		cat "$3" | mail -s "$2" "$1"
	fi
}

# \brief Write 4 passes of 0xaa, 0x55, 0xff, and 0x00
# \param $1 The device
function bdpurge_erase_badblocks
{
	local DEVICE="$1"
	
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	local BLOCK_SIZE=$(bdpurge_gcd $((2**20)) $DEVICE_SIZE)
	
	if $PRETEND; then
		printf "%s\n" "badblocks -b $BLOCK_SIZE -o /dev/null -p 0 -s -w $DEVICE"
	else
		badblocks -b $BLOCK_SIZE -o /dev/null -p 0 -s -w $DEVICE 2>&1 | tee "$TMP/badblocks.output"
		BADBLOCKS_STATUS=$?
		local NOW=$(date --rfc-3339=seconds)
		if [ $BADBLOCKS_STATUS -eq 0 ]; then
			printf "%s\n\n" "The badblocks erasure of \"$DEVICE\" completed successfully." >"$TMP"/complete.email
			bdpurge_analyze "$DEVICE" >>"$TMP"/complete.email
			cat "$TMP/badblocks.output" | sed -r -e 's/^.*\r(.*)$/\1/' >"$TMP/badblocks.output2"
			bdpurge_hail -n 100 "$TMP/badblocks.output2" >>"$TMP"/complete.email
			bdpurge_email $EMAIL "$HOSTNAME: bdpurge: badblocks of \"$DEVICE\" completed successfully at $NOW" "$TMP"/complete.email
		else
			printf "%s\n\n" "The badblocks erasure of \"$DEVICE\" failed." >"$TMP"/complete.email
			bdpurge_analyze "$DEVICE" >>"$TMP"/complete.email
			cat "$TMP/badblocks.output" | sed -r -e 's/^.*\r(.*)$/\1/' >"$TMP/badblocks.output2"
			bdpurge_hail -n 100 "$TMP/badblocks.output2" >>"$TMP"/complete.email
			bdpurge_email $EMAIL "$HOSTNAME: bdpurge: badblocks of \"$DEVICE\" failed at $NOW" "$TMP"/complete.email
			return 1
		fi
	fi
}

# \brief Erase data at both ends of a block device
# \param $1 The device
# \param $2 The amount of data to erase at each end
function bdpurge_erase_quick
{
	local DEVICE="$1"
	local PURGE_SIZE="$2"
	
	echo -n "Erasing the first and last $PURGE_SIZE bytes on \"$DEVICE\"..."
	
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	
	if [ $(($PURGE_SIZE*2)) -ge "$DEVICE_SIZE" ]; then
		if $PRETEND; then
			echo
			echo dd status=none if=/dev/zero bs=512 count=$(($DEVICE_SIZE/512)) of="$DEVICE" >/dev/null 2>/dev/null
		else
			dd status=none if=/dev/zero bs=512 count=$(($DEVICE_SIZE/512)) of="$DEVICE" >/dev/null 2>/dev/null
		fi
	else
		if $PRETEND; then
			echo
			echo dd status=none if=/dev/zero bs=512 count=$(($PURGE_SIZE/512)) of="$DEVICE"
			echo dd status=none if=/dev/zero bs=512 count=$(($PURGE_SIZE/512)) of="$DEVICE" seek=$(($(($DEVICE_SIZE/512))-$(($PURGE_SIZE/512))))
		else
			dd status=none if=/dev/zero bs=512 count=$(($PURGE_SIZE/512)) of="$DEVICE"
			dd status=none if=/dev/zero bs=512 count=$(($PURGE_SIZE/512)) of="$DEVICE" seek=$(($(($DEVICE_SIZE/512))-$(($PURGE_SIZE/512))))
		fi
	fi
	
	if ! $PRETEND; then
		echo " done."
	fi
}

# \brief Perform a quick erase of each partition and the entire block device
# \param $1 The device
# \param $2 The amount of data to erase at each end
function bdpurge_erase_quickr
{
	local DEVICE="$1"
	local PURGE_SIZE="$2"
	
	for PARTITION in "$DEVICE"[0-9]* "$DEVICE"p[0-9]*; do
		if [ "$PARTITION" = "${DEVICE}[0-9]*" ]; then continue; fi
		if [ "$PARTITION" = "${DEVICE}p[0-9]*" ]; then continue; fi
		bdpurge_erase_quick "$PARTITION" "$PURGE_SIZE"
	done
	bdpurge_erase_quick "$DEVICE" "$PURGE_SIZE"
}

# \brief Erase the entire device with the trim command
# \param $1 The device
function bdpurge_erase_trim
{
	local DEVICE="$1"
	
	echo -n "Erasing \"$DEVICE\" with the trim command... "
	
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	
	if $PRETEND; then
		echo
		echo blkdiscard "$DEVICE"
	else
		blkdiscard "$DEVICE" || return
		echo "done."
	fi
}

# \param $1 The device
function bdpurge_erase_urandom
{
	local DEVICE="$1"
	
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	local BLOCK_SIZE=$(bdpurge_gcd $((2**20)) $DEVICE_SIZE)
	local BLOCK_COUNT=$(($DEVICE_SIZE/$BLOCK_SIZE))
	
	if $PRETEND; then
		printf "%s\n" "dd if=/dev/urandom bs=$BLOCK_SIZE conv=sync of=$DEVICE count=$BLOCK_COUNT"
	else
		dd if=/dev/urandom bs=$BLOCK_SIZE conv=sync of="$DEVICE" count=$BLOCK_COUNT
	fi
}

# \param $1 The device
function bdpurge_erase_zero
{
	local DEVICE="$1"
	
	local DEVICE_SIZE=$(blockdev --getsize64 "$DEVICE")
	local BLOCK_SIZE=$(bdpurge_gcd $((2**20)) $DEVICE_SIZE)
	local BLOCK_COUNT=$(($DEVICE_SIZE/$BLOCK_SIZE))
	
	if $PRETEND; then
		printf "%s\n" "dd if=/dev/zero bs=$BLOCK_SIZE conv=sync of=$DEVICE count=$BLOCK_COUNT"
	else
		dd if=/dev/zero bs=$BLOCK_SIZE conv=sync of="$DEVICE" count=$BLOCK_COUNT
	fi
}

# \brief Cleans up the environment and exits
# \param $1 The exit code
# \param $2 The exit message
#
# If DEBUG=true then temporary files are not deleted.
function bdpurge_exit
{
	if [ "$EXITING" = "true" ]; then return; fi
	EXITING=true
	local EXIT="$1"
	local MESSAGE="$2"
	if [ "$EXIT" = "" ]; then
		EXIT=0
	fi
	if [ "$MESSAGE" = "" ]; then
		MESSAGE="An unrecoverable error has occurred"
	fi
	if ! $DEBUG; then
		rm -rf --one-file-system "$TMP"
	else
		echo "Debug mode is enabled. Temporary files in \"$TMP\" will *not* be deleted."
	fi
	case $EXIT in
		0) exit;;
		*) echo "$CODENAME: $MESSAGE" >&2; exit $EXIT;;
	esac
}

# \brief Calculates the greatest common divisor of $1 and $2
# \param $1 The first number
# \param $2 The second number
function bdpurge_gcd
{
	local A=$1
	local B=$2
	local R=$2
	while [ $R -ne 0 ]; do
		R=$(( $A % $B ))
		A=$B
		B=$R
	done
	echo -n "$A"
}

# \brief Prints the first and last n lines from $1
# \param $1 The file to read lines from
function bdpurge_hail
{
	local HAIL_BYTES=false
	local HAIL_LINES=10
	
	while [ $# -ne 0 ]; do
		case "$1" in
			"-c"|"--bytes")
				HAIL_BYTES="$2"
				shift 2
				;;
			"-n"|"--lines")
				HAIL_LINES="$2"
				shift 2
				;;
			*) break;;
		esac
	done
	
	if [ $# -eq 0 ]; then
		echo "Usage:"
		echo "$0 [file]"
		return 1
	fi
	
	if [ "$HAIL_BYTES" != "false" ]; then
		if [ $(cat "$1" | wc -c) -gt $(( 2 * $HAIL_BYTES + 3 )) ]; then
			head -c "$HAIL_BYTES" "$1"
			echo -n "..."
			tail -c "$HAIL_BYTES" "$1"
		else
			cat "$1"
		fi
	else
		if [ $(cat "$1" | wc -l) -gt $(( 2 * $HAIL_LINES + 1 )) ]; then
			head -n "$HAIL_LINES" "$1"
			echo "..."
			tail -n "$HAIL_LINES" "$1"
		else
			cat "$1"
		fi
	fi
}

# \brief Displays the help and exits the program
function bdpurge_help
{
	#     01234567890123456789012345678901234567890123456789012345678901234567890123456789
	echo "Description:"
	echo "  Erase a block device"
	echo
	echo "Usage:"
	echo "  $CODENAME [options] BLOCK_DEVICE"
	echo
	echo "Options:"
	echo "  -e, --email mail@example.com"
	echo "    Send an email upon completion of the badblocks method."
	echo "  -f, --force"
	echo "    Erase the block device even if it is mounted."
	echo "  -h, --help"
	echo "    Display this help message and exit."
	echo "  -m, --method quickr"
	echo "    Set the method used to erase the block device. Default is quickr:"
	echo "      badblocks  Write 4 passes of 0xaa, 0x55, 0xff, and 0x00"
	echo "      quick      Erase both ends of the device"
	echo "      quickr     Erase both ends of each partition and the device"
	echo "      trim       Erase the entire device with the trim command"
	echo "      urandom    Write 1 pass with data from /dev/urandom"
	echo "      zero       Write 1 pass with data from /dev/zero"
	echo "  -n, --nice N"
	echo "    Sets the niceness to N (default 0)."
	echo "  -p, --pretend"
	echo "    Performs a dry run and prints out what purge commands would be performed."
	echo "    Data is not actually purged."
	echo "  -s, --size 16777216"
	echo "    The amount of data to erase at both ends of the block device."
	echo
	echo "Examples:"
	echo "  $CODENAME -r -p /dev/disk/by-id/ata-ST8000AS0002-1NA17Z_00000000"
	echo
	echo "Version:"
	echo "  $NAME $VERSION"
	echo "  $COPYRIGHT"
	echo "  Licensed under $LICENSE"
	exit
}

# \brief Erase data on a block device
# \param $1 The device
function bdpurge_main
{
	local DEVICE="$1"
	
	bdpurge_analyze "$DEVICE"
	local ANALYZE_STATUS=$?
	if [ $ANALYZE_STATUS -ne 0 ]; then
		return $ANALYZE_STATUS
	fi
	
	echo "Press CTRL+C to cancel"
	for SECONDS in $(seq 10 -1 1); do
		printf "\r${SECONDS} ... "
		sleep 1
	done
	printf "\r0 ... \n"
	
	case "$METHOD" in
		"badblocks")
			bdpurge_erase_badblocks "$DEVICE"
			;;
		"quick")
			bdpurge_erase_quick "$DEVICE" "$PURGE_SIZE"
			;;
		"quickr")
			bdpurge_erase_quickr "$DEVICE" "$PURGE_SIZE"
			;;
		"trim")
			bdpurge_erase_trim "$DEVICE" || exit
			;;
		"urandom")
			bdpurge_erase_urandom "$DEVICE"
			;;
		"zero")
			bdpurge_erase_zero "$DEVICE"
			;;
		*)
			echo "error: invalid method \"$METHOD\""
			return 1
			;;
	esac
	
	if $PRETEND; then
		echo partprobe "$DEVICE"
	else
		partprobe "$DEVICE"
	fi
}

#------------------------------------------------------------------------------
# default configuration

DEBUG=false
EMAIL=""
FORCE=false
METHOD=quickr
NICE=0
PRETEND=false
PURGE_SIZE=16777216
TMP="/tmp"

#------------------------------------------------------------------------------
# config files

if [ -r /etc/$CODENAME.conf ]; then
	. /etc/$CODENAME.conf
fi
if [ -r ~/.$CODENAME.conf ]; then
	. ~/.$CODENAME.conf
fi

#------------------------------------------------------------------------------
# command line arguments

if [ $# -eq 0 ]; then
	bdpurge_help
	exit 1
fi

THIS="$0"
while [ $# -ne 0 ]; do
	case "$1" in
		"-e"|"--email")
			EMAIL="$2"
			shift 2
			;;
		"-f"|"--force")
			FORCE=true
			shift
			;;
		"-h"|"--help")
			bdpurge_help
			exit
			;;
		"-m"|"--method")
			METHOD="$2"
			shift 2
			;;
		"-n"|"--nice")
			NICE="$2"
			shift 2
			;;
		"-p"|"--pretend")
			PRETEND=true
			shift
			;;
		"-s"|"--size")
			PURGE_SIZE="$2"
			if [ $(($(($PURGE_SIZE/512))*512)) -ne "$PURGE_SIZE" ]; then
				PURGE_SIZE=$(($(($PURGE_SIZE/512))*512))
				echo "warning: The purge size is not a multiple of 512, rounding down from $2 to $PURGE_SIZE" 
			fi
			shift 2
			;;
		*)
			break;;
	esac
done

if [ $# -eq 0 ]; then
	echo "error: Missing BLOCK_DEVICE argument."
	exit 1
fi
if [ $# -gt 1 ]; then
	bdpurge_help
	exit 1
fi
DEVICE="$1"

#------------------------------------------------------------------------------
# prepare environment

bdpurge_check_dependencies blockdev dd partprobe sleep
trap bdpurge_exit EXIT SIGHUP SIGINT SIGQUIT SIGABRT SIGKILL SIGTERM
TMP="$TMP/$CODENAME.$$"
mkdir -m 0700 -p "$TMP"
renice $NICE $$ >/dev/null

#------------------------------------------------------------------------------
# begin execution

bdpurge_main "$1"