#!/bin/bash
#
# flash-disk-check
#
# Check that the claimed capacity is real, and all blocks can hold data.
# It is not a full proof that the flash memory is good.
#
# Copyright (C) by Volker Kuhlmann
#               http://volker.top.geek.nz/contact.html
# All rights reserved.
# Released under the terms of the GNU General Public License (GPL) Version 2.
# See http://www.gnu.org/ for details.
#
# Volker Kuhlmann
#   1.0, 19 Apr 2015  Created.
#   1.0.1, 19 Apr 2015  Added comment.
#   1.1, 01 May 2015
#	Write to status files named with time stamp.
#	Document --debug. Improve help output. Minor bug fixes.
#	Write all dd stderr in write_chunk_tempfile() to debug log.
#
chr0(){ echo @;};chr1(){ echo .;}

VERSION="1.1, 01 May 2015"
AUTHOR="Volker Kuhlmann <VolkerKuhlmann`chr0`gmx`chr1`de>"
COPYRIGHT="Copyright (C) 2015"


####
#### Constants and initialised variables
####

# Test chunk size, in bytes.
# A hash/checksum is calculated for each chunk.
# Keep this powers of 2! Suggested 1024k or more.
CHUNKSIZE=32M

# Temp directory root to use.
# Put this into RAM disk!!
TEMPDIR_ROOT="/dev/shm"

# File where checksums of data chunks are stored in.
STATFILEROOT="flash_status"
STATFILE="${STATFILEROOT}_`date +%Y%m%dT%H%M%S`"

# File where checksums of chunks written to the test device are stored in.
# This is appended to STATFILE.
STATWRITE="_write"

# File where checksums of chunks read from the test device are stored in.
# This is appended to STATFILE.
STATVERIFY="_verify"

# Debug log file.
# This is appended to STATFILE.
STATDEBUG="_debug"

# Which chunk write function to use.
write_chunk=write_chunk_tempfile
write_chunk=write_chunk_tee


####
#### Version, Usage, Help
####

show_version() {
    echo "${0##*/} version $VERSION
$COPYRIGHT by $AUTHOR"
}

show_usage() {
    echo "
Usage: ${0##*/} DEVICE [CAPACITY]
Version $VERSION
$COPYRIGHT by $AUTHOR
"
    test -n "$1" && echo "$1
"
}

show_help() {
    show_usage
    echo "\
Check that the claimed capacity of a USB stick / flash card DEVICE is real, and
that all blocks can hold data.
Optionally give the size to test as CAPACITY.

OPTIONS:
  -h|--help     Show help.
  --version     Show version.
  -i|--info     Show details about operation of this script.
  --chunk SIZE  Use chunk size SIZE for testing.
                  Rounded down to full KiB (1024 bytes). Default $CHUNKSIZE.
  -y|--yes      Answer yes to all questions.
  --debug       Write chunk I/O speeds to $STATFILEROOT*$STATDEBUG.

CAPACITY and SIZE accept k, M, G as multiplier.
Files created in current directory: $STATFILEROOT*
"
#Creates files in current directory: $STATFILE, $STATFILE$STATVERIFY
}

Version() { show_version; exitwith ErrVersion; }
Usage() { test -n "$1" && exitwith ErrUsage show_usage "$1";
    	    show_usage; exitwith ErrOK; }
Help() { test -n "$1" && exitwith ErrHelp show_help; show_help; exitwith ErrOK; }

Info() {
    echo "
Test method:
Read chunks from the random number generator and write them to the test disk,
while also writing hashes to $STATFILE$STATWRITE.
Read chunks from test disk, writing hashes to $STATFILE$STATVERIFY, which have
to match those calculated while writing.

It is a good indication that the device has its claimed capacity and can hold
data on all chunks if all chunks verify OK.

More good tests:
 * Repeat test, writing chunks in reverse (end to start).
 * Write some chunks again at the beginning, to see it it affects the end.
 * Within a chunk: Check 128-byte sub-chunks, and all bytes in a sub-chunk, for
   byte-swapping or byte-mirroring.

Currently dd is run with bs=chunksize. It may be more efficient to run it with
bs=1M and appropriate count= for chunksize - but a quick test shows that is not
the case.

Writing data to a USB flash drive read directly from /dev/urandom is noticeably
slower than writing urandom data read from RAM disk (about 72% with my test
device). Reading urandom data in the background into 2 alternating RAM disk
buffers might speed that up again.

On fast flash memory the write speed is limited by the speed at which the random
number generator can be read.

This script uses up all your system's entropy so don't create encryption keys
immediately after.
"
    exitwith ErrOK
}


####
#### Error/Exit codes
####

exitwith() {
    exec 1>&2  # Write stdout on stderr instead.
  case "$1" in
    ErrOK|ErrVersion)
	exit 0;;
    ErrUsage|ErrHelp)
	# Output generated by function/program $2, if given.
	test -n "$2" && { shift; "$@"; }
	exit 1;;
 # more codes in here
    ErrTestPass)
    	echo "PASSED all tests."
	exit 0;;
    ErrNoDevCapacity)
	echo "Couldn't get device capacity for '$2'."
	exit 3;;
    ErrNothingToTest)
        echo "No chunks to test on '$2'."
	exit 3;;
    ErrIllValue)
    	echo "Bad value: '$2'."
	exit 4;;
    ErrDir)
    	echo "Is a directory: '$2'."
	exit 4;;
    ErrTestFail)
    	echo "FAILED at least one test."
    	exit 5;;
 # more codes in here
    ErrBadoption)
    	echo "Bad option '$2'."
	echo "Call with -h for help."
	exit 9;;
    ErrMissingParameter)
	echo "A required parameter for option $2 is missing."
	echo "Call with -h for help."
	exit 9;;
    *)
	echo "Internal error: exitwith() called with illegal error code '$1'."
	exit 19;;
  esac
}


####
#### Parse command line parameters
####

parse_cmd_line() {
    local cmdline arg
    unset -v debug verbose chunksize yes
    cmdline=("$@")
    test "$1" = "--debug" && debug=1  # Pre-check for debug.
    #set -- -opt1 "$@"  # pre-init options
    while [ -n "$1" ]; do
#	test $debug && echo "Current arg: $1"
	case "$1" in
	    --) shift; break;;	    # Stop if word is "--".
#	    -[^-]?*) arg="$1";;     # "-long" or "--long" for long options.
#    	    -[^-][A-Za-ce-z]*)	    # "-long" or "--long", except for some
#		arg="$1";;  	    # shortcuts.
	    -*) arg="${1#-}";;	    # Strip one hyphen.
	    *) break;;	    	    # Stop if word does not start with "-".
	esac
#	test $debug && echo "    cmd arg: $arg"
	case "$arg" in
	    -debug)	    debug=1;;
	    -version)	    Version;;
	    -usage) 	    Usage;;
	    h|help|-help)   Help;;
#	    t|-test)	    dryrun=1 _drycreate=1;;
	    
	    i|-info)  	    Info;;
	    -chunk) 	    chunksize="$2"; shift;;
	    y|-yes)    	    yes=1;;

	    "")     break;;  # allow "-" as file arg
	    *)	    exitwith ErrBadoption "$1";;
	    #*)	    break;;  # stop option scanning with first unknown option
	esac
	shift
    done
    # save option and file arguments for later:
    FILEARGS=("$@")
    OPTARGS=("${cmdline[@]:0:${#cmdline[@]} - ${#FILEARGS[@]}}")
}


####
#### Functions
####

# Get the current time with sub-second resolution.
get_stamp() {
    date "+%s.%N"
}

# Obtain the device capacity, in bytes, on stdout.
# Return code != 0 if size is not a number.
# Uses fdisk -l to obtain the device capacity (works on GPT tables).
get_device_size() {
    local dev="$1"
    
    # fdisk output e.g.
    #  Disk /dev/sdk: 8053 MB, 8053063680 bytes

    fdisk -l "$1" 2>/dev/null \
    	| grep -i "^Disk.*$dev" \
	| {
	    read -r line
	    val="${line% bytes}"
	    val="${val##* }"
	    echo "$val"
	    
	    # Test for number, set function return value.
	    test "$val" -gt 0 2>/dev/null
	}
}

# Interpret and check the given capacity expression.
# Expand kilo, Mega, Giga units. (Ignore bytes.)
# Example: "3 * 16M" = 50331648 (bytes).
# $1: expression
# stdout: expression's expanded value
expand_size() {
    local expr="$1" old= new
    
    # Expand the units.
    while [ "$expr" != "$old" ]; do
	old="$expr"
	expr="${expr/G/ *1024*1024*1024}"
	expr="${expr/M/ *1024*1024}"
	expr="${expr/k/ *1024}"
	#expr="${expr/b/ }"  # not a proper unit
    done
    
    # Evaluate expression, show error if necessary.
    evalexpr() {
	echo "$(($expr))" # Can't suppress error to stderr without function.
    }
    new="`evalexpr 2>/dev/null`"
    test $? -eq 0 -a "$new" -gt 0 || exitwith ErrIllValue "$expr"
    echo "$new"
}

# Create temp dir.
# Out: TEMPDIR, the dir that was created.
create_temp_dir() {
    local tempdir
    tempdir="`date +%Y%m%d_%H%M%S`"
    tempdir="$TEMPDIR_ROOT/testflash_$tempdir"
    tempdir="`mktemp -d "${tempdir}_XXXXXXXXXXXX"`"
    if [ $? -ne 0 ]; then
    	echo >&2 "Failed to create temp dir!"
	exit 9
    fi
    TEMPDIR="$tempdir"
    
    trap delete_temp_dir EXIT

    echo "Temp dir: $TEMPDIR"
    #ls -la "$TEMPDIR"
}
delete_temp_dir() {
    rm "$TEMPDIR/"* 2>/dev/null || true
    rmdir "$TEMPDIR"
}

# Write a chunk to the device under test.
# In: $1 chunk number (0 based), $2 device, $3 status file.
write_chunk_tempfile() {
    local chunk="$1" dev="$2" stat="$3"
    
    test -n "$TEMPDIR" || create_temp_dir
    
    local chunkdata="$TEMPDIR/chunkdata"
    
    # Read a chunk of random data, write it to the device, and write a checksum
    # to the status file.
    dd bs="$chunksize" if=/dev/urandom of="$chunkdata" count=1 2>>"$debugfile"
    dd bs="$chunksize" if="$chunkdata" of="$dev" seek="$chunk" count=1 \
    	oflag=direct conv=notrunc 2>>"$debugfile"
    
    hash="`md5sum < "$chunkdata"`"
    hash="${hash%% *}"
    echo -e "$chunk\t$hash" >> "$stat"
}
# Alternative implementation, without temp file. Not much faster.
# In: $1 chunk number (0 based), $2 device, $3 status file.
write_chunk_tee() {
    local chunk="$1" dev="$2" stat="$3"
    
    local chunkdata="$TEMPDIR/chunkdata"
    
    exec 4> >(
	hash="`md5sum`"
	hash="${hash%% *}"
	echo -e "$chunk\t$hash" >> "$stat"
    )
    
    # Read a chunk of random data, write it to the device, and write a checksum
    # to the status file.
    # The iflag=fullblock is important or the pipe closes early.
    dd bs="$chunksize" if=/dev/urandom count=1 2>>"$debugfile" \
    	| tee -a /dev/fd/4 \
    	| dd bs="$chunksize" of="$dev" seek="$chunk" count=1 \
    	    iflag=fullblock oflag=direct conv=notrunc 2>>"$debugfile"
    exec 4>&-
}

# Verify a chunk from the device under test.
# In: $1 chunk number (0 based), $2 device, $3 status file.
verify_chunk() {
    local chunk="$1" dev="$2" stat="$3"

    local chunkdata="$TEMPDIR/chunkdata"
    
    hash="`dd bs="$chunksize" if="$dev" skip="$chunk" count=1 iflag=direct \
    	2>>"$debugfile" | md5sum`"
    hash="${hash%% *}"
    echo -e "$chunk\t$hash" >> "$stat"
}

# Calculate transfer speeds.
# In: start="$1" end="$2" volume="$3"
# Out: mbps, mbpsSI, runtime; the output must be eval'ed!
fp_speed() {
    echo "$1" "$2" "$3" | awk '{
    	printf "mbps=%.2f\nmbpsSI=%.2f\nruntime=%.1f\n",
	    $3/($2-$1)/1024/1024, $3/($2-$1)/1000/1000, $2-$1
	}'
}

# Calculate floating point numbers.
# In: $1 value
# Out: valkB, valKiB, valMB, valMiB, valGB, valGiB; the output must be eval'ed!
fp_num() {
    echo "$1" | awk '{
    	printf "valkB=%.3f\nvalKiB=%.3f\nvalMB=%.3f\nvalMiB=%.3f\n",
	    $1/1000, $1/1024, $1/1000/1000, $1/1024/1024
    	printf "valGB=%.3f\nvalGiB=%.3f\n",
	    $1/1000/1000/1000, $1/1024/1024/1024
	}'
}

# Ask user before overwriting a device.
doublecheck_device() {
    test "$yes" = 1 && return
    
    if [ -b "$device" ]; then
    	echo -n "$device is a block special device - proceed? (y/N) "
	read line
	test "$line" = "y" || exitwith ErrOK
    elif [ -d "$device" ]; then
    	exitwith ErrDir "$device"
    elif [ -e "$device" -a ! -f "$device" ]; then
    	echo -n "$device is not a plain file - proceed? (y/N) "
	read line
	test "$line" = "y" || exitwith ErrOK
    fi
}

# Perform a write/verify pass over the device.
run_test() {
    rm "$STATFILE$STATWRITE" "$STATFILE$STATVERIFY" 2>/dev/null

    echo "Writing all chunks:"
    local s1=`get_stamp`
    local nchunk=0
    while [ $nchunk -lt $chunks ]; do
	$write_chunk $nchunk "$device" "$STATFILE$STATWRITE"
	let nchunk++
    done
    local s2=`get_stamp`
    local mbps mbpsSI runtime
    eval `fp_speed $s1 $s2 $(($chunks * $chunksize))`
    echo "Written all chunks in $runtime seconds: $mbps MiB/s, $mbpsSI MB/s."

    echo "Reading all chunks:"
    s1=`get_stamp`
    nchunk=0
    while [ $nchunk -lt $chunks ]; do
	verify_chunk $nchunk "$device" "$STATFILE$STATVERIFY"
	let nchunk++
    done
    s2=`get_stamp`
    eval `fp_speed $s1 $s2 $(($chunks * $chunksize))`
    echo "Read all chunks in $runtime seconds: $mbps MiB/s, $mbpsSI MB/s."

    echo "Verifying all chunks:"
    diff -q "$STATFILE$STATWRITE" "$STATFILE$STATVERIFY" >/dev/null
    if [ $? -ne 0 ]; then
	echo "ERROR: corrupted chunks encountered on device $device!
To show, use .e.g.   diff -U0 $STATFILE$STATWRITE $STATFILE$STATVERIFY"
	false
    else
	echo "OK: All chunks pass on device: '$device'"
	echo "Find checksums in $STATFILE$STATVERIFY"
	rm "$STATFILE$STATWRITE"
	
	echo -n "Check all chunks have a different checksum: "
	lines="`cut -f2 "$STATFILE$STATVERIFY" | sort | uniq | wc -l`"
	if [ $lines -eq $chunks ]; then
	    echo "OK"
	else
	    echo "FAIL"
	    echo "ERROR: Got $lines different checksums for $chunks chunks!"
	    echo "       This may indicate fake memory areas."
	    false
	fi
    fi
}


####
#### Main
####

test $# -ge 1 || Help err

parse_cmd_line "$@"

# Device name from cmd line (compulsory).
device="${FILEARGS[0]}"

# Device capacity from cmd line (optional).
test $# -ge 2 -a -n "${FILEARGS[1]}" && devsize="${FILEARGS[1]}"

# Work out device capacity.
if [ -n "$devsize" ]; then
    devsize="`expand_size "$devsize"`" || exit $?
else
    devsize="`get_device_size "$device"`"
    test $? -eq 0 || exitwith ErrNoDevCapacity "$device"
fi
eval `fp_num "$devsize"`
echo "Device capacity $devsize bytes ($valMiB MiB, $valGiB GiB)."
test "$devsize" -gt 0 || exitwith ErrIllValue "$devsize"

# Work out chunk size.
if [ -n "$chunksize" ]; then
    chunksize="`expand_size "$chunksize"`" || exit $?
    chunksize=$(($chunksize / 1024))
    chunksize=$(($chunksize * 1024))
else
    chunksize="`expand_size "$CHUNKSIZE"`" || exit $?
fi
eval `fp_num "$chunksize"`
echo "Chunk size $valMiB MiB, $valKiB KiB, $chunksize bytes."
test "$chunksize" -gt 0 || exitwith ErrIllValue "$chunksize"

# Work out number of chunks.
chunks=$(($devsize / $chunksize))
untested=$(($devsize - ( $chunks * $chunksize ) ))
echo "$chunks chunks, $untested bytes untested (at end)."
test $chunks -gt 0 || exitwith ErrNothingToTest "$device"

# Debug output. Contains dd transfer speeds.
debugfile=/dev/null
test -n "$debug" && debugfile="$STATFILE$STATDEBUG"

# Run test.
unset TEMP_DIR
doublecheck_device
run_test || exitwith ErrTestFail
exitwith ErrTestPass