#!/bin/bash
#!/usr/bin/env bash

get_netid () {
    netid=
	echo "Enter netid: "
	read netid
	while [[ $netid =~ ^\s*$ ]];
	do
		echo "Your netid cannot be blank. Enter netid: "
		read netid
	done
	clear
	}

get_ticket_number () {
    ticket_number=
	echo "Enter ticket number: "
	read ticket_number
    echo $ticket_number
	clear
}

confirm_message () {
    local typed=
	while [[ ! $typed = "$2" ]];
	do
	    echo $1
		read typed
	done
}

get_eraselevel () {
    eraselevel=
	echo "What level of erase are you performing on the drive? (0 = baseline, 1 = secure erase, 2 = decommission):"
    echo "Type 'help' for an explanation of each level."
	read eraselevel
	while [[ ! $eraselevel = "0" ]] && [[ ! $eraselevel = "1" ]] && [[ ! $eraselevel = "2" ]];
	do
        if [[ $eraselevel = "help" ]];
        then
            echo "Level 0 / Baseline: clears the partitions of the drive. Data can still be recovered but this is the optimal level for simple reuse."
            echo "Level 1 / Secure Erase: HDD: Scrambles the data on the drive over 3 passes. SSD: Sends a signal to the drive to forget the encryption key, making it unreadable."
            echo "Level 2 / decommission: Scrambles the data on the drive over 7 passes. Afterward, put the drive in the bin near the Hardware office to have it destroyed."
        else
		    echo "Invalid level, correct values can be 0, 1, or 2"
        fi
        read eraselevel
	done
}

get_devicetype () {
    devicetype=
	echo "What is the device type? (HDD_SATA, SSD_NVME, SSD_SATA):"
    echo "Type 'help' for an explanation of each type."
	read devicetype
	while [[ ! $devicetype = "HDD_SATA" ]] && [[ ! $devicetype = "SSD_NVME" ]] && [[ ! $devicetype = "SSD_SATA" ]];
	do
        if [[ $devicetype = "help" ]];
        then
            echo "HDD_SATA: Spinning disk platters on a SATA connection. Typically 3.5 in or 2.5 in."
            echo "SSD_SATA: Solid State drive on a SATA connection. Typically 2.5 in."
            echo "SSD_NVME: Solid State drive on a M.2 connection. Looks like a small PCB."
        else
		    echo "Invalid type, correct values can be HDD_SATA, SSD_NVME, SSD_SATA"
        fi
        read devicetype
	done
}

get_device () {
    device=
    pcie_enable
    echo "Printing current attached devices..."
    echo ""
    if [[ $devicetype = "HDD_SATA" ]] || [[ $devicetype = "SSD_SATA" ]]; 
    then
        lsblk | grep -E '^NAME|^sd'
    elif [[ $devicetype = "SSD_NVME" ]];
    then
        lsblk | grep -E '^NAME|^nvme'
    fi
    echo ""
    echo "Which is the device from this list? (Type 'help' for help.)"

    while [[ $device =~ ^\s*$ ]];
	do
	    verify_device
	done
}

verify_device (){
	read device
	while [[ ! $device =~ ^sd[a-z]$ ]] && [[ ! $device =~ ^nvme0n[0-9]$ ]];
	do
        if [[ $device = "help" ]];
        then
            echo "The UNIX filesystem thinks of storage devices as directories, which are under /dev/"
            echo "If you have a SATA connection, you will be looking for sd{a-z}."
            echo "If you have a NVME connection, you will be looking for nvme0n{0-9}."
        else
		    echo "Invalid format, device should follow naming conventions. (i.e. sd{a-z}, nvme0n{0-9})"
        fi
		read device
	done
	if [[ $devicetype = "HDD_SATA" ]] || [[ $devicetype = "SSD_SATA" ]]; 
	then
        if [[ $device =~ ^nvme0n[0-9]$ ]];
        then
            echo "Device was specified to be a SATA HDD or SSD, but a NVME device was chosen."
            device=
        fi
	elif [[ $devicetype = "SSD_NVME" ]];
	then
	    if [[ $device =~ ^sd[a-z]$ ]];
        then
            echo "Device was specified to be a NVME SSD, but a SATA device was chosen.  Please ensure the device is plugged into the motherboard via PCIe slot and not SATA."
            device=
        fi
	fi

    if [ -e "/dev/${device}" ]; then
        echo "Picking device /dev/${device}."
    else
        echo "/dev/${device} does not exist, please ensure you are typing the device name correctly."
        device=
    fi
}

erase_device_lv0 () {
    echo "Wiping partition tables of ${devicetype} : ${device}."
    
    ret_value=$(sgdisk -g -Z /dev/$device &> /dev/null; echo $?)
    
    if [[ ! $ret_value = "0" ]]
    then
        case $ret_value in
            "1")
            echo "There was an issue with the argument. Ensure you typed the drive name correctly."
            ;;
        
            "2")
            echo "An error occurred while reading the partition table."
            ;;
        
            "3")
            echo "Non-GPT disk detected and no -g option, but operation requires a write action."
            ;;
            
            "4")
            echo "An error prevented saving changes."
            ;;
            
            "5")
            echo "An error occurred while reading standard input."
            ;;
            
            "8")
            echo "Disk replication operation (-R) failed."
            ;;
        
            *)
            echo "Unspecified error."
            ;;
        esac
    else
        echo "Partition tables wiped. Device is ready for reuse."
    fi
}

erase_device_lv1 () {
    #HDD = 3-Pass DOD Short
    #SSD Sata = hdparm secure erase
    #SSD NVMe = nvme secure erase 
    if [[ $devicetype = "HDD_SATA" ]] && [[ ! $eraselevel = "2" ]];
    then
        echo "Running 3-Pass DOD Short on ${devicetype} : ${device}."
        echo "This may take a while!"

        ret_value=$(nwipe -m dodshort --autonuke --nogui /dev/$device &> /dev/null; echo $?)
        
        if [[ ! $ret_value = "0" ]]
        then
            case $ret_value in
                "-1")
                echo "Unspecified error while nwiping disk."
                ;;

                *)
                echo "Unspecified error while nwiping disk."
                ;;
            esac
        else
            echo "Finished wiping data with nwipe."
        fi
    elif [[ $devicetype = "SSD_SATA" ]];
    then
        echo "Running hdparm format on ${devicetype} : ${device}."
        echo "Setting security pass to NULL."

        #Set the password of the drive to NULL
        ret_value=$(hdparm --user-master m --security-set-pass NULL /dev/$device &> /dev/null; echo $?)
        
        if [[ ! $ret_value  = "0" ]];
        then
            hdparm_check_error "${ret_value}"
            echo "Failed to set password of drive. Attempting to continue."
        else
            echo "Set password of drive to NULL."
        fi

        #Attempt a secure erase.
        ret_value=$(hdparm --user-master m --security-erase-enhanced NULL /dev/$device &> /dev/null; echo $?)

        if [[ ! $ret_value  = "0" ]];
        then
            hdparm_check_error "${ret_value}"
        else
            echo "Enhanced secure erase failed. This is not unexpected depending on the drive type."
            echo "Attempting normal erase instead."

            #Attempt a secure erase, less enhanced.
            ret_value=$(hdparm --user-master m --security-erase NULL /dev/$device &> /dev/null; echo $?)
            
            if [[ ! $ret_value  = "0" ]];
            then
                hdparm_check_error "${ret_value}"
            else
                echo "Failed to secure erase. Please contact your supervisor."
            fi
        fi
    elif [[ $devicetype = "SSD_NVME" ]];
    then
        echo "Running nvme_cli format on ${devicetype} : ${device}."
        #Erase using nvme cli to send secure erase command to drop the crypto stuff.
        ret_value=$(nvme format --force -r -s 2 /dev/$device &> /dev/null; echo $?)
        
        if [[ ! $ret_value = "0" ]]
        then
            echo "Failed to cryptographically erase drive, blanket erasing instead."
            #Sad, can't do it. Blank out drive instead using the drive controller.
            ret_value=$(nvme format --force -r -s 1 /dev/$device &> /dev/null; echo $?)
            
            if [[ ! $ret_value = "0" ]]
            then
                echo "Failed to cryptographically erase or blanket erase. Please contact your supervisor."
            else
                echo "Blanket erased drive."
            fi
        else
            echo "Cryptographically erased drive."
        fi
    else
    echo "There was an issue with the device type: ${devicetype}. Ensure you typed the drive name correctly."
    fi
}

erase_device_lv2 () {
    #7-Pass DOD Wipe through nwipe.
    echo "Running 7-Pass DOD on ${devicetype} : ${device}."
    echo "This may take a while!"

    ret_value=$(nwipe -m dod --autonuke --nogui /dev/$device &> /dev/null; echo $?)
    
    if [[ ! $ret_value = "0" ]]
    then
        case $ret_value in
            "-1")
            echo "Unspecified error while nwiping disk."
            ;;

            *)
            echo "Unspecified error while nwiping disk."
            ;;
        esac
    else
        echo "Finished wiping data with nwipe."
        echo "Please put device in decommission bin in the hardware room."
    fi
}

hdparm_check_error () {
    case $1 in
        "1")
        echo "Operation not permitted."
        ;;
    
        "2")
        echo "No such file or directory."
        ;;
    
        "5")
        echo "I/O error."
        ;;
        
        "16")
        echo "Device or resource busy."
        ;;
        
        "22")
        echo "Invalid argument."
        ;;
    
        *)
        echo "Unspecified error."
        ;;
    esac
}

#Rescan for PCIe devices
pcie_enable (){
    echo 1 > /sys/bus/pci/rescan
    echo "Enabled PCIe card, sleeping for 5 seconds."
    sleep 5
}

#Remove the PCIe NVMe reader
pcie_disable (){
    echo 1 > /sys/bus/pci/devices/0000:02:00.0/remove
    echo "Disabled PCIe card."
}

main (){
    # get_netid
    # get_ticket_number
    
    echo "By running this script, you are confirming that it has been two weeks past the date written on the slip attached to the storage device. DO NOT erase the device before two weeks have passed."
    
    confirm_message "Please type 'confirm' to acknowledge you have read this and that it has been two weeks." "confirm"
    
    get_eraselevel #sets $eraselevel to 0-2 based on how to erase. Higher levels include lower levels.
    get_devicetype #sets $devicetype to HDD_SATA/SSD_SATA/SSD_NVME
    get_device #sets $device to one of the /dev/xyz devices.
    
    clear

    echo "Starting process for erasing ${devicetype} : ${device} with erase level ${eraselevel}."

    if [[ $eraselevel = "1" ]] || [[ $eraselevel = "2" ]]; 
	then
	    local randomconf
        randomconf=$(shuf -i 10000-99999 -n 1)
	    confirm_message "This level is destructive and irreversible. Please type ${randomconf} to continue" "${randomconf}"
	else
        confirm_message "Please type 'confirm' to begin." "confirm"
    fi
    
    
    case $eraselevel in
        "0")
        erase_device_lv0 #Wipe partition
        ;;
        
        "1")
        erase_device_lv1 #Secure Erase
        erase_device_lv0 #Wipe partition
        ;;
        
        "2")
        erase_device_lv1 #Secure Erase (skips nwipe)
        erase_device_lv2 #7-Pass nwipe
        erase_device_lv0 #Wipe partition
        ;;
    
        *)
        echo "Invalid Erase Level"
        ;;
    esac    

    sleep 3

    echo "Finished erasing ${devicetype} : ${device} with erase level ${eraselevel}."

    pcie_disable
}

main