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

"""
 *  Copyright (C) 2011-2015, 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 sys
import dbus
from time import time
from os.path import exists, join
from optparse import OptionParser
from ConfigParser import ConfigParser

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()
havestate  = bool( savedstate.read("/var/lib/nagios3/protocol_traffic.%s" % options.interface.lower()) 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( "/var/lib/nagios3/protocol_traffic.%s" % options.interface, "wb" ) )

sys.exit(exit)
