#! /usr/bin/env python
#
# check_temperature - Nagios temperature check for DS18B20 sensor on RaspberryPi
#
# Version 1.1, latest version, documentation and bugtracker available at:
#              https://gitlab.lindenaar.net/scripts/nagios-plugins
#
# Copyright (c) 2017 Frederik Lindenaar
#
# This script is free software: you can redistribute and/or modify it under the
# terms of version 3 of the GNU General Public License as published by the Free
# Software Foundation, or (at your option) any later version of the license.
#
# This script is distributed in the hope that it will be useful but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, visit <http://www.gnu.org/licenses/> to download it.

from sys import exit
from os.path import basename, splitext
from glob import glob
from time import time, sleep
from argparse import ArgumentParser as StandardArgumentParser, FileType, \
              _StoreAction as StoreAction, _StoreConstAction as StoreConstAction
import logging

# Constants (no need to change but allows for easy customization)
VERSION="1.1"
PROG_NAME=splitext(basename(__file__))[0]
PROG_VERSION=PROG_NAME + ' ' + VERSION
SENSOR_SCALE=1000
W1_SENSOR_DEV_DIR =  '/sys/bus/w1/devices/'
W1_SENSOR_DEV_PREFIX = '28-'
W1_SENSOR_DEV_SUFFIX = '/w1_slave'
W1_SENSOR_READ_RETRIES=10
CPU_SENSOR_DEV = '/sys/class/thermal/thermal_zone0/temp'

LOG_FORMAT='%(levelname)s - %(message)s'
LOG_FORMAT_FILE='%(asctime)s - ' + LOG_FORMAT
LOGGING_NONE=logging.CRITICAL + 10
NAGIOS_OK = ( 'OK', 0)
NAGIOS_WARN = ( 'WARNING', 1)
NAGIOS_CRITICAL = ( 'CRITICAL', 2 )
NAGIOS_UNKNOWN = ( 'UNKNOWN', 3 )

# Setup logging
logging.basicConfig(format=LOG_FORMAT)
logging.addLevelName(LOGGING_NONE, 'NONE')
logger = logging.getLogger(PROG_NAME)
logger.setLevel(logging.CRITICAL)

################[ wrapper to stop ArgumentParser from exiting ]################
# based on http://stackoverflow.com/questions/14728376/i-want-python-argparse-to-throw-an-exception-rather-than-usage/14728477#14728477
# the only way to do this is overriding the error method and throw and Exception
class ArgumentParserError(Exception): pass

class ArgumentParser(StandardArgumentParser):
    """ArgumentParser not exiting with non-Nagios format message upon errors"""
    def error(self, message):
        raise ArgumentParserError(message)

##################[ Action to immediately set the log level ]##################
class SetLogLevel(StoreConstAction):
    """ArgumentParser action to set log level to provided const value"""
    def __call__(self, parser, namespace, values, option_string=None):
        logging.getLogger(PROG_NAME).setLevel(self.const)

####################[ Action to immediately log to a file ]####################
class SetLogFile(StoreAction):
    """ArgumentParser action to log to file (sets up FileHandler accordingly)"""
    def __call__(self, parser, namespace, values, option_string=None):
        super(SetLogFile, self).__call__(parser,namespace,values,option_string)
        formatter = logging.Formatter(LOG_FORMAT_FILE)
        handler = logging.FileHandler(values)
        handler.setFormatter(formatter)
        logger = logging.getLogger(PROG_NAME)
        logger.propagate = False
        logger.addHandler(handler)

###############################################################################

def convert_celcius(temp_read):
    """Converts raw temperature sensore value to degrees Celcius"""
    return float(temp_read) / float(SENSOR_SCALE)
CONVERT_CELCIUS = ( convert_celcius, 'C', 'Celcius' )


def convert_farenheit(temp_read):
    """Converts raw temperature sensore value to degrees Farenheit"""
    return float(temp_read * 9) / float(5 * SENSOR_SCALE) + 32.0
CONVERT_FARENHEIT = ( convert_farenheit, 'F', 'Farenheit' )


def isempty(string):
    """Checks whether string 'str' provided is unset or empty"""
    return string is None or len(string) == 0


def read_rpi_cpu_temp(args):
    """Reads CPU temperature and converts it to desired unit, returns temperature"""
    with open(args.file, 'r') as f:
        lines = f.readlines()
        logger.debug('Temperature sensor data read from %s: %s', f.name, lines)

    temp_read = int(lines[0])
    temp = args.converter[0](temp_read)
    logger.debug('Temperature sensor value %d is %.2f%s', temp_read,
    temp, args.converter[1])
    return temp, 1


def get_w1_sensor_device_filename(args, dev_dir=W1_SENSOR_DEV_DIR,
                            prefix=W1_SENSOR_DEV_PREFIX, suffix=W1_SENSOR_DEV_SUFFIX):
    """Auto-determine sensor datafile name (unless args.file is set)"""
    if isempty(args.file):
        search_pat = dev_dir + ('/' if dev_dir[-1]!='/' else '')
        search_pat+= prefix + '*' if isempty(args.serial) else '*' + args.serial
        logger.debug('looking for sensors with search pattern %s', search_pat)

        device_folders = glob(search_pat)
        if len(device_folders) == 1:
            filename = device_folders[0] + suffix
        else:
            if len(device_folders) == 0:
                errmsg = 'no supported temperature sensors in %s' % dev_dir
            else:
                serials = [ basename(x)[len(prefix):] if x.find(prefix)>=0
                                else basename(x) for x in device_folders ]
                errmsg = 'found multiple temperature sensors (%s), please '\
                         'specify which one to use' % ', '.join(serials)
            logger.critical(errmsg)
            raise ValueError(errmsg)
    else:
        filename = args.file
    logger.debug('using temperature sensor at %s', filename)
    return filename


def read_w1_temp(args):
    """Reads sensor data and converts it to desired unit, returns temperature"""
    device_file = get_w1_sensor_device_filename(args)
    lines=[ '' ]
    tries = 0
    while tries <= args.retries and lines[0].strip()[-3:] != 'YES':
        if tries > 0:
            logger.warn('Temperature sensor data not stable, reading once more')
            sleep(0.2)
        tries += 1
        with open(device_file, 'r') as f:
            lines = f.readlines()
            logger.debug('Temperature sensor data read from %s: %s', f.name, lines)

    if lines[0].strip()[-3:] != 'YES':
        errmsg = 'no stable temperature sensor data after %d tries' % tries
    else:
        equals_pos = lines[1].find('t=')
        if equals_pos == -1:
            errmsg = 'temperature sensor data format is not supported'
        else:
            temp_read = int(lines[1][equals_pos+2:])
            temp = args.converter[0](temp_read)
            logger.debug('Temperature sensor value %d is %.2f%s', temp_read,
                         temp, args.converter[1])
            return temp, tries

    logger.critical(errmsg)
    raise ValueError(errmsg)


def parse_args():
    """Parse command line and get parameters from environment, if present"""

    # Setup argument parser, the workhorse gluing it all together
    parser = ArgumentParser(
        epilog='(*) by default the script will look for the first device that '
               'matches %s* in %s, if multiple entries are found -s or -f must '
               'be used to specify which sensor to read.' %
               (W1_SENSOR_DEV_PREFIX, W1_SENSOR_DEV_DIR),
        description='Nagios check plugin for 1-wire temp. sensor on RaspberryPi'
    )
    parser.add_argument('-V', '--version',action="version",version=PROG_VERSION)

    pgroup = parser.add_mutually_exclusive_group(required=False)
    pgroup.add_argument('-C', '--celcius',  action='store_const',
                        dest='converter',   const=CONVERT_CELCIUS,
                        help='measure, critical and warn values in Celcius '
                             '(default)',   default=CONVERT_CELCIUS)
    pgroup.add_argument('-F', '--farenheit',action='store_const',
                        dest='converter',   const=CONVERT_FARENHEIT,
                        help='measure, critical and warn values in Farenheit')

    parser.add_argument('-w', '--warn',     type=float,
                        help='temperature for warning status')
    parser.add_argument('-c','--critical',  type=float,
                        help='temperature for critical status')

    pgroup = parser.add_mutually_exclusive_group(required=False)
    pgroup.add_argument('-q', '--quiet',    default=logging.CRITICAL,
                        action=SetLogLevel, const=LOGGING_NONE,
                        help='quiet (no output, only exit with exit code)')
    pgroup.add_argument('-v', '--verbose',  help='more verbose output',
                        action=SetLogLevel, const=logging.INFO)
    pgroup.add_argument('-d', '--debug',    help='debug output (more verbose)',
                        action=SetLogLevel, const=logging.DEBUG)

    parser.add_argument('-l', '--logfile',  action=SetLogFile,
                        help='send logging output to logfile')

    subparser = parser.add_subparsers(title='Supported temperature sensors')

    cpuparser = ArgumentParser(add_help=False)
    cpuparser.add_argument('-f', '--file', default=CPU_SENSOR_DEV,
                           help='input file (or device) to obtain data from'
                                ' (defaults to %s)' % CPU_SENSOR_DEV)
    cmdparser = subparser.add_parser('rpi_cpu', parents=[cpuparser],
                                    help='read built-in Raspberry Pi CPU temperature')
    cmdparser.set_defaults(func=read_rpi_cpu_temp, cmdparser=cmdparser, retries=0)

    w1parser = ArgumentParser(add_help=False)
    pgroup = w1parser.add_mutually_exclusive_group(required=False)
    pgroup.add_argument('-s', '--serial',
                        help='(unique part of) temperature sensor serial (*)')
    pgroup.add_argument('-f', '--file',
                        help='input file (or device) to obtain data from (*)')
    w1parser.add_argument('-r', '--retries', type=int,default=W1_SENSOR_READ_RETRIES,
                        help='number of times to retry reading sensor data when'
                             ' unstable (defaults to %d)' % W1_SENSOR_READ_RETRIES)
    cmdparser = subparser.add_parser('w1_ds18b20', parents=[w1parser],
                                    help='read 1-wire connected DS18b20 sensor')
    cmdparser.set_defaults(func=read_w1_temp,   cmdparser=cmdparser)

    # parse arguments and post-process command line options
    args = parser.parse_args()

    # if we got here all seems OK
    return args


def nagios_exit(status, message, data=None):
    """exit 'nagios-style', print status and message followed by perf. data"""
    if logger.isEnabledFor(logging.CRITICAL):
        if data is not None and len(data) > 0:
            perfdata = ' | ' + ' '.join([ "'%s'=%s" % (k,
                ';'.join(['' if x is None else str(x) for x in v])
                        if isinstance(v,list) else v) for k,v in data ])
        else:
            perfdata = ''
        print 'Temperature %s: %s%s' % (status[0], message, perfdata)
    exit(status[1])


if __name__ == '__main__':
    try:
        args = parse_args()
    except ArgumentParserError as e:
        nagios_exit(NAGIOS_UNKNOWN,'error with setup: ' + e.message)
    except (KeyboardInterrupt, EOFError) as e:
        print
        nagios_exit(NAGIOS_UNKNOWN,'initialization aborted')

    try:
        starttime = time()
        temperature, tries = args.func(args)
        endtime = time()

    except (KeyboardInterrupt) as e:
        nagios_exit(NAGIOS_UNKNOWN,'temperature sensor read aborted by user')

    except (IOError, ValueError) as e:
        nagios_exit(NAGIOS_UNKNOWN,'temperature sensor read failed: %s' % e)

    elapse = endtime-starttime
    logger.info('Got temperature reading of %.2f degrees %s in %fs',
                temperature, args.converter[2], elapse)

    unit = args.converter[1]
    message = 'current temperature is %.2f%s' % (temperature, unit)
    if args.critical is not None and temperature > args.critical:
        nagiosresult = NAGIOS_CRITICAL
        message+= ' and exceeds critical threshold %.2f%s' %(args.critical,unit)
    elif args.warn is not None and temperature > args.warn:
        nagiosresult = NAGIOS_WARN
        message+= ' and exceeds warning threshold %.2f%s' % (args.warn,unit)
    else:
        nagiosresult = NAGIOS_OK

    nagios_exit(nagiosresult, message, [
          ('temperature', [ '%f%s' % (temperature, unit),
                            args.warn, args.critical, None, None]),
          ('retries', [ tries-1, None, args.retries, 0, None ]),
          ('checktime', [ '%fs' % elapse, None, None, 0, None])
    ])