#!/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, sys
import dbus
from time import time
from os.path import exists, join
from optparse import OptionParser
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()

parser = OptionParser()

parser.add_option( "-d", "--dbus",
    help="DBus service to connect to.",
    default="org.openattic.systemd"
    )

parser.add_option( "-i", "--interface",
    help="The interface to query.", default=''
    )

options, progargs = parser.parse_args()

protocols = {
    'iSCSI':   ('tcp', 3260),
    'SMB':     ('tcp', 445),
    'SSH':     ('tcp', 22),
    'HTTP':    ('tcp', 80),
    'HTTPS':   ('tcp', 443),
    'FTPCTRL': ('tcp', 21),
    'FTPDATA': ('tcp', 20),
    'NFSTCP':  ('tcp', 2049),
    'NFSUDP':  ('udp', 2049),
    }

savedstate = ConfigParser()
protocol_traffic = "{}/protocol_traffic.{}".format(settings["NAGIOS_STATE_DIR"], options.interface.lower())
havestate  = bool( savedstate.read(protocol_traffic) and savedstate.has_section("state") )

if not exists( join( "/sys/class/net", options.interface, "statistics" ) ):
    print "Interface %s does not exist!" % options.interface
    sys.exit(2)

kernelstats = {
    "in": {
        "bytes": int( open( join( "/sys/class/net", options.interface, "statistics", "rx_bytes"   ), "r" ).read() ),
        "pkgs":  int( open( join( "/sys/class/net", options.interface, "statistics", "rx_packets" ), "r" ).read() ),
        },
    "out": {
        "bytes": int( open( join( "/sys/class/net", options.interface, "statistics", "tx_bytes"   ), "r" ).read() ),
        "pkgs":  int( open( join( "/sys/class/net", options.interface, "statistics", "tx_packets" ), "r" ).read() ),
        },
    }

try:
    nagios = dbus.SystemBus().get_object(options.dbus, "/nagios")
    stats = nagios.iptables_get_stats()
except dbus.exceptions.DBusException:
    print "Could not query Systemd"
    sys.exit(3)

exit = 0

ifspeed = dbus.SystemBus().get_object(options.dbus, "/ifconfig").get_speed(options.interface)


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

data = {}
for rule in stats:
    tags = rule["comment"].split(':')
    if not tags or tags[0] != "OPENATTIC":
        continue

    obj = data
    for tag in tags[1:-1]:
        if tag not in obj:
            obj[tag] = {}
        obj = obj[tag]
    obj[tags[-1]] = rule

havedata = ( options.interface.upper() in data and data[ options.interface.upper() ] )

if not havedata or not havestate:
    if not savedstate.has_section("state"):
        savedstate.add_section("state")

    # add rules
    for protocol, (socketproto, portno) in protocols.iteritems():
        if not havedata:
            nagios.iptables_install_rules( options.interface, socketproto, portno, protocol )
            for direction in ("in", "out"):
                sectname = protocol + ':' + direction
                if not savedstate.has_section(sectname):
                    savedstate.add_section(sectname)
                for field in ("bytes", "pkgs"):
                    savedstate.set(sectname, field, 0)

        elif not havestate:
            for direction in ("in", "out"):
                sectname = protocol + ':' + direction
                savedstate.add_section(sectname)
                for key, val in data[ options.interface.upper() ][ protocol.upper() ][ direction.upper() ].iteritems():
                    savedstate.set(sectname, key, val)

    print "Checks for %s have been initialized, please wait until Nagios checks again." % options.interface
    exit = 3

else:
    ifacedata = data[ options.interface.upper() ]
    perfdata  = {}

    dt = time() - savedstate.getfloat("state", "timestamp")

    sums = {
        "in":  { "bytes": 0, "pkgs": 0 },
        "out": { "bytes": 0, "pkgs": 0 },
        }

    for protocol, (socketproto, portno) in protocols.iteritems():
        for direction in ("in", "out"):
            sectname = protocol + ':' + direction
            mydata   = ifacedata[ protocol.upper() ][ direction.upper() ]

            for field in ("bytes", "pkgs"):
                perfkey = "_".join( (protocol.lower(), direction, field) )

                dvalue  = wrapdiff(int(mydata[field]), savedstate.getint(sectname, field))
                sums[direction][field] += dvalue
                perfdata[perfkey] = dvalue / dt

            for key, val in mydata.iteritems():
                savedstate.set(sectname, key, val)

    # Kernel stats
    for direction in ("in", "out"):
        for field in ("bytes", "pkgs"):
            perfkey = "_".join( ("kernel", direction, field) )
            dvalue  = wrapdiff(int(kernelstats[direction][field]), savedstate.getint("state", direction+'_'+field))
            perfdata[perfkey] = dvalue / dt

            perfkey = "_".join( ("other", direction, field) )
            perfdata[perfkey] = (dvalue - sums[direction][field]) / dt

    # make sure the order of performance values in the output does never change.
    perfkeys = perfdata.keys()
    perfkeys.sort()

    out = "Traffic on %s" % options.interface

    if ifspeed != -1:
        # ifspeed = 1000 means 1000 MegaBit. Calc to bytes.
        bytespeed = ifspeed * 1000000 / 8.
        perc = float(perfdata["kernel_in_bytes"] + perfdata["kernel_out_bytes"]) / float(bytespeed) * 100
        out += " is at %d%% (max %d MBit/s)" % (perc, ifspeed)

        if perc > 50:
            exit = 1
        elif perc > 80:
            exit = 2

    print "%s|%s" % (
        out, ' '.join([ "%s=%.2f" % (key, perfdata[key]) for key in perfkeys ])
        )

savedstate.set("state", "timestamp", time())

for direction in ("in", "out"):
    for field in ("bytes", "pkgs"):
        savedstate.set("state", direction + '_' + field, kernelstats[direction][field])

savedstate.write( open( protocol_traffic, "wb" ) )

sys.exit(exit)
