From 90c92bbc3a531c3fedda00321486fab48c1c8440 Mon Sep 17 00:00:00 2001
From: dak180 <dak180@users.sourceforge.net>
Date: Mon, 5 Feb 2024 10:39:11 -0500
Subject: [PATCH] Add an option to print a list of drives.

---
 README.md      |  6 ++++-
 disk-burnin.sh | 71 +++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 19d465b..830d868 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,10 @@ The only required command-line argument is the device specifier, e.g.:
 
 ...will run the burn-in test in 'dry run mode' on device /dev/ada0
 
+If you are not sure of the device specifiers currently on the system you can get a list with the `-L` switch.
+
+`./disk-burnin.sh -L`
+
 You can run the script in 'dry run mode' to check the sleep duration calculations and to insure that the sequence of commands suits your needs. In 'dry runs' the script does not actually perform any SMART tests or invoke the `sleep` or `badblocks` programs. The script is distributed with 'dry runs' enabled, so you will need to pass the `-t` switch, in order to actually perform tests on drives.
 
 The script can automatically invoke `tmux` sessions for a space separated list of drive specifiers:
@@ -65,6 +69,6 @@ Requires the smartmontools, available at https://www.smartmontools.org
 
 Uses: `grep`, `pcregrep`, `awk`, `sed`, `tr`, `sleep`, `badblocks`
 
-Tested with the static analysis tool at https://www.shellcheck.net to insure that the code is POSIX-compliant and free of issues.
+Tested with the static analysis tool at https://www.shellcheck.net to insure that the code is free of issues.
 
 Written by Keith Nash, March 2017. Modified by Yifan Liao and dak180.
diff --git a/disk-burnin.sh b/disk-burnin.sh
index e47b5c7..9da9234 100755
--- a/disk-burnin.sh
+++ b/disk-burnin.sh
@@ -86,14 +86,35 @@ Options:
 
 -l
 	Log files directory.
+
+-L
+	Print list of Drive Device Specifiers and exit
 ...
 EOF
 }
 
+# Check if we are running on BSD
+if [[ "$(uname -mrs)" =~ .*"BSD".* ]]; then
+	systemType="BSD"
+fi
+
+# Get the version numbers for smartctl
+major_smartctl_vers="$(smartctl -jV | jq -Mre '.smartctl.version[] | values' | sed '1p;d')"
+minor_smartctl_vers="$(smartctl -jV | jq -Mre '.smartctl.version[] | values' | sed '2p;d')"
+if [[ "${major_smartctl_vers}" -gt "7" ]]; then
+	smartctl_vers_74_plus="true"
+elif [[ "${major_smartctl_vers}" -eq "7" ]] && [[ "${minor_smartctl_vers}" -ge "4" ]]; then
+	smartctl_vers_74_plus="true"
+elif [ -z "${major_smartctl_vers}" ]; then
+	echo "smartctl version 7 or greater is required" >&2
+	smartctl -V
+	exit 1
+fi
+
 Log_Dir="."
 Dry_Run=1
 
-while getopts ":d:m:l:b:th" OPTION; do
+while getopts ":d:m:l:Lth" OPTION; do
 	case "${OPTION}" in
 		d)
 			driveID="${OPTARG}"
@@ -104,6 +125,10 @@ while getopts ":d:m:l:b:th" OPTION; do
 		l)
 			Log_Dir="${OPTARG}"
 		;;
+		L)
+			get_drive_list
+			exit 0
+		;;
 		t)
 			Dry_Run=0
 		;;
@@ -132,6 +157,12 @@ commands+=(
 tmux
 )
 fi
+if [ "${systemType}" = "BSD" ]; then
+commands+=(
+sysctl
+# nvmecontrol
+)
+fi
 for command in "${commands[@]}"; do
 	if ! type "${command}" &> /dev/null; then
 		echo "${command} is missing, please install"
@@ -252,6 +283,44 @@ Poll_Interval="15"
 #
 ######################################################################
 
+function get_drive_list() {
+	local drives
+	local localDriveList
+
+	if [ "${systemType}" = "BSD" ]; then
+		localDriveList="$(sysctl -n kern.disks | sed -e 's:nvd:nvme:g')"
+	else
+		# shellcheck disable=SC2010
+		localDriveList="$(ls -l "/sys/block" | grep -v 'devices/virtual' | sed -e 's:[[:blank:]]\{1,\}: :g' | cut -d ' ' -f "9" | sed -e 's:n[0-9]\{1,\}$::g' | uniq )"
+		# lsblk -n -l -o NAME -E PKNAME | tr '\n' ' '
+	fi
+
+	if [ "${systemType}" = "BSD" ]; then
+		# This sort breaks on linux when going to four leter drive ids: "sdab"; it works fine for bsd's numbered drive ids though.
+		readarray -t "drives" <<< "$(for drive in ${localDriveList}; do
+			if [ "${smartctl_vers_74_plus}" = "true" ] && [ "$(smartctl -ji "/dev/${drive}" | jq -Mre '.smart_support.enabled | values')" = "true" ]; then
+				printf "%s\n" "${drive}"
+			elif smartctl -i "/dev/${drive}" | sed -e 's:[[:blank:]]\{1,\}: :g' | grep -q "SMART support is: Enabled"; then
+				printf "%s\n" "${drive}"
+			elif grep -q "nvme" <<< "${drive}"; then
+				printf "%s\n" "${drive}"
+			fi
+		done | sort -V | sed '/^nvme/!H;//p;$!d;g;s:\n::')"
+	else
+		readarray -t "drives" <<< "$(for drive in ${localDriveList}; do
+			if [ "${smartctl_vers_74_plus}" = "true" ] && [ "$(smartctl -ji "/dev/${drive}" | jq -Mre '.smart_support.enabled | values')" = "true" ]; then
+				printf "%s\n" "${drive}"
+			elif smartctl -i "/dev/${drive}" | sed -e 's:[[:blank:]]\{1,\}: :g' | grep -q "SMART support is: Enabled"; then
+				printf "%s\n" "${#drive} ${drive}"
+			elif grep -q "nvme" <<< "${drive}"; then
+				printf "%s\n" "${#drive} ${drive}"
+			fi
+		done | sort -Vbk 1 -k 2 | cut -d ' ' -f 2 | sed '/^nvme/!H;//p;$!d;g;s:\n::')"
+	fi
+
+	echo "${drives[@]}"
+}
+
 function echo_str() {
 	echo "$1" | tee -a "${Log_File}"
 }