From e4378cd040533776ff18d7c793643a9ec01e5aaa Mon Sep 17 00:00:00 2001
From: Austin Horstman <khaneliman12@gmail.com>
Date: Fri, 14 Jan 2022 19:17:48 -0600
Subject: [PATCH] Add input validation to diskpart

---
 startup.sh | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/startup.sh b/startup.sh
index 5505e0e..392e9cf 100644
--- a/startup.sh
+++ b/startup.sh
@@ -125,19 +125,31 @@ esac
 
 # selection for disk type
 diskpart () {
-# show disks present on system
-lsblk -n --output TYPE,KNAME,SIZE | awk '$1=="disk"{print NR,"/dev/"$2" - "$3}' # show disks with /dev/ prefix and size
 echo -ne "
 ------------------------------------------------------------------------
-    THIS WILL FORMAT AND DELETE ALL DATA ON THE DISK             
-    Please make sure you know what you are doing because         
-    after formating your disk there is no way to get data back      
+    THIS WILL FORMAT AND DELETE ALL DATA ON THE DISK
+    Please make sure you know what you are doing because
+    after formating your disk there is no way to get data back
 ------------------------------------------------------------------------
 
-Please enter full path to disk: (example /dev/sda):
 "
-read option
-echo "DISK=$option" >> setup.conf
+
+PS3='
+Select the disk to install on: '
+options=($(lsblk -n --output TYPE,KNAME,SIZE | awk '$1=="disk"{print "/dev/"$2"|"$3}'))
+select opt in "${options[@]}"
+do
+
+# Positive check
+if echo -e '%s\n' "${options[@]}" | grep -Fqw ${opt} 2> /dev/null; then
+    echo -e "\n${opt%|*} selected \n"
+    echo "DISK=${opt%|*}" >> setup.conf
+    break
+else
+    echo -e "\nInvalid selection, please try again. \n"
+fi
+
+done
 
 drivessd
 set_option DISK $option