#!/usr/bin/env python
# shorewall_accounting
# A munin plugin for tracking traffic as recorded by shorewall accounting rules
# Written by Chris AtLee <chris@atlee.ca>
#  Modified by Tanguy Pruvot for KMGT <tanguy.pruvot@gmail.com>
# Released under the GPL v2
import sys, commands, re
accountingLineExp = re.compile(r"^\s*\d+[KMG]*\s+(\d+)([KMGT]*)\s+(\w+).*$")

def getBytesByChain():
    trafficCmd = "shorewall"
    status, output = commands.getstatusoutput("/sbin/shorewall show accounting 2>/dev/null")
    if status != 0:
        raise OSError("Error running command (%s)[%i]: %s" % (trafficCmd, status, output))
    chains = {}
    for line in output.split("\n"):
        m = accountingLineExp.match(line)
        if m is not None:
            target = m.group(3)
            bytes = int(m.group(1))

	    if m.group(2) == "K":
	    	bytes *= 1024
            elif m.group(2) == "M":
		bytes *= 1024 * 1024
	    elif m.group(2) == "G":
		bytes *= 1024 * 1024 * 1024
	    elif m.group(2) == "T":
		bytes *= 1024 * 1024 * 1024 * 1024

            if target in chains:
                chains[target] += bytes
            else:
                chains[target] = bytes
#	else:
#	    print "IGNORED LINE %s" % ( line)

    retval = []
    chainNames = chains.keys()
    chainNames.sort()
    for name in chainNames:
        retval.append((name, chains[name]))
    return retval

if len(sys.argv) > 1:
    if sys.argv[1] == "autoconf":
        print "yes"
        sys.exit(0)
    elif sys.argv[1] == "config":
        print "graph_title Shorewall accounting"
        print "graph_category network"
        print "graph_vlabel bits per ${graph_period}"
        for chain,bytes in getBytesByChain():
            print "%s.min 0" % chain
            print "%s.type DERIVE" % chain
            print "%s.label %s" % (chain, chain)
            print "%s.cdef %s,8,*" % (chain, chain)
        sys.exit(0)

for chain, bytes in getBytesByChain():
    print "%s.value %i" % (chain, bytes)