#!/bin/ksh

##
## USAGE definition
##
# usage: bht [options] [list of hdds]
#   -b, --blocks=   the size of blocks in bytes. default is 32768
#   -n, --stride=   the number of blocks per test. default is 512
#   -e, --email=    email address to send notifications
#   -s, --status    show status of running tests
#   -d, --directory path to directory with bht data. defaults to current working directory
#   -t, --terminate terminte all running tests
USAGE=$'[-?\n@(#)$Id: 0.1\n]'
USAGE+="[-author?Art of Server <the.art.of.server@gmail.com>]"
USAGE+="[-copyright?Copyright (C) 2018 Art of Server.]"
USAGE+="[-license?GNU General Public License v3.]"
USAGE+="[+NAME?bht --- bulk hard-drive tester.]"
USAGE+="[+DESCRIPTION?bht is a bulk hard drive testing "
USAGE+="script that runs badblocks on multiple hard drives "
USAGE+="in parallel and allows to user to walk away and "
USAGE+="login later to check on the results.]"
USAGE+="[b:blocks]:[blocks:=32768?the size of blocks in bytes.]"
USAGE+="[n:stride]:[stride:=512?the number of blocks per test.]"
USAGE+="[e:email]:[email?email address to send notifications.]"
USAGE+="[s:status?show status of current tests.]"
USAGE+="[d:directory]:[directory?path to directory with bht data. defaults to current working directory.]"
USAGE+="[t:terminate?terminate all running tests.]"

##
## configuration settings
##

# options flags
OPT_BLOCKS=0
OPT_STRIDE=0
OPT_EMAIL=0
OPT_STATUS=0
OPT_BHTDIR=0
OPT_TERMINATE=0
# option variables
BLOCKS_SIZE=32768
STRIDE_SIZE=512
EMAIL=
BHTDIR="$(pwd)"
# misc
NL=$'\n'

##
## functions
##

#-----------------------------------------------------------------------------
# description: show status of all hard drive tests
# inputs: $1=path to data directory
# outputs: print status to stdout
# return: 0==success
#-----------------------------------------------------------------------------
function show_status {
    typeset _dir=$1
    typeset _hdd
    typeset _status
    typeset _smart
    typeset _hdd_type

    for D in $(find ${_dir}/disks -mindepth 1 -type d)
    do
	# extract hdd_type
	_hdd_type=$(<${D}/hdd_type)
	# extract hdd from path $D
	_hdd=${D##*/}
	if [[ -f ${D}/badblocks.out ]] ; then
	    # grab the first part of the status string
	    _status=$(awk '{gsub(/\b+/,"\n",$0); print}' ${D}/badblocks.out \
		| awk -F: '$1~"^(Testing with|Reading and)" {line=$1} END{print line}') 
	    # insert space
	    _status+=" "
	    # now grab the latest of the last part
	    _status+=$(awk '{gsub(/\b+/,"\n",$0); print}' ${D}/badblocks.out \
		| awk '!/^$/ {line=$0} END{print line}')
	    print -- "$_hdd:"
	    print -- "\tbadblocks[$_status]"
	    # check SMART data. use post-test.smart if available. fall back on pre-test.smart
	    if [[ -f ${D}/post-test.smart ]] ; then
	        _smart="$(SMART_check ${D}/post-test.smart)"
	    else
	        _smart="$(SMART_check ${D}/pre-test.smart)"
	    fi
	    print -- "${_smart}" | awk '$0="\t"$0'
	fi
    done
    return 0
}
    
#-----------------------------------------------------------------------------
# description: terminate all running badblocks tests
# inputs: $1=path to data directory
# outputs: print termination message to stdout
# return: 0==success
#-----------------------------------------------------------------------------
function terminate_tests {
    typeset _dir=$1
    typeset _hdd
    typeset _pid

    for D in $(find ${_dir}/disks -mindepth 1 -type d)
    do
	# extract hdd from path $D
	_hdd=${D##*/}
	if [[ -f ${D}/badblocks.pid ]] ; then
	    read _pid <${D}/badblocks.pid
	    # check pid is actually for badblocks
	    if [[ "$(ps -eo pid,comm \
		| awk -v pid=$_pid '$1==pid && $2=="badblocks" {print $1}')" \
		-eq $_pid ]] ; then
		# print message
		print -- "$_hdd: terminating PID=$_pid."
		kill $_pid
	    fi
	fi
    done
    return 0
}

#-----------------------------------------------------------------------------
# description: print HDD type (SATA or SAS)
# inputs: $1=path to file containing output from "smartctl -a"
# outputs: print "SAS" or "SATA"
# return: none
#-----------------------------------------------------------------------------
function print_HDD_type {
    typeset _file=$1
    typeset is_SATA=0
    typeset is_SAS=0

    if [[ $(awk -F'[: ]+' '$1~"^SATA" && $2=="Version" {print $4}' $_file) == "SATA" ]] ; then
	is_SATA=1
    fi

    if [[ $(awk -F'[: ]+' '$1~"^Transport" && $2=="protocol" {print $3}' $_file) == "SAS" ]] ; then
	is_SAS=1
    fi

    if (( is_SATA == 1 )) && (( is_SAS == 1 )) ; then
	print "ERR"
    elif (( is_SATA == 1 )) ; then
	print "SATA"
    elif (( is_SAS == 1 )) ; then
	print "SAS"
    else
	print "NA"
    fi
    return 0
}

#-----------------------------------------------------------------------------
# description: check SMART attributes of HDD
# inputs: $1=path to file containing output from "smartctl -a"
# outputs: print some SMART attributes to stdout
# return: none
#-----------------------------------------------------------------------------
function SMART_check {
    typeset _file=$1
    typeset -i _5_Reallocated_Sector_Ct
    typeset -i _9_Power_On_Hours
    typeset -i _200_Multi_Zone_Error_Rate
    typeset _power_on_time
    typeset _hdd_type=$(print_HDD_type $_file)

    print "HDD Type:[$_hdd_type]"
    if [[ $_hdd_type == "SATA" ]] ; then
	_5_Reallocated_Sector_Ct=$(awk '$1=="5" && $2=="Reallocated_Sector_Ct" {print $10}' $_file)
	_9_Power_On_Hours=$(awk '$1=="9" && $2=="Power_On_Hours" {print $10}' $_file)
	_200_Multi_Zone_Error_Rate=$(awk '$1=="200" && $2=="Multi_Zone_Error_Rate" {print $10}' $_file)
	print "SMART:[Reallocated_Sector_Ct=$_5_Reallocated_Sector_Ct]"
	print "SMART:[Power_On_Hours=$_9_Power_On_Hours]"
	print "SMART:[Multi_Zone_Error_Rate=$_200_Multi_Zone_Error_Rate]"
    elif [[ $_hdd_type == "SAS" ]] ; then
	_power_on_time=$(awk '$0~/power on time/ {print $6}' $_file)
	if [[ -n $_power_on_time ]] ; then
	    print "SMART:[power_on_time(hours:minutes)=$_power_on_time]"
	fi
    fi
}

#-----------------------------------------------------------------------------
# description: run badblocks tests on hard drives in background
# inputs: $1=path to block device, $2=folder to write output
# outputs: none
# return: 0==success
#-----------------------------------------------------------------------------
function run_badblocks {
    typeset _hdd=$1
    typeset _dir=$2
    typeset _status
    typeset _mesg
    typeset _hdd_type=$(<${_dir}/hdd_type)

    print "$_hdd" > $_dir/log
    badblocks -b $BLOCKS_SIZE -c $STRIDE_SIZE -wsv $_hdd >${_dir}/badblocks.out 2>&1 &
    print $! > ${_dir}/badblocks.pid
    # wait for badblocks to finish
    wait
    # collect SMART data after the test
    if [[ $_hdd_type == "SATA" ]] ; then
	smartctl -a $D > ${_dir}/post-test.smart
    elif [[ $_hdd_type == "SAS" ]] ; then
	smartctl --xall $D > ${_dir}/post-test.smart
    fi
    # when we're done, send out notification if we have EMAIL
    if [[ -n $EMAIL ]] ; then
	_mesg="Model=${HDD_DATA[$_hdd].model}, Serial=${HDD_DATA[$_hdd].serial}$NL"
	# grab the first part of the status string
	_status=$(awk '{gsub(/\b+/,"\n",$0); print}' ${_dir}/badblocks.out \
	    | awk -F: '$1~"^(Testing with|Reading and)" {line=$1} END{print line}')
	# add a space
	_status+=" "
	# now grab the latest of the last part
	_status+=$(awk '{gsub(/\b+/,"\n",$0); print}' ${_dir}/badblocks.out \
	    | awk '!/^$/ {line=$0} END{print line}')
	# append to _mesg
	_mesg+="badblocks[$_status]$NL"
	# append SMART check
	_mesg+="$(SMART_check ${_dir}/post-test.smart)"
	print "$_mesg"|mailx -s "badblocks testing of $_hdd ended" $EMAIL
    fi
    return 0
}

##
## main section
##

SELF=$0
PID=$$

# process command line args
while getopts "$USAGE" option
do
    case $option in
	b|blocks)   OPT_BLOCKS=1 && BLOCKS_SIZE=$OPTARG	;;
	n|stride)   OPT_STRIDE=1 && STRIDE_SIZE=$OPTARG	;;
	e|email)    OPT_EMAIL=1 && EMAIL=$OPTARG	;;
	s|status)   OPT_STATUS=1			;;
	d|directory) OPT_BHTDIR=1 && BHTDIR=$OPTARG	;;
	t|terminate) OPT_TERMINATE=1			;;
    esac
done

#print "OPT_BLOCKS=$OPT_BLOCKS, OPT_STRIDE=$OPT_STRIDE, OPT_EMAIL=$OPT_EMAIL, OPT_STATUS=$OPT_STATUS, OPT_BHTDIR=$OPT_BHTDIR, OPT_TERMINATE=$OPT_TERMINATE"
#print "BLOCKS_SIZE=$BLOCKS_SIZE, STRIDE_SIZE=$STRIDE_SIZE, EMAIL=$EMAIL, BHTDIR=$BHTDIR"

# remaining non-options should be list of hard drives to test
shift "$((OPTIND-1))"
typeset -a HDD_LIST=( $@ )
#print "HDD_LIST size = ${#HDD_LIST[*]}, list = ${HDD_LIST[*]}"

# check that all required external utilities are available
for p in smartctl badblocks sha256sum mkdir id lsscsi mktemp grep awk mv find mailx ps pvs
do
    if ! $(which $p 2>/dev/null 1>&2) ; then
	print "ERR: $p is not available or not in your PATH. please install $p and try again."
	exit 1
    fi
done

# check that we have root privileges
if (( $(id -u) != 0 )) ; then
    print "ERR: this program must be run with root privileges."
    exit 1
fi

# check conflicting options
# --status should not be used with --blocks, --stride, --email, --terminate
if (( $OPT_STATUS == 1 && ( $OPT_BLOCKS == 1 || $OPT_STRIDE == 1 || $OPT_EMAIL == 1 ) )) ; then
    print "ERR: cannot specify --status with --terminate, --blocks, --stride, nor --email."
    exit 1
fi 
# --terminate should not be used with --blocks, --stride, --email
if (( $OPT_TERMINATE == 1 && ( $OPT_BLOCKS == 1 || $OPT_STRIDE == 1 || $OPT_EMAIL == 1 ) )) ; then
    print "ERR: cannot specify --terminate with --blocks, --stride, nor --email."
    exit 1
fi 

# if OPT_STATUS=1 or OPT_TERMINATE=1, then BHTDIR must already exist
if (( $OPT_STATUS == 1 )) || (( $OPT_TERMINATE == 1 )) ; then
    if [[ ! -d $BHTDIR ]] ; then
	print "ERR: data directory $BHTDIR doesn't exist."
	print "ERR: cannot get status of non-existent data."
	exit 1
    fi
    # must have r and x permissions
    if [[ ! -x $BHTDIR ]] ; then
	print "ERR: data directory $BHTDIR is not accessible by you."
	exit 1
    fi
    if [[ ! -r $BHTDIR ]] ; then
	print "ERR: data directory $BHTDIR is not readable by you."
	exit 1
    fi
    # check for magic file
    if [[ -f ${BHTDIR}/.bht ]] ; then
	# compare checksum
	if [[ $(sha256sum $0| awk '{print $1}') != $(awk '{print $1}' ${BHTDIR}/.bht) ]] ; then
	    print "WARN: $BHTDIR has data generated by a different version of bht."
	fi
    else
	print "ERR: $BHTDIR directory doesn't appear to have bht magic?"
	exit 1
    fi
    if (( $OPT_STATUS == 1 )) ; then
	# run show_status
	show_status $BHTDIR
	ret=$?
	if [[ $ret -ne 0 ]] ; then
	   print "ERR: show_status() returned an error code [$ret]."
	    exit 1
	fi
    fi
    if (( $OPT_TERMINATE == 1 )) ; then
	# run terminate_tests
	terminate_tests $BHTDIR
	ret=$?
	if [[ $ret -ne 0 ]] ; then
	   print "ERR: terminate_tests() returned an error code [$ret]."
	    exit 1
	fi
    fi
    # exit success
    exit 0
fi

# prepare to run tests
# 1) if $BHTDIR doesn't exist, create it and setup magic file
# 2) move to $BHTDIR
# 3) check that all elements in HDD_LIST are valid HDDs
if [[ ! -e $BHTDIR ]] || [[ ! -d $BHTDIR ]] ; then
    mkdir -p $BHTDIR
fi
sha256sum $0 > ${BHTDIR}/.bht

print "INFO: changing to $BHTDIR."
cd $BHTDIR

# make sure $BHTDIR is rwx
[[ -x $BHTDIR ]] || ( print "ERR: data directory $BHTDIR is not accessible by you." && exit 1 )
[[ -w $BHTDIR ]] || ( print "ERR: data directory $BHTDIR is not write-able by you." && exit 1 )
[[ -r $BHTDIR ]] || ( print "ERR: data directory $BHTDIR is not read-able by you." && exit 1 )

for i in ${!HDD_LIST[*]}
do
    # make sure it is a block device
    if [[ ! -b ${HDD_LIST[i]} ]] ; then
	print "WARN: ${HDD_LIST[i]} is not a block device. removing from list."
	unset HDD_LIST[i]
	continue
    fi
    # make sure it is a "disk" device in lsscsi
    if (( $(lsscsi | awk -v d="${HDD_LIST[i]}" \
	'BEGIN{i=0} $2=="disk" && $NF==d {i++}END{print i}') == 0 )) ; then
	print "WARN: ${HDD_LIST[i]} is not a disk in lsscsi. removing from list."
	unset HDD_LIST[i]
	continue
    fi
    # make sure it is not mounted
    if (( $(awk -v dev="^${HDD_LIST[i]}[0-9]+" \
	'BEGIN{i=0} $1~dev {i++} END{print i}' /etc/mtab) > 0 )) ; then
	print "WARN: ${HDD_LIST[i]} appears to be mounted. removing from list."
	unset HDD_LIST[i]
	continue
    fi
    # make sure it is not part of a volume group
    if (( $(pvs --no-headings | awk -v dev="^${HDD_LIST[i]}" \
	'BEGIN{i=0} $1~dev {i++} END{print i}') > 0 )) ; then
	print "WARN: ${HDD_LIST[i]} appears to be part of LVM. removing from list."
	unset HDD_LIST[i]
	continue
    fi
    ## NEED TO ADD TEST TO CHECK RAID/mdraid
    # if we have zpool, make sure this disk is not part of a ZFS pool
    if $(which zpool 1>/dev/null 2>&1) && \
	(( $(zpool status -L -P | awk -v dev="^${HDD_LIST[i]}" \
	'BEGIN{i=0} $1~dev {i++} END{print i}') > 0 )) ; then
	print "WARN: ${HDD_LIST[i]} appears to be part of ZFS pool. removing from list."
	unset HDD_LIST[i]
	continue
    fi
done
if (( ${#HDD_LIST[*]} == 0 )) ; then
    print "ERR: no hard drives found in list. aborting."
    exit 1
fi

# confirm with user before running destructive test on HDDs
print "ATTN: this hard drive testing process will wipe all data"
print "ATTN: on the following hard drives:"
print "ATTN: ${HDD_LIST[*]}"
print -n "ATTN: ARE YOU SURE YOU WANT TO PROCEED?: "
read PROCEED
if [[ $PROCEED != @(~(+i:yes|y|affirmative)) ]] ; then
    print "INFO: did not get affirmative response from user."
    print "INFO: hard drive testing aborted."
    exit 0
fi

# create "disks" subdirectory for collecting data
[[ ! -d ${BHTDIR}/disks ]] && mkdir ${BHTDIR}/disks

print "INFO: collecting SMART data from each drive."
typeset -A HDD_DATA
for D in ${HDD_LIST[*]}
do
    TMPFILE=$(mktemp /var/tmp/bht.XXXXXX)
    smartctl -a $D > $TMPFILE
    HDD_DATA[$D].type=$(print_HDD_type $TMPFILE)
    if [[ ${HDD_DATA[$D].type} == SAS ]] ; then
	smartctl --xall $D > $TMPFILE
    fi
    HDD_MODEL=$(awk -F: '$1~"^(Device Model|Product)" {gsub(" ","",$2); print $2}' $TMPFILE)
    HDD_SN=$(awk -F: '$1~"^Serial [Nn]umber" {gsub(" ","",$2); print $2}' $TMPFILE)
    # check that we have HDD_MODEL and HDD_SN before proceeding
    if [[ -z $HDD_MODEL ]] || [[ -z $HDD_SN ]] ; then
	print "ERR: could not identify model and/or serial number for $D. Abort."
	exit 1
    fi
    [[ ! -d ${BHTDIR}/disks/${HDD_MODEL}_${HDD_SN} ]] && mkdir ${BHTDIR}/disks/${HDD_MODEL}_${HDD_SN}
    mv -f $TMPFILE ${BHTDIR}/disks/${HDD_MODEL}_${HDD_SN}/pre-test.smart
    print "INFO: $HDD_MODEL / $HDD_SN"
    HDD_DATA[$D].dir=${BHTDIR}/disks/${HDD_MODEL}_${HDD_SN}
    HDD_DATA[$D].model=${HDD_MODEL}
    HDD_DATA[$D].serial=${HDD_SN}
    print "${HDD_DATA[$D].type}" > ${HDD_DATA[$D].dir}/hdd_type
    #HDD_DATA[$D].pre_smart=${BHTDIR}/disks/${HDD_MODEL}_${HDD_SN}/pre-test.smart
done

#print "HDD_DATA = ${HDD_DATA[*]}"
for D in ${!HDD_DATA[*]}
do
    print "INFO: Running badblocks on $D."
    print "INFO: ==> output to ${HDD_DATA[$D].dir}."
    print "INFO: ==> [${HDD_DATA[$D].model}|${HDD_DATA[$D].serial}]"
    run_badblocks $D "${HDD_DATA[$D].dir}" &
done

exit 0

##
## end of program
##