#!/bin/bash

# Copyright 2018, 2019 eomanis
# 
# This file is part of disk-test.
#
# disk-test is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation.
#
# disk-test is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with disk-test.  If not, see <http://www.gnu.org/licenses/>.

# TODO Default settings via config file(s), such as /etc/disk-test.conf
# TODO Bash completion

# Set shell options
set -o errexit
set -o noclobber
set -o pipefail
set -o nounset

# Semantic versioning
declare -r versionMajor=0
declare -r versionMinor=2
declare -r versionPatch=3
declare -r versionLabel=""

getVersion () {
	echo -n "${versionMajor}.${versionMinor}.${versionPatch}"
	test "$versionLabel" != "" && echo "-$versionLabel" || echo ""
}

# Make sure the required programs are available
allRequiredProgramsPresent=true
type base64 &> /dev/null   || { echo "ERROR  Required program \"base64\" is not available" >&2;   allRequiredProgramsPresent=false; }
type blockdev &> /dev/null || { echo "ERROR  Required program \"blockdev\" is not available" >&2; allRequiredProgramsPresent=false; }
type cmp &> /dev/null      || { echo "ERROR  Required program \"cmp\" is not available" >&2;      allRequiredProgramsPresent=false; }
type dd &> /dev/null       || { echo "ERROR  Required program \"dd\" is not available" >&2;       allRequiredProgramsPresent=false; }
type grep &> /dev/null     || { echo "ERROR  Required program \"grep\" is not available" >&2;     allRequiredProgramsPresent=false; }
type openssl &> /dev/null  || { echo "ERROR  Required program \"openssl\" is not available" >&2;  allRequiredProgramsPresent=false; }
type pv &> /dev/null       || { echo "ERROR  Required program \"pv\" is not available" >&2;       allRequiredProgramsPresent=false; }
type realpath &> /dev/null || { echo "ERROR  Required program \"realpath\" is not available" >&2; allRequiredProgramsPresent=false; }
type sed &> /dev/null      || { echo "ERROR  Required program \"sed\" is not available" >&2;      allRequiredProgramsPresent=false; }
type stat &> /dev/null     || { echo "ERROR  Required program \"stat\" is not available" >&2;     allRequiredProgramsPresent=false; }
type tr &> /dev/null       || { echo "ERROR  Required program \"tr\" is not available" >&2;       allRequiredProgramsPresent=false; }
$allRequiredProgramsPresent || exit 1
unset allRequiredProgramsPresent

# No arguments, or --help as single argument
if test $# -eq 0 || { test $# -eq 1 && test "$1" = "--help"; }; then
	echo -n "disk-test "; getVersion
	echo -n "Block device read-write test utility written in bash

Usage:
disk-test
  [-a|--actions actions]
  [-b|--block-size blockSize]
  [-c|--count-bytes numberOfBytes]
  [-s|--skip-bytes numberOfBytes]
  [--]
  deviceOrFile
disk-test [--help]

Tests if a device can write and correctly read arbitrary data over its
entire reported size, or over a specified range.

For more information see manual page disk-test(1).
"
	exit 1
fi

# Declare defaults
declare -r actionsDefault="r"
declare -r blockSizeDefault="8M"
declare -r skipBytesDefault=0

# Set defaults
actions="$actionsDefault"
blockSize="$blockSizeDefault"
skipBytes="$skipBytesDefault"
# countBytes will be calculated/validated later

# Parse the arguments
trailingArgs=false
while test $# -gt 0; do
	if ! $trailingArgs && test "$1" = "--"; then
		trailingArgs=true; shift
	elif $trailingArgs || ! { test ${#1} -ge 1 && test "${1:0:1}" = "-"; }; then
		# Trailing arguments section reached, or the argument does not
		# start with "-": Must be the device
		if ! test -z ${device+x}; then
			echo "ERROR  Multiple devices specified: \"$device\", \"$1\"" >&2
			exit 1
		fi
		device="$1"; shift
	elif test "$1" = "-a" || test "$1" = "--actions"; then
		test $# -ge 2 || { echo "ERROR  No value given for argument \"$1\"" >&2; exit 1; }
		shift
		actions="$1"; shift
	elif test "$1" = "-b" || test "$1" = "--block-size"; then
		test $# -ge 2 || { echo "ERROR  No value given for argument \"$1\"" >&2; exit 1; }
		shift
		blockSize="$1"; shift
	elif test "$1" = "-s" || test "$1" = "--skip-bytes"; then
		test $# -ge 2 || { echo "ERROR  No value given for argument \"$1\"" >&2; exit 1; }
		shift
		skipBytes="$1"; shift
	elif test "$1" = "-c" || test "$1" = "--count-bytes"; then
		test $# -ge 2 || { echo "ERROR  No value given for argument \"$1\"" >&2; exit 1; }
		shift
		countBytes="$1"; shift
	else
		echo "ERROR  Unknown argument \"$1\"" >&2; exit 1
	fi
done
unset trailingArgs

# Validate the device
if test -z ${device+x}; then
	echo "ERROR  No device specified" >&2
	exit 1
fi
if test -L "$device"; then
	echo -n " INFO  Device \"$device\" is a symbolic link to " >&2
	device="$(realpath --physical --canonicalize-missing "$device")"
	echo "\"$device\"" >&2
fi
if ! test -b "$device" && ! test -f "$device"; then
	echo "ERROR  Neither a block device nor a regular file: \"$device\"" >&2
	exit 1
fi

# Determine the device's or file's size
declare -i deviceSizeTmp
if test -b "$device"; then
	if ! deviceSizeTmp="$(blockdev --getsize64 "$device")" || test "$deviceSizeTmp" = ""; then
		echo "ERROR  Failed to determine the size of device \"$device\"" >&2
		exit 1
	fi
else
	if ! deviceSizeTmp="$(stat -c "%s" "$device")" || test "$deviceSizeTmp" = ""; then
		echo "ERROR  Failed to determine the size of file \"$device\"" >&2
		exit 1
	fi
fi
declare -r -i deviceSize="$deviceSizeTmp"
unset deviceSizeTmp

# Validate --actions
if ! echo "$actions" | grep -E '^[wzvr]*$' &> /dev/null; then
	echo "ERROR  Invalid character(s) in actions string \"$actions\"; allowed characters: wzvr" >&2
	exit 1
fi
if echo "$actions" | grep "v" &> /dev/null \
	&& ! echo "$actions" | sed -re 's/^([^v]*)v.*$/\1/;s/[^wz]+//g' | grep -E '^[wz]+$' &> /dev/null; then
	echo "ERROR  Invalid sequence of actions \"$actions\": Read-and-verify action (v) without preceding write action (wz)" >&2
	exit 1
fi

# Validate --block-size (basic validation only)
if test "$blockSize" = "" || ! echo "$blockSize" | grep -E '^[1-9]' &> /dev/null; then
	# The block size must not be empty, and must start with a digit 1..9
	echo "ERROR  Invalid value for --block-size: \"$blockSize\"" >&2
	exit 1
fi

# Validate --skip-bytes
if ! echo "$skipBytes" | grep -E '^(0|[1-9][0-9]*)$' &> /dev/null; then
	echo "ERROR  Value for --skip-bytes must be an integer >= 0, was \"$skipBytes\"" >&2
	exit 1
fi
if test "$skipBytes" -gt "$deviceSize"; then
	echo "ERROR  Value for --skip-bytes must not exceed the device's size in bytes ($deviceSize), was $skipBytes" >&2
	exit 1
fi

# Calculate (if required) and validate --count-bytes
countBytesMax="$(( "$deviceSize" - "$skipBytes" ))"
if test -z ${countBytes+x}; then
	countBytes="$countBytesMax"
fi
if ! echo "$countBytes" | grep -E '^(0|[1-9][0-9]*)$' &> /dev/null; then
	echo "ERROR  Value for --count-bytes must be an integer >= 0, was \"$countBytes\"" >&2
	exit 1
fi
if test "$countBytes" -gt "$countBytesMax"; then
	echo "ERROR  Value for --count-bytes must not exceed the number of bytes from --skip-bytes to the end of the device ($countBytesMax), was $countBytes" >&2
	exit 1
fi
unset countBytesMax

# Print a summary
echo " INFO  Testing block device or file: \"$device\"" >&2
echo " INFO  Performing ${#actions} action(s): \"$actions\"" >&2
test "$skipBytes" -ne "$skipBytesDefault" && echo " INFO  Skipping the first $skipBytes byte(s)" >&2
test "$countBytes" -ne "$deviceSize" && echo " INFO  Testing a range of $countBytes byte(s)" >&2
test "$blockSize" != "$blockSizeDefault" && echo " INFO  Using custom block size \"$blockSize\"" >&2

# Handle some corner cases up front
if test "$countBytes" -eq 0; then
	echo " WARN  Range to test is 0 bytes wide, nothing to do" >&2
	exit 0
fi
if test ${#actions} -eq 0; then
	echo " WARN  No actions specified, nothing to do" >&2
	exit 0
fi

# Perform actions
declare -i actionIndex=0
declare action
declare -i stepreturn
declare seed
declare mostRecentWriteAction
while test "$actionIndex" -lt ${#actions}; do
	action="${actions:$actionIndex:1}"
	echo -n " INFO  Action $(( actionIndex + 1 ))/${#actions} (${action}) on device \"$device\": " >&2
	echo "$actions" | sed -re 's/^(.{'$actionIndex'})./\1'"$(echo "$action" | tr 'a-z' 'A-Z')"'/' >&2
	
	if test "$action" = "w"; then
		mostRecentWriteAction="$action"
		if ! seed="$(dd if=/dev/urandom iflag=fullblock bs=128 count=1 2>/dev/null | base64 --wrap 0)"; then
			echo "ERROR  Could not generate random seed for pseudo-random data generation" >&2
			exit 1
		fi
		dd if=/dev/zero iflag=fullblock,count_bytes count="$countBytes" ibs="$blockSize" 2> /dev/null \
			| openssl enc -aes-256-ctr -pass 'pass:'"$seed" -nosalt -pbkdf2 -iter 1 \
			| pv --progress --timer --eta --rate --bytes --size "$countBytes" --name " INFO  Writing pseudo-random" \
			| dd of="$device" conv=notrunc oflag=direct,seek_bytes seek="$skipBytes" obs="$blockSize" 2> /dev/null \
			&& stepreturn=$? || stepreturn=$?
		if test $stepreturn != 0; then
			echo "ERROR  Writing pseudo-random data failed with return value $stepreturn" >&2
			exit 2
		fi
	elif test "$action" = "z"; then
		mostRecentWriteAction="$action"
		dd if=/dev/zero iflag=fullblock,count_bytes count="$countBytes" ibs="$blockSize" 2> /dev/null \
			| pv --progress --timer --eta --rate --bytes --size "$countBytes" --name " INFO  Writing zeros" \
			| dd of="$device" conv=notrunc oflag=direct,seek_bytes seek="$skipBytes" obs="$blockSize" 2> /dev/null \
			&& stepreturn=$? || stepreturn=$?
		if test $stepreturn != 0; then
			echo "ERROR  Writing zeros failed with return value $stepreturn" >&2
			exit 2
		fi
	elif test "$action" = "v"; then
		if test "$mostRecentWriteAction" = "w"; then
			dd if=/dev/zero iflag=fullblock,count_bytes count="$countBytes" ibs="$blockSize" 2> /dev/null \
				| openssl enc -aes-256-ctr -pass 'pass:'"$seed" -nosalt -pbkdf2 -iter 1 \
				| pv --progress --timer --eta --rate --bytes --size "$countBytes" --name " INFO  Reading and comparing" \
				| cmp --quiet --ignore-initial="${skipBytes}:0" --bytes="$countBytes" -- "$device" - \
				&& stepreturn=$? || stepreturn=$?
			if test $stepreturn != 0; then
				echo "ERROR  Reading and comparing pseudo-random data failed with return value $stepreturn" >&2
				exit 2
			fi
		elif test "$mostRecentWriteAction" = "z"; then
			dd if=/dev/zero iflag=fullblock,count_bytes count="$countBytes" ibs="$blockSize" 2> /dev/null \
				| pv --progress --timer --eta --rate --bytes --size "$countBytes" --name " INFO  Reading and comparing" \
				| cmp --quiet --ignore-initial="${skipBytes}:0" --bytes="$countBytes" -- "$device" - \
				&& stepreturn=$? || stepreturn=$?
			if test $stepreturn != 0; then
				echo "ERROR  Reading and comparing zeros failed with return value $stepreturn" >&2
				exit 2
			fi
		else
			echo "ERROR  Preceding write action is not known" >&2
			exit 1
		fi
	elif test "$action" = "r"; then
		dd if="$device" iflag=fullblock,skip_bytes,count_bytes skip="$skipBytes" count="$countBytes" ibs="$blockSize" 2> /dev/null \
			| pv --progress --timer --eta --rate --bytes --size "$countBytes" --name " INFO  Reading" \
			| dd of=/dev/null conv=notrunc obs="$blockSize" 2> /dev/null \
			&& stepreturn=$? || stepreturn=$?
		if test $stepreturn != 0; then
			echo "ERROR  Reading failed with return value $stepreturn" >&2
			exit 2
		fi
	else
		echo "ERROR  Unknown action \"$action\" (what the hell, argument validation should have caught that)" >&2
		exit 1
	fi
	actionIndex=$(( actionIndex + 1 ))
done
echo " INFO  Finished, no errors: \"$device\"" >&2