#!/bin/bash

# To-do:
# Incorporate override switch to override current process for specified drive
# Incorporate warning about raw disk I/O flag
# Way to detect previous failed run?

# Hard Drive Burn-in:
# https://forums.freenas.org/index.php?threads/how-to-hard-drive-burn-in-testing.21451/

# exec bash
clear

# User Editable Variables
Drive=$1
Save_Path=/mnt/SystemDataset

# Test if $Drive exists
if [ $(ls /dev/$Drive) ]; then
	echo "Drive $Drive exists."
else
	echo "Drive $Drive is not found."
	echo "Exiting..."
	exit -1
fi

# Get Drive Details
Drive_Model=$(smartctl -a /dev/$Drive | grep "Model" | awk '{print $3}')
Drive_Serial_Number=$(smartctl -a /dev/$Drive | grep "Serial" | awk '{print $3}')

# Test if $Drive is already in use
Test_String="^$Drive$"
Drive_In_Use=$(gpart show | grep "=>" | awk '{print $4}' | grep "$Test_String")
if [ -z $Drive_In_Use ]; then
	echo "Drive $Drive not found in 'gpart show'."
	echo "Drive $Drive is cleared for testing..."
else
	echo "Drive $Drive is currently in use according to 'gpart show'."
	echo "May I suggest a quick wipe using the GUI?  That will remove it from 'gpart show'."
	echo "Exiting..."
	exit -1
fi

# Test if the "In-Progress" file already exists
File_Count=$(ls -l *$Drive_Serial_Number*In-Progress | wc -l)
if [ $File_Count -gt 0 ]; then
	echo "Cannot begin new test on $1.  Test already in progress"
	exit -1
fi

# Test if inside TMUX already
if [ -n "$TMUX" ]; then 
	echo "Inside TMUX.  Continuing Execution..."; 
else
	echo "Spinning up new TMUX Session '$1' and exiting."
	Command="$(readlink -nf $0) $1"
	echo "Parameters: tmux new-session -s $1 -d '$Command'"
	tmux new-session -s $1 -d "$Command"
	tmux ls
	exit
fi

# Create "In-Progress" Status File
In_Progress_File_Name=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_In-Progress"
touch $In_Progress_File_Name


# Cancel previous test
smartctl -X /dev/$Drive

# Save SMART Details to Disk
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive
touch $Save_File"_SMART_Details.txt"
smartctl -a /dev/$Drive > $Save_File"_SMART_Details.txt"


# Start short self-test
CompleteDate="$(smartctl -t short /dev/$Drive | grep after)"
CompleteDate=${CompleteDate#*after*}
echo Performing Short Self-Test on $Drive_Model S/N:$Drive_Serial_Number
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_SMART_Test_Short"
echo $CompleteDate > $Save_File


# Sleep during self-test
echo Sleeping until: $CompleteDate
current_epoch=$(date +%s)
target_epoch=$(date -j -f "%a %b %d %T %Z %Y" "`echo $CompleteDate`" +%s)
echo $target_epoch - $current_epoch + 2
sleep_seconds=$(( $target_epoch - $current_epoch + 2))
#smartctl -X /dev/$Drive
sleep $sleep_seconds
echo SMART short test completed
echo
echo 


# Start conveyance self-test
# Sleep during self-test
#clear
CompleteDate=$(smartctl -t conveyance /dev/$Drive | grep after | cut -c 25-100)
echo Performing Conveyance Self-Test on $Drive_Model S/N:$Drive_Serial_Number
echo Sleeping until: $CompleteDate
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_SMART_Test_Conveyance"
echo $CompleteDate > $Save_File
current_epoch=$(date +%s)
target_epoch=$(date -j -f "%a %b %d %T %Z %Y" "`echo $CompleteDate`" "+%s")
sleep_seconds=$(( $target_epoch - $current_epoch + 2))
#smartctl -X /dev/$Drive
sleep $sleep_seconds
echo SMART conveyance test completed
echo
echo

# Start long self-test
# Sleep during self-test
#clear
CompleteDate=$(smartctl -t long /dev/$Drive | grep after | cut -c 25-100)
echo Performing Long Self-Test on $Drive_Model S/N:$Drive_Serial_Number
echo Sleeping until: $CompleteDate
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_SMART_Test_Long"
echo $CompleteDate > $Save_File
current_epoch=$(date +%s)
target_epoch=$(date -j -f "%a %b %d %T %Z %Y" "`echo $CompleteDate`" "+%s")
sleep_seconds=$(( $target_epoch - $current_epoch + 2))
#smartctl -X /dev/$Drive
sleep $sleep_seconds
echo SMART long test completed
echo
echo

# Enable raw disk I/O
sysctl kern.geom.debugflags=0x10
RAW_Disk_IO_File_Name=$Save_Path/RAW_Disk_IO_Enabled
touch $RAW_Disk_IO_File_Name
# Default value = 0


# Start Destructive BadBlocks test
# NOTE:  This is a blocking command.  Additional commands
#        will not run until this command is finished.
#        No need to use a sleep command.
#clear
echo Performing Destructive BadBlocks Test on $Drive_Model S/N:$Drive_Serial_Number
echo This test may take days to complete.
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_BadBlocks"
echo "tmux attach -t $Drive" > $Save_File
badblocks -wsv /dev/$Drive >> $Save_File
echo
echo

# Start long self-test
# Sleep during self-test
#clear
CompleteDate=$(smartctl -t long /dev/$Drive | grep after | cut -c 25-100)
echo Performing Long Self-Test on $Drive_Model S/N:$Drive_Serial_Number
echo Sleeping until: $CompleteDate
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_SMART_Test_Long"
echo $CompleteDate > $Save_File
current_epoch=$(date +%s)
target_epoch=$(date -j -f "%a %b %d %T %Z %Y" "`echo $CompleteDate`" "+%s")
sleep_seconds=$(( $target_epoch - $current_epoch + 2))
#smartctl -X /dev/$Drive
sleep $sleep_seconds
echo SMART long test completed
echo
echo

# Save SMART Details to Disk
Save_File=$Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive
touch $Save_File"_SMART_Details.txt"
smartctl -a /dev/$Drive > $Save_File"_SMART_Details.txt"

# Remove "In-Progress" Status File
rm $In_Progress_File_Name

# Create "Completed" Status File
touch $Save_Path/$Drive_Model"_"$Drive_Serial_Number"_"$(date -u +"%Y-%m-%dT%H.%M.%SZ")"_"$Drive"_Completed"

# If no other "In-Progress" files exist, then reset the RAW Disk IO flag to zero
# and remove the RAW Disk IO Flag file
File_Count=$(ls -l *In-Progress | wc -l)
if [ $File_Count -gt 0 ]; then
	echo "Cannot disable RAW Disk IO Flag:"
       echo "  sysctl kern.geom.debugflags"
else
	sysctl kern.geom.debugflags=0
	rm $RAW_Disk_IO_File_Name
	echo "Successfully disabled RAW Disk IO Flag:"
       echo "  sysctl kern.geom.debugflags"
fi

#exit