#!/usr/bin/env python3


import pySMART as pys
from math import log
import sys

# Faal als deze atttibutes onder deze waarden komen.
# Hier zet je 1 van drie neer:
# - Een ATA 'ATTRIBUTE_NAME' uit `smartctl -A`
#   Let op, kies in 'attributes_use_raw' of je de 'RAW' of juist de 'VALUE' wilt gebruiken.
# - De SCSI diagnostics (SAS) attributen van de smart interface
#   https://github.com/truenas/py-SMART/blob/master/pySMART/diagnostics.py
# - Een NVME waarde uit de pySMART nvme interface
#   https://github.com/truenas/py-SMART/blob/master/pySMART/interface/nvme/__init__.py

attributes_fail_below = {
    "Wear_Leveling_Count": 50,
    "availableSpare": 50,
}

attributes_fail_above = {
    "Power_Cycle_Count": 1500,
    "Reallocated_Sector_Ct": 100,
    "Current_Pending_Sector": 0,
    "Uncorrected_Reads": 5,
    "Uncorrected_Writes": 0,
    "Uncorrected_Verifies": 0,
    "Reported_Uncorrect": 0,
    "Uncorrectable_Error_Cnt": 0,
    "Airflow_Temperature_Cel": 45,
    "Temperature_Celsius": 45,
    "End-to-End_Error": 100,
}

# Gebruik (voor desmart ATA) juist de RAW ipv de VALUE van deze atributen.
attributes_use_raw = [
    "Power_Cycle_Count",
    "Reallocated_Sector_Ct",
    "Current_Pending_Sector",
    "Reported_Uncorrect",
    "Uncorrectable_Error_Cnt",
    "Airflow_Temperature_Cel",
    "Temperature_Celsius",
]

# Een simpel lijstje met slechte firmwareversies.
bad_firmware = [
    "JXTE004Q",
]

# Allemaal extra logging aanzetten
verbose = False

# Script start hier, config is klaar.

# Alle nagios statussen.
statii = {
    "OK": 0,
    "WARNING": 1,
    "CRITICAL": 2,
    "UNKNOWN": 3,
}


def check_smart():

    global statii, attributes_fail_below, attributes_fail_below, bad_firmware

    def maybeebytes(size: int, si: bool = True) -> str:
        base = 1000 + 24 * (not si)
        order = int(log(size, base))
        factor = " KMGTPEZY"[order].strip()
        maybee = " i"[si].strip()
        newsize = round(size / base**order)
        return f"{newsize} {factor}{maybee}B"

    def verbose(message):
        global verbose
        if verbose:
            print(message)

    # Spindels en SSD's
    def check_custom_attributes(attrs: list) -> str:
        # Loop door alle attbibutes, en check of ze in attibute_fail staan.
        global attributes_use_raw
        for attr in attrs:
            if attr.name in attributes_use_raw:
                # Tot nu toe alleen bij temperatuur gezien, staat tekst achter het getal.
                attrval = attr.raw.split(' ', 1)[0]
            else:
                attrval = attr.value

            if (
                attr.name in attributes_fail_below
                and int(attrval) < attributes_fail_below[attr.name]
            ):
                return f"{attr.name} below custom value: {attributes_fail_below[attr.name]}/{attrval}"
            if (
                attr.name in attributes_fail_above
                and int(attrval) > attributes_fail_above[attr.name]
            ):
                return f"{attr.name} above custom value: {attributes_fail_above[attr.name]}/{attrval}"
        return ""

    # Ander smart smaakje voor SSD's
    def check_custom_diagnostics(diags: "pys.Diagnostics") -> str:
        for attr, val in attributes_fail_below.items():
            rval = getattr(diags, attr, val)
            if rval != None and rval < val:
                return f"{attr} below custom value: {rval}/{val}"
        for attr, val in attributes_fail_above.items():
            rval = getattr(diags, attr, val)
            if rval != None and rval > val:
                return f"{attr} above custom value: {rval}/{val}"
        return ""

    # NVME's
    def check_custom_if_attributes(attrs: pys.interface.nvme.NvmeAttributes) -> str:
        # Andersom van SSD's, loop door de attributeS_fail en check of die waarden in de interface staan.
        for attr, val in attributes_fail_below.items():
            rval = getattr(attrs, attr, val)
            if rval < val:
                return f"{attr} below custom value: {rval}/{val}"
        for attr, val in attributes_fail_above.items():
            rval = getattr(attrs, attr, val)
            if rval > val:
                return f"{attr} above custom value: {rval}/{val}"
        return ""

    def gettype(dev: pys.device.Device) -> str:
        if dev.is_ssd:
            if dev.dev_interface == 'nvme':
                return "NVME"
            return "SSD"
        return "HDD"


    status = statii["OK"]
    message = ""
    skip = 0

    devlist = pys.DeviceList()
    devs = devlist.devices

    verbose(f"found {len(devs)} devices")

    for dev in devs:
        devstring = f"Device {dev.name}, {dev.vendor} model: {dev.model}, size: {maybeebytes(dev.size)}, sn: {dev.serial}, {gettype(dev)} on interface {dev.interface}"
        verbose(devstring)

        # debian10 heeft te oude python(3.7), daar zit PEP572 niet in zit. Dus raar oplossen.
        custval = check_custom_attributes([x for x in dev.attributes if x != None])
        custdiags = check_custom_diagnostics(dev.diagnostics)
        custvalif = check_custom_if_attributes(dev.if_attributes)

        if not (dev.smart_capable and dev.smart_enabled):
            verbose("-> No smart data")
            skip += 1
        elif dev.assessment != "PASS":
            verbose("-> FAILED")
            status = statii["CRITICAL"]
            message += f"FAILED: {devstring}: {dev.assessment} \n"
        elif custval != "":
            verbose("-> FAILING")
            status = statii["WARNING"]
            message += f"WARNING: {devstring}: {custval} \n"
        elif custdiags != "":
            verbose("-> FAILING")
            status = statii["WARNING"]
            message += f"WARNING: {devstring}: {custdiags} \n"
        elif custvalif != "":
            verbose("-> FAILING")
            status = statii["WARNING"]
            message += f"WARNING: {devstring}: {custvalif} \n"
        elif dev.firmware in bad_firmware:
            verbose("-> BAD FIRMWARE")
            status = statii["WARNING"]
            message += f"WARNING: {devstring}: BAD FIRMWARE{dev.firmware} \n"
        else:
            verbose("-> OK")

    if status == statii["OK"]:
        message = (
            f"OK: {len(devs)- skip} device(s) passing SMART check, {skip} skipped. 🎉"
        )

    return (status, message)


if __name__ == "__main__":
    verbose_args = {'--verbose', '-v', 'verbose'}
    if len(verbose_args & set(sys.argv)) > 0:
        verbose = True

    s, m = check_smart()
    print(m)
    sys.exit(s)