#!/usr/bin/env python
# -*- coding: utf-8 -*-
# kate: space-indent on; indent-width 4; replace-tabs on;

"""
 *  Copyright (C) 2011-2016, it-novum GmbH <community@openattic.org>
 *
 *  openATTIC is free software; you can redistribute it and/or modify it
 *  under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2.
 *
 *  This package 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.
"""

import os
import sys
from time import time, mktime
from datetime import datetime
from ConfigParser import ConfigParser
from configobj import ConfigObj

distro_config = [ '/etc/default/openattic', '/etc/sysconfig/openattic' ]
for filename in distro_config:
    if os.path.isfile(filename):
        config = ConfigObj(filename)
        break

sys.path.append(config['OADIR'])
from nagios.conf.distro import distro_settings

settings = distro_settings()

try:
    # Uncomment the following line to disable the device UUID check.
    #raise ImportError
    from pyudev import Context, Device
except ImportError:
    HAVE_UDEV = False
else:
    HAVE_UDEV = True


# Resolve the real device path, dereferencing symlinks as necessary.
disk = os.path.realpath(sys.argv[-1]).replace("/dev/", "")

# Read the state file (if possible).
state = ConfigParser()
diskstats = "{}/diskstats.{}".format(settings["NAGIOS_STATE_DIR"], disk)
havestate = bool( state.read(diskstats) ) and state.has_section("state")

exit = 0

with open("/sys/block/%s/stat" % disk, "rb") as fd:
    currstate = dict(zip((
        "rd_ios", "rd_merges", "rd_sectors", "rd_ticks",
        "wr_ios", "wr_merges", "wr_sectors", "wr_ticks",
        "ios_in_prog", "tot_ticks", "rq_ticks"
      ), [
        int(count) for count in fd.read().split()
    ]))
    currstate["timestamp"] = time()


# Sanity-Check the state file. We expect the Device UUID and creation timestamp to
# match those from the statfile (if possible).

if HAVE_UDEV:
    ctx = Context()
    dev = Device.from_name(ctx, "block", disk)

    try:
        createstamp = int(mktime((datetime.now() - dev.time_since_initialized).timetuple()))
    except AttributeError:
        createstamp = None
else:
    createstamp = None

if havestate:
    if not state.has_section("device"):
        havestate = False

    if createstamp is not None and (
       not state.has_option("device", "initialized") or createstamp != state.getint("device", "initialized")):
        havestate = False

    if HAVE_UDEV and "DM_UUID" in dev.keys() and (
       not state.has_option("device", "uuid") or dev.get("DM_UUID") != state.get("device", "uuid") ):
        havestate = False


# If something's off about the state file (non-existant, empty, doesn't match), populate it with basic data.
# (Note this can't be an else: statement because we need to re-check.)
if not havestate:
    if not state.has_section("device"):
        state.add_section("device")

    if createstamp is not None:
        state.set("device", "initialized", createstamp)

    if HAVE_UDEV and "DM_UUID" in dev.keys():
        state.set("device", "uuid", dev.get("DM_UUID"))


def wrapdiff(curr, last):
    """ Calculate the difference between last and curr.

        If last > curr, try to guess the boundary at which the value must have wrapped
        by trying the maximum values of 64, 32 and 16 bit signed and unsigned ints.
    """
    if last <= curr:
        return curr - last

    boundary = None
    for chkbound in (64,63,32,31,16,15):
        if last > 2**chkbound:
            break
        boundary = chkbound
    if boundary is None:
        raise ArithmeticError("Couldn't determine boundary")
    return 2**boundary - last + curr


if havestate:
    # <http://www.mjmwired.net/kernel/Documentation/block/stat.txt>
    # | The "sectors" in question are the standard UNIX 512-byte
    # | sectors, not any device- or filesystem-specific block size.
    bytes_per_sector = 512

    interval = float(currstate["timestamp"] - state.getfloat("state", "timestamp"))

    rd_ios = wrapdiff(currstate["rd_ios"], state.getfloat("state", "rd_ios"))
    wr_ios = wrapdiff(currstate["wr_ios"], state.getfloat("state", "wr_ios"))

    rd_ticks = wrapdiff(currstate["rd_ticks"], state.getfloat("state", "rd_ticks"))
    wr_ticks = wrapdiff(currstate["wr_ticks"], state.getfloat("state", "wr_ticks"))

    rd_sectors = wrapdiff(currstate["rd_sectors"], state.getfloat("state", "rd_sectors"))
    wr_sectors = wrapdiff(currstate["wr_sectors"], state.getfloat("state", "wr_sectors"))

    tot_ticks  = wrapdiff(currstate["tot_ticks"], state.getfloat("state", "tot_ticks"))

    rd_iops = rd_ios / interval
    wr_iops = wr_ios / interval

    rd_bps  = rd_sectors / interval * bytes_per_sector
    wr_bps  = wr_sectors / interval * bytes_per_sector

    tot_ios = rd_ios + wr_ios
    tot_iops = tot_ios / interval

    if tot_iops:
        servicetime = tot_ticks / tot_ios
    else:
        servicetime = 0

    if tot_ios:
        tot_avg_wait = (rd_ticks + wr_ticks) / tot_ios
    else:
        tot_avg_wait = 0

    if rd_ios:
        rd_avg_wait = rd_ticks / rd_ios
        rd_avg_size = rd_sectors * bytes_per_sector / rd_ios
    else:
        rd_avg_wait = 0
        rd_avg_size = 0

    # Normalized IOPS are calculated by taking the request size into account,
    # and basically say how many requests of size 4096 it would have taken
    # to accomplish what one bigger request has read/written.
    rd_normratio = rd_avg_size / 4096.
    rd_normiops  = rd_normratio * rd_iops

    if wr_ios:
        wr_avg_wait = wr_ticks / wr_ios
        wr_avg_size = wr_sectors * bytes_per_sector / wr_ios
    else:
        wr_avg_wait = 0
        wr_avg_size = 0

    wr_normratio = wr_avg_size / 4096.
    wr_normiops  = wr_normratio * wr_iops
    tot_normiops = rd_normiops + wr_normiops

    util_percent = tot_ticks / (interval * 1000.) * 100.

    if util_percent > 30:
        exit = 1
    if util_percent > 50:
        exit = 2

    out =  ("Disk load for %(disk)s is at %(util_percent).2f%%.|"
            "rd_iops=%(rd_iops).2f wr_iops=%(wr_iops).2f "
            "tot_iops=%(tot_iops).2f "
            "rd_normiops=%(rd_normiops).2f wr_normiops=%(wr_normiops).2f "
            "tot_normiops=%(tot_normiops).2f "
            "rd_normratio=%(rd_normratio).2f wr_normratio=%(wr_normratio).2f "
            "rd_bps=%(rd_bps).2fB/s wr_bps=%(wr_bps).2fB/s "
            "tot_avg_wait=%(tot_avg_wait)fs "
            "rd_avg_wait=%(rd_avg_wait)fs wr_avg_wait=%(wr_avg_wait)fs "
            "rd_avg_size=%(rd_avg_size).2fB wr_avg_size=%(wr_avg_size).2fB "
            "load_percent=%(util_percent).2f%%;30;50;0;100 ") % {
        "disk":         disk,
        "util_percent": util_percent,
        "rd_iops":      rd_iops,
        "wr_iops":      wr_iops,
        "tot_iops":     tot_iops,
        "rd_normiops":  rd_normiops,
        "wr_normiops":  wr_normiops,
        "tot_normiops": tot_normiops,
        "rd_normratio": rd_normratio,
        "wr_normratio": wr_normratio,
        "rd_bps":       rd_bps,
        "wr_bps":       wr_bps,
        "rd_avg_wait":  rd_avg_wait / 1000.,
        "wr_avg_wait":  wr_avg_wait / 1000.,
        "rd_avg_size":  rd_avg_size,
        "wr_avg_size":  wr_avg_size,
        "tot_avg_wait": tot_avg_wait / 1000.,
        }

    if tot_ios and tot_normiops and util_percent:
        ioindex   = int(tot_iops / util_percent)
        normindex = int(tot_normiops / util_percent)
        out += "ioindex=%d normindex=%d " % (ioindex, normindex)

    print out

else:
    print "Need state info, please wait until Nagios checks again."
    exit = 3
    if not state.has_section("state"):
        state.add_section("state")

# Copy all values from currstate to the statfile and save it.
for key in currstate:
    state.set("state", key, currstate[key])

state.write( open( diskstats, "wb" ) )

sys.exit(exit)
