#!/usr/bin/env python

import os, sys, re, stat
import commands

bold = "\033[1m"
green = "\033[01;32m"
red = "\033[01;31m"
yellow = "\033[01;33m"
nocolor = "\033[0m"

sdlist = []
sdarray = commands.getoutput('ls -A /dev/sd[a-z]')
mounted = commands.getoutput('df -l --output=source /root | grep -v Filesystem')

for sd in sdarray.split():
    mode = os.stat(sd).st_mode
    if stat.S_ISBLK(mode):
        # is /dev/sd* and is a block device
        sdlist.append(sd)
    if os.path.exists(str(sd) + "1"):
        print("{0} has partitions!").format(yellow + sd + nocolor)

# if the sd* is currently mounted, remove it from the list (sdlist) of drives to run on.
for dev in sdlist:
    if dev in mounted:
        print(yellow + "%s is mounted and will be ignored." + nocolor) % (dev)
        sdlist.remove(dev)

for dev in sdlist:
    print("Started test on {0}{1}{2}...").format(bold, dev, nocolor)
    # begin tests on each dev
    commands.getoutput("smartctl -t short " + dev)

def status(dev):
    isgoing = commands.getoutput("smartctl -a " + dev + " | grep -i 'routine in progress'")
    if not 'routine in progress' in isgoing:
        devpassed = commands.getoutput("smartctl -H " + dev + " | grep 'PASSED' | awk 'NF>1 {print $NF}'")
        devserial = commands.getoutput("smartctl -i " + dev + " | grep 'Serial Number' | awk 'NF>1 {print $NF}'")
        devmodel = commands.getoutput("smartctl -i " + dev + " | grep 'Device Model' | awk 'NF>1 {print $NF}'")
        devhours = commands.getoutput("smartctl -A " + dev + " | grep 'Power_On_Hours' | sed 's/(.*)//g' | awk 'NF>1 {print $NF}'")
        if devpassed != "PASSED":
            devpassed = str(red + "FAILED" + nocolor)
        print("Test completed on {0}{1}{2}. Result: {3}").format(bold, dev, nocolor, devpassed)
        print("Serial: {0} | Model: {1} | Power On Hours: {2}\n").format(devserial, devmodel, devhours)
        sdlist.remove(dev)

while len(sdlist) > 0:
    for dev in sdlist:
        status(dev)

print("{0}All drive tests done.{1}\n").format(green, nocolor)