pencheck

#!/bin/bash
#
# pencheck - create, copy and check random files to test storage integrity
#
# Copyright (C) 2020 Rodrigo Silva (MestreLion) <linux@rodrigosilva.com>
# License: GPLv3 or later, at your choice. See <http://www.gnu.org/licenses/gpl>
#------------------------------------------------------------------------------

set -Eeuo pipefail  # exit on any error
trap '>&2 echo "non-success status: line $LINENO, status $?: $BASH_COMMAND"' ERR

#------------------------------------------------------------------------------

count=10
size=1G
bs=1Mi
verbose=1

#------------------------------------------------------------------------------

myname=${0##*/}
mydir=$(dirname "$(readlink -f "$0")")

# checksum command must have --binary, --check and --ignore-missing
# TODO: use xxh*sum when (and if) it adds --ignore-missing
# https://github.com/Cyan4973/xxHash/issues/811
checksum=md5sum
dirbase=$myname
verify=${myname}-verify.sh
checkfile=${myname}-${checksum}.txt

#------------------------------------------------------------------------------

bold()    { tput bold; printf '%s' "$@"; tput sgr0; }
red()     { tput setaf 1; bold "$@"; }
green()   { tput setaf 2; bold "$@"; }
fatal()   { if (($#)); then { red "$@"; echo; } >&2; fi; exit 1; }
message() { if (($# && verbose)); then green '* ' "$@"; echo; fi; }
argerr()  { printf "%s: %s\n" "$myname" "${1:-error}" >&2; usage 1; }
invalid() { argerr "invalid ${2:-option}: ${1:-}"; }
missing() { argerr "missing ${1:+$1 }argument${2:+ in $2}."; }
integer() { [[ "$1" ]] || missing "${3:-NUM}" "${2:-}"; [[ "$1" != *[!0-9]* ]] ||
            argerr "${2:+${3:-NUM} argument in $2 is }not an integer: $1"; }
sizefmt() { [[ "$1" ]] || missing "${3:-SIZE}" "${2:-}";
            local num=$(numfmt --from=auto -- "$1");
            [[ "$num" ]] && echo "$num" || return 1; }
exists()  { type "$@" >/dev/null 2>&1; }
si_iec()  { numfmt --to=${2:-si} --suffix 'B' "$1"; }
si()      { si_iec "$1" si; }
iec()     { si_iec "$1" iec-i; }
bytesfmt(){ numfmt --grouping "$@"; }
sizemsg() { echo "$(${2:-si} $1) ${3:+$3 }($(bytesfmt $1) bytes)"; }

# An extremely fast way of generating pseudo-random garbage test data
# Use a fast openssl algo to encrypt /dev/zero, MUCH faster than /dev/urandom or shred
# See also ./reference/randspeed
garbage() {
	# User-defined --block-size does not apply here when dealing with /dev/zero
	if exists pv; then
		pv "${pvopts[@]}" --buffer-size 1M --stop-at-size --size "$size"
	else
		dd "${ddopts[@]}" bs=1M iflag=count_bytes,fullblock count="$size"
	fi < /dev/zero |
	openssl enc -aes-128-ctr -pbkdf2 -nosalt -pass \
		"pass:$(dd status=none if=/dev/urandom bs=16 count=1 | base64)"
}
freespace() {
	local path=$1
	local required=$2
	local available=$(df --block-size=1 -- "$path" | awk 'NR==2{print $4}')
	if ((available < required)); then
		fatal "not enough space on $path: " \
			$(si "$available") " available, " \
			$(si "$required")  " required."
	fi
}

#------------------------------------------------------------------------------

usage() {
	if [[ "${1:-}" ]] ; then exec >&2; fi
	cat <<-USAGE
	Usage: $myname [options] MOUNTDIR
	USAGE
	if [[ "${1:-}" ]] ; then
		cat <<- USAGE
		Try '$myname --help' for more information.
		USAGE
		exit 1
	fi
	cat <<-USAGE

	Create, copy and check random files to test storage integrity

	Options:
	  -h|--help    - show this page.
	  -q|--quiet   - suppress informative messages.

	  -c|--count NUM       - Number of test files. [Default: $count]
	  -s|--size SIZE       - Size of each test file. [Default: $size]
	  -d|--dirbasename DIR - Directory basename under MOUNTDIR root to copy
	                           the test files. [Default: $dirbase]
	  -b|--block-size SIZE - Use SIZE as block size on write operations
	                           when applicable. [Default: $bs]

	Create files with random data in MOUNTDIR/$dirbase, MOUNTDIR usually
	being the mountpoint directory of a USB Thumb Drive, and verify their
	integrity using the '$checksum' utility. Also creates the checksum file
	'$checkfile' and a companion script '$verify' so
	integrity can easily be tested again in the future.

	It is suggested to adjust size and count so test files occupy almost
	all of the storage free space, so for an empty 16GB Thumb Drive $myname
	should preferably be used with '--size 1G --count 15'. Test files can
	later be selectively deleted to free just enough space for your own data,
	keeping the Thumb Drive always almost full to maximize integrity test
	coverage.

	All SIZE units are parsed through 'numfmt' auto mode, so SI suffixes such
	as K, M, G are allowed. '1K' means 1000; use 'Ki' (and 'Mi', 'Gi', etc)
	for 1024-based units.

	Copyright (C) 2020 Rodrigo Silva (MestreLion) <linux@rodrigosilva.com>
	License: GPLv3 or later. See <http://www.gnu.org/licenses/gpl.html>
	USAGE
	exit 0
}

# Pre-parse for -h|--help, ignoring if after '--'
for arg in "$@"; do
	if [[ "$arg" == '--' ]]; then break; fi
	if [[ "$arg" == "-h" || "$arg" == "--help" ]]; then usage; fi
done
args=()
while (($#)); do
	case "$1" in
	-q|--quiet      ) verbose=0;;
	-s|--size       ) shift; size=${1:-}   ;;
	-c|--count      ) shift; count=${1:-}  ;;
	-b|--block-size ) shift; bs=${1:-}     ;;
	-d|--dirbasename) shift; dirbase=${1:-};;
	--size=*        ) size=${1#*=}   ;;
	--count=*       ) count=${1#*=}  ;;
	--block-size=*  ) bs=${1#*=}     ;;
	--dirbasename=* ) dirbase=${1#*=};;
	--) shift; break;;
	-*) invalid "$1";;
	* ) args+=( "$1" );;
	esac
	shift || break
done
args+=("$@")

#------------------------------------------------------------------------------

if ((${#args[@]} < 1)); then missing "MOUNTDIR"; fi
if ((${#args[@]} > 1)); then invalid "${args[1]}" argument; fi

mountdir=${args[0]%/}
[[ -d "$mountdir" ]] || argerr "MOUNTDIR is not a valid directory: $mountdir"

dirbase=${dirbase%/}  # will also consider '/' as missing, which is sensible
[[ "$dirbase" ]] || missing 'DIR' '--dirbase'

outdir=$mountdir/$dirbase

# integer automatically checks for missing
integer "$count" '--count'

# sizefmt converts to bytes, and also checks missing values and valid integers
# Being in a subshell requires outer || exit
size=$(sizefmt "$size" '--size') || exit
bs=$(  sizefmt "$bs"   '--block-size') || exit

# account for sha256sum.txt and verify script
totalsize=$((count * (size + 82) + 1000))
message "Use $(sizemsg $bs iec 'block size')"
freespace "$mountdir" "$totalsize"

if ((verbose)); then
	checksumopts=()
	rsyncopts=(--info=progress2)
	ddopts=(status=progress)
	pvopts=()
	exec 3>&1
else
	checksumopts=(--quiet)
	rsyncopts=()
	ddopts=(status=none)
	pvopts=(--quiet)
	exec 3>/dev/null
fi

#------------------------------------------------------------------------------

message "Create $count test files of $(sizemsg $size) each"

tmpdir=$(mktemp --directory) || fatal "could not create temp dir"
trap 'rm -rf -- "$tmpdir"' EXIT
freespace "$tmpdir" "$totalsize"

workdir=$tmpdir/$dirbase
script=$workdir/$verify
mkdir --parents -- "$workdir"
for i in $(seq --equal-width "$count"); do
	garbage "$size" > "$workdir"/"$myname"-"$i".bin
done

message "Calculate test files checksum using '$checksum'"
cd -- "$workdir"
"$checksum" --binary *.bin | tee "$workdir"/"$checkfile" >&3
cat > "$script" <<-EOF
	#!/bin/sh
	# Check data consistency of test files
	# Created by pencheck: https://github.com/MestreLion/scripts
	echo "Checking data consistency of test files, this might take a while..."
	cd -- "\$(dirname "\$(readlink -f "\$0")")"
	$checksum --check --ignore-missing "\$@" -- ./$checkfile
EOF
chmod +x -- "$script"

message "Verify test files"
sync -- "$workdir"/*
sh "$script" "${checksumopts[@]}"
cd - &>/dev/null

message "Transfer test data to: $outdir"
# test again, just in case $mountdir is the same filesystem as $tmpdir
freespace "$mountdir" "$totalsize"
# Copy script and checksum file first
mkdir --parents -- "$outdir"
mv -t "$outdir" -- "$script" "$workdir"/"$checkfile"
# If only Debian's rsync had --drop-cache :-(
rsync "${rsyncopts[@]}" --recursive -- "$workdir" "$mountdir"  # NOT $outdir!

message "Wait for files to be actually written on device"
sync --file-system -- "$outdir"/*  # intentionally sync whole fs

message "Test storage"
# Some filesystems may be automounted as noexec (e.g. FAT), hence invoking sh
sh "$outdir"/"$verify" "${checksumopts[@]}"