#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
"""check_iostat - plugin shows the I/O usage of a specified disk"""

import argparse
import sys
import os
import subprocess

__author__ = "John van Zantvoort"
__copyright__ = "John van Zantvoort"
__email__ = "John van Zantvoort"
__license__ = "MIT"
__version__ = "1.0.1"

RETURN_CODE = {'OK': 0, 'WARNING': 1, 'CRITICAL': 2, 'UNKNOWN': 3, 'DEPENDENT': 4}

PATH = [os.path.expanduser(x) for x in os.getenv('PATH').split(':')]
PATH = [os.path.abspath(x) for x in PATH]

def level_args(string):
    cols = string.split(',')
    if len(cols) != 4:
        raise argparse.ArgumentTypeError("incorrect number of arguments")

    rvlist = list()
    for x in cols:
        try:
            x = float(x)
        except ValueError:
            raise argparse.ArgumentTypeError("incorrect type of argument")
        rvlist.append(x)

    return dict(zip(['tps', 'kbread', 'kbwritten', 'await'], rvlist))

def disk_arg(string):
    if not os.path.exists("/dev/%s" % string):
        raise argparse.ArgumentTypeError("Device incorrectly specified")
    return string

def which(command):
    for dirname in PATH:
        fpath = os.path.join(dirname, command)
        if os.path.exists(fpath):
            return fpath
    return command

def parse_args(args):
    """
    Parse command line parameters

    :param args: command line parameters as list of strings
    :return: command line parameters as :obj:`argparse.Namespace`
    """
    retv = dict()
    description  = "This plugin shows the I/O usage of the specified disk,"
    description += "using the iostat external program. It prints three"
    description += "statistics: Transactions per second (tps), Kilobytes per "
    description += "second read from the disk (KB_read/s) and and written to "
    description += "the disk (KB_written/s)"

    usage='''%(prog)s -d disk -c <tps>,<read>,<wrtn>,<await> -w <tps>,<read>,<wrtn>,<await>'''

    epilog = "available disks: %s" % ", ".join(lsblk())

    parser = argparse.ArgumentParser(description=description, usage=usage, epilog=epilog)

    parser.add_argument(
        '--version',
        action='version',
        version='check_iostat {ver}'.format(ver=__version__))

    parser.add_argument(
            '-d', '--disk',
            dest="DISK",
            type=disk_arg,
            required=True,
            )

    parser.add_argument(
            '-c', '--critical',
            dest="CRITICAL",
            required=True,
            type=level_args,
            help="Sets the CRITICAL level for tps, KB_read/s, KB_written/s and average wait, respectively")

    parser.add_argument(
            '-w', '--warning',
            dest="WARNING",
            required=True,
            type=level_args,
            help="Sets the WARNING level for tps, KB_read/s,  KB_written/s and average wait, respectively")

    opts = vars(parser.parse_args(args))

    for keyn, keyv in opts.items():
        if keyv is not None:
            retv[keyn] = keyv


    return retv

def _iostat(disk, *args):
    command = [which('iostat')]
    command.append(disk)
    command.extend(args)
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdoutdata, stderrdata) = process.communicate()
    if process.returncode != 0:
        sys.exit("command failed")
    return stdoutdata, stderrdata

def _lsblk(*args):
    command = [which('lsblk')]
    command.extend(args)
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    (stdoutdata, stderrdata) = process.communicate()
    if process.returncode != 0:
        sys.exit("command failed")
    return stdoutdata, stderrdata

def lsblk():
    retv = list()
    (stdoutdata, stderrdata) = _lsblk('--output', 'NAME,TYPE', '--ascii', '--raw')
    for line in stdoutdata.split('\n'):
        cols = line.strip('\n').strip().split()
        if len(cols) >= 2:
            if cols[1] == 'disk':
                retv.append(cols[0])
    return retv


def iostat(disk):
    retv = dict()
    legend_kbdisk = ['Device', 'tps', 'kB_read/s', 'kB_wrtn/s', 'kB_read',
                     'kB_wrtn']
    legend_kbdiskes = ['Device', 'rrqm/s', 'wrqm/s', 'r/s', 'w/s', 'rkB/s',
                       'wkB/s', 'avgrq-sz', 'avgqu-sz', 'await', 'r_await',
                       'w_await', 'svctm', '%util']

    (stdoutdata, stderrdata) = _iostat(disk, '-k', '2', '2')
    inline = str()
    for line in stdoutdata.split('\n'):
        if line.strip().startswith(disk):
            inline = line
    kbdisk = dict(zip(legend_kbdisk, inline.split()))

    (stdoutdata, stderrdata) = _iostat(disk, '-x', '-k', '2', '2')
    inline = str()
    for line in stdoutdata.split('\n'):
        if line.strip().startswith(disk):
            inline = line
    kbdiskes = dict(zip(legend_kbdiskes, inline.split()))
    retv['tps'] = kbdisk['tps']
    retv['kbread'] = kbdisk['kB_read/s']
    retv['kbwritten'] = kbdisk['kB_wrtn/s']
    retv['await'] = kbdiskes['await']
    return retv


def check(indata, refdata):
    legend = ['tps', 'kbread', 'kbwritten', 'await']
    message = str()
    is_ok = True
    for item in legend:
        inx = float(indata.get(item))
        refx = float(refdata.get(item))
        if inx >= refx:
            is_ok = False
            message += "%s: %s " % (item, inx)
    return is_ok, message
    

def is_critical(data, args):
    refdata = args.get('CRITICAL')
    status, msg = check(data, refdata)

    if status:
        return False, msg
    else:
        return True, msg

def is_warning(data, args):
    refdata = args.get('WARNING')
    status, msg = check(data, refdata)

    if status:
        return False, msg
    else:
        return True, msg

def main(args):
    args = parse_args(args)
    data = iostat(args['DISK'])
    exit_code = RETURN_CODE['OK']
    status = "OK"

    crit_status, msg = is_critical(data, args)
    warn_status, msg = is_warning(data, args)

    message = " - I/O stats "
    message += "tps=%(tps)s KB_read/s=%(kbread)s " % data
    message += "KB_written/s=%(kbwritten)s Average_wait/s=%(await)s | " % data
    message += "'tps'=%(tps)s; 'KB_read/s'=%(kbread)s; " % data
    message += "'KB_written/s'=%(kbwritten)s; 'await'=%(await)s;" % data

    if crit_status:
        status = 'CRITICAL'
        exit_code = RETURN_CODE['CRITICAL']

    elif warn_status:
        status = 'WARNING'
        exit_code = RETURN_CODE['WARNING']

    print(status + message)
    sys.exit(exit_code)

if __name__ == "__main__":
    main(sys.argv[1:])
