#!/usr/bin/env python3

# Marc Schoechlin <ms@256bit.org>

import re
import os
import sys
import argparse
import subprocess
import stat

####################################################################################
#### HELPERS


def is_device(dev):
    dev = "/dev/%s" % dev
    try:
        mode = os.stat(dev).st_mode
    except:
        return False
    return stat.S_ISBLK(mode)

def get_health(device):

    if not is_device(device):
        sys.stderr.write("%s :: get_health :: FAILED - device /dev/%s does not exist\n" % (__file__, device))
        return

    cmd = "smartctl -H /dev/%s" % device
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = process.communicate()

    result = None
    for line in out.decode('utf8').splitlines():
        m = re.match(r"^.*test result:\s*(.*)$", line)
        if not m:
            continue
        result = m.group(1)

    if result is None or not "PASSED" in result:
        result = "FAILED: no health result"
        sys.stderr.write("%s :: get_health :: %s:\n--\n%s\n%s\n--\n" % (__file__, cmd, out.decode('utf8'), err.decode('utf8')))

    print(result)

def get_attribute(device, attribute):

    if not is_device(device):
        sys.stderr.write("%s :: get_attribute :: FAILED - device /dev/%s does not exist\n" (__file__, device))
        return

    cmd = "smartctl -A /dev/%s" % device
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = process.communicate()
    result = None
    #   1 Raw_Read_Error_Rate     0x000a   100   100   000    Old_age   Always       -       0
    for line in out.decode('utf8').splitlines():
        if attribute == "Temperature_Celsius":
            line = re.sub("\(Min\/Max \d+\/\d+\)", "", line) 

        m = re.match(r"^\s*\d+\s+([^\s]+?)\s+0x.*\s+(([^\s]+))\s*$", line)
        if not m:
            continue
        c_attrib = m.group(1)
        c_value = m.group(2)
        if attribute == c_attrib:
            result = c_value
    print(result)

def get_info(device, info):

    if not is_device(device):
        sys.stderr.write("%s :: get_info :: FAILED - device /dev/%s does not exist\n" (__file__, device))
        return

    cmd = "smartctl -i /dev/%s" % device
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = process.communicate()
    result = None
    for line in out.decode('utf8').splitlines():
        m = re.match("^%s:\s+(.*)$" % info, line)
        if not m:
            continue
        result = m.group(1)
    print(result)

####################################################################################
#### MAIN

parser = argparse.ArgumentParser(
    description='gather smart data'
)

parser.add_argument(
    '--debug',
    help='Output debug information',
    action='store_true',
)


parser_group = parser.add_mutually_exclusive_group(required=True)

parser.add_argument(
   '--device',
   nargs='?',
   type=str,
   help='the hardware device',
   default=None,
   required=True,
)

parser_group.add_argument(
    '--health',
    help='Health information',
    action='store_true',
)

parser_group.add_argument(
   '--attribute',
   nargs='?',
   type=str,
   help='the device attribute (smartctl -A <device>)',
   default=None,
)
parser_group.add_argument(
   '--info',
   nargs='?',
   type=str,
   help='the device info (smartctl -i <device>)',
)


args = parser.parse_args()
args.device = re.sub("/dev/+", "", args.device)

if args.health:
    get_health(args.device)

if args.attribute:
    get_attribute(args.device, args.attribute)

if args.info:
    get_info(args.device, args.info)

# vim: ai et ts=2 shiftwidth=2 expandtab
