prng

#!/bin/bash
DEFAULT_SIZE=100M
usage(){
    echo "Usage: prng stream SEED"
    echo "       prng stdout SEED [SIZE]"
    echo "       prng file SEED [SIZE]"
    echo "       prng files COUNT [SIZE]"
    echo "       prng checksums file SEED [SIZE]"
    echo "       prng checksums files COUNT [SIZE]"
    echo "       prng verify stdin SEED [SIZE]"
    echo "       prng verify file SEED [SIZE]"
    echo "       prng verify files COUNT [SIZE]"
    echo "       prng test-usb BLOCKDEVICE [BLOCKSIZE]"
    exit 1
}
RANDOM_STREAM=$(dirname $0)/random-stream

stream(){
    seed=$1
    seed=$((4357+seed*2))
	bytes=$2
	if [ -z "$bytes" ]; then
		echo -n "inf:"
    else
		echo -n "${bytes}:"
    fi
    echo -n "MT19937:$((seed))::"
    $RANDOM_STREAM $seed
}
stdout(){
    seed=$1
    bytes=$2
    [ $2 -ge 100 ] || {
        echo "Size must be at least 100 bytes" >/dev/stderr
        exit 1
    }
    stream $seed $bytes | head -c $bytes
}
file(){
    stdout $1 $2 >"$3"
}
checksum_file(){
    stdout $1 $2 | sha256sum | cut -d' ' -f1 | tr -d "\n"
    echo " " $3
}

size(){
	size=$1
	if [ -z "$size" ]; then
		size=$DEFAULT_SIZE
	fi
	dehumanize $size
}
dehumanize(){
    echo $1 | awk \
      'BEGIN{IGNORECASE = 1}
       function printpower(n,b,p) {printf "%u\n", n*b^p}
       /[0-9]$/{print $1};
       /K(iB)?$/{printpower($1,  2, 10)};
       /M(iB)?$/{printpower($1,  2, 20)};
       /G(iB)?$/{printpower($1,  2, 30)};
       /T(iB)?$/{printpower($1,  2, 40)};
       /KB$/{    printpower($1, 10,  3)};
       /MB$/{    printpower($1, 10,  6)};
       /GB$/{    printpower($1, 10,  9)};
       /TB$/{    printpower($1, 10, 12)}'
}
write_block(){
    seed=$1
    blocksize=$2
    offset=$(($3 - 1))
    blocks=$4
    device=$5
    echo "writing block $((offset+1))/$blocks" >/dev/stderr
    stdout $((seed + offset)) $blocksize | dd seek=$offset of=$device bs=$blocksize status=none
}
verify_block(){
    seed=$1
    blocksize=$2
    offset=$(($3 - 1))
    blocks=$4
    device=$5
    dd skip=$offset if=$device bs=$blocksize count=1 status=none | $0 verify stdin $((seed + offset)) $blocksize
    if [ $? -eq 0 ]; then
        echo "verified block $((offset+1))/$blocks" >/dev/stderr
    else
        echo "block $((offset+1))/$blocks incorrect" >/dev/stderr
        exit 1
    fi
}

[ $# -ge 1 ] || usage
command=$1
shift

case $command in
    stream)
        [ $# -eq 1 ] || usage
        stream "$@"
    ;;
    stdout)
        [ $# -ge 1 ] || usage
        [ $# -le 2 ] || usage
		seed=$1
        s=$(size $2)
		stdout $seed $s
    ;;
    file)
        [ $# -ge 1 ] || usage
        [ $# -le 2 ] || usage
        seed=$1
        s=$(size $2)
        file $seed $s "prng-$2.${seed}"
    ;;
    files)
        [ $# -ge 2 ] || usage
        [ $# -le 3 ] || usage
        count=$1
        s=$(size $2)
        seq 1 $count | while read seed; do
            file $seed $s "prng-$2.${seed}"
        done
        $0 checksums files "$@" >checksums
    ;;
	checksums)
        command=$1
        shift
        case $command in
            file)
                [ $# -ge 1 ] || usage
                [ $# -le 2 ] || usage
                seed=$1
                s=$(size $2)
                checksum_file $seed $s "prng-$2.${seed}"
            ;;
            files)
                [ $# -ge 2 ] || usage
                [ $# -le 3 ] || usage
                count=$1
                s=$(size $2)
                seq 1 $count | while read seed; do
                    checksum_file $seed $s "prng-$2.${seed}"
                done
            ;;
            *)
                usage
            ;;
        esac
    ;;
    verify)
        case $1 in
            stdin)
                shift 1 
                [ $# -ge 1 ] || usage
                [ $# -le 2 ] || usage
                seed=$1
                s=$(size $2)
                ACTUAL=$(sha256sum)
                EXPECTED=$(stdout $seed $s | sha256sum)
                if [ "$EXPECTED" = "$ACTUAL" ]; then
                    exit 0
                else
                    echo "Got: $ACTUAL" >/dev/stderr
                    echo "Expected: $EXPECTED" >/dev/stderr
                    exit 2
                fi
                ;;
            file|files)
                $0 checksums "$@" | sha256 -c -
                ;;
            *)
                usage
                ;;
        esac
    ;;
    test-usb)
        [ $# -ge 1 ] || usage
        [ $# -le 2 ] || usage
        [ `whoami` = 'root' ] || {
            echo "Accessing block devices directly requires running as root" >/dev/stderr
            exit 1
        }
        seed=1
        device=$1
        blocksize=$(size $2)
        size=$(blockdev --getsize64 $device || stat --printf="%s" $device)
        if [[ -z "$2" && $(( $size % $blocksize )) -ne 0 ]]; then
            echo "apparent size is: $size. writing data." >/dev/stderr
            stdout $seed $size | pv --size ${size} | dd of=$device
            echo "wrote data. verifying" >/dev/stderr
            dd if=$device | pv --size ${size} | $0 verify stdin $seed $size
            if [ $? -eq 0 ]; then
                echo "all correct" >/dev/stderr
                exit 0
            else
                echo "not all correct" >/dev/stderr
                exit 1
            fi
        exit $?
        else
            blocks=$((size / blocksize))
            echo "apparent size is: $size. writing $blocks blocks of data of $blocksize each." >/dev/stderr
            if [ $((blocks * blocksize)) -eq $((size)) ]; then
                echo "  (blocks divide the disk evenly)" >/dev/stderr
            else
                echo "  (last $((size - ( blocks * blocksize ) )) bytes will be skipped)" >/dev/stderr
            fi

            # Test first and last block at the very start
            echo "writing first and last block to test failure quickly." >/dev/stderr
            write_block $seed $blocksize 1 $blocks $device
            write_block $seed $blocksize $blocks $blocks $device
            verify_block $seed $blocksize 1 $blocks $device || exit 1
            verify_block $seed $blocksize $blocks $blocks $device || exit 1
            echo "first and last block verified. will now write+test every block" >/dev/stderr
            seq 2 $((blocks - 1)) | while read offset; do
                write_block $seed $blocksize $offset $blocks $device
            done
            seq 1 $blocks | while read offset; do
                verify_block $seed $blocksize $offset $blocks $device || exit 1
            done
            echo "all correct" >/dev/stderr
            exit 0
        fi
    ;;
    *)
        usage
    ;;
esac