#!/usr/bin/env python3

# Marc Schoechlin <ms@256bit.org>

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

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


def get_root():
  return os.environ.get("ZABBIX_EXTENTIONS_OS_TOOLS_CHROOT", "/")


def is_device(dev):

    m = re.match(r"^([a-zA-Z0-9]+)([\sa-zA-Z0-9,-_\/]+)?$", dev)
    if m:
      dev = "/dev/%s" % m.group(1)
      try:
          mode = os.stat(dev).st_mode
      except:
          return False
      return stat.S_ISBLK(mode)
    else:
      sys.stderr.write("%s :: get_health :: FAILED - device >>%s<< parameter is not valid\n" % (__file__, device))
      return False

def get_health(device):

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

    cmd = f"chroot {get_root()} smartctl -H /dev/{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:|Health Status:)\s*(.*)$", line)
        if not m:
            continue
        result = m.group(2)

    if result is None:
        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_health :: FAILED - device >>%s<< does not exist\n" % (__file__, device))
        return

    cmd = f"chroot {get_root()} smartctl -A /dev/{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(r"\(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_health :: FAILED - device >>%s<< does not exist\n" % (__file__, device))
        return

    cmd = f"chroot {get_root()} smartctl -i /dev/{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"^%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(r"/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
