#!/bin/bash
# Ports Bandwidth, this is a Munin plugin by _KaszpiR_ kaszpir@gmail.com
# code under GPL
#
# USE AT YOUR OWN RISK
#
# ver 0.2 2009.09.28
# - fixed issue with resetting stats = makes wrong data displayed
# - change to derive
# - initial release
#
# Known issues:
# - complicated :)
# - plugin was written for munin in up to 60min cycle

CHAIN="TRAFFIC" # name of the chain to create
IPT=`which iptables`
AWK=`which awk`
IFACE="eth0" #interface  to bind, currenlty only one supported
PORTS_RAW="64738 64739 27015 27016 27017 27018 27019 27020 26900 26901 26902 27005 27006" # define numeric ports here


if [ $EUID -ne 0 ]; then
    echo "You must be root to run this script."
    exit 1
fi

if [ "$IPT" == ""  ]; then
    echo "ERROR: iptables or not found."
    exit 2
fi
if [ "$AWK" == ""  ]; then
    echo "ERROR: awk or not found."
    exit 2
fi

PORT=`echo $0 | awk -F_ '{print $2}'`
if [ "$PORT" == "" ]; then
    echo "Failed to parse port from script name, aborting"
#    exit 3
fi
OUT=/tmp/$CHAIN.$PORT.log #output file used to save data, use this as the output file when running scripts like munin
CHECK="/tmp/munin.chains.txt"

######################################################################
function i_add()
{
if [ -e $CHECK ]; then
echo "ERROR: You have already added counters to the $IPT, Run $0 del and then try again."
exit 1;
fi;

    $IPT -N $CHAIN

    #on standalone server we join to INPUT AND OUTPUT, on router you can add FORWARD
    $IPT -I INPUT -j $CHAIN
    $IPT -I OUTPUT -j $CHAIN

    #where are we adding
    $IPT -A $CHAIN -o $IFACE # upstream
    $IPT -A $CHAIN -i $IFACE # downstream

    for port in $PORTS_RAW; do
        $IPT -A $CHAIN -p tcp --sport $port
        $IPT -A $CHAIN -p udp --sport $port

        $IPT -A $CHAIN -p tcp --dport $port
        $IPT -A $CHAIN -p udp --dport $port
    done
    touch $CHECK

}
######################################################################
function i_del()
{
    $IPT -F $CHAIN
    $IPT -D INPUT -j $CHAIN
    $IPT -D OUTPUT -j $CHAIN
    $IPT -X $CHAIN
    rm -f $CHECK
}
######################################################################
function i_reset()
{
    $IPT -Z $CHAIN
}
######################################################################
function i_save()
{
    $IPT -L $CHAIN -nvx > $OUT
}

######################################################################
function i_show()
{
    $IPT -L $CHAIN -nvx
}

######################################################################
function i_munin()
{
#output useful for munin, jsut grep by src tcp/udp port
##  136993 15798233            udp  --  *      *       0.0.0.0/0            0.0.0.0/0           udp dpt:64738
    $IPT -L $CHAIN -nvx | grep -v all|tr ":" "_"|awk '$1 ~ /^[0-9]+$/ {printf "%s%s.value %d\n",$9, $10, $2 }'|grep $PORT

}
######################################################################
function i_config()
{
echo "graph_title $PORT Bandwidth"
echo "graph_args --base 1000 -l 0"
echo "graph_vlabel bytes"
echo "graph_category network"
echo "graph_info Displays bandwidth usage per port parsed from script name."

echo "tcpdpt_$PORT.type DERIVE"
echo "udpdpt_$PORT.type DERIVE"
echo "tcpspt_$PORT.type DERIVE"
echo "udpspt_$PORT.type DERIVE"

echo "tcpdpt_$PORT.label TCP DST $PORT"
echo "udpdpt_$PORT.label UDP DST $PORT"
echo "tcpspt_$PORT.label TCP SRC $PORT"
echo "udpspt_$PORT.label UDP SRC $PORT"

}

######################################################################

case "$1" in
"add")
    i_add;
    i_show;
    ;;

"reset")
    i_save;
    i_reset;
    ;;

"del")
    i_save;
    i_del;
    ;;

"update")
    i_save;
    i_del;
    i_add;
    ;;

"show")
    i_save;
    i_show;
    ;;

"autoc*")
    echo "No. Edit ${0}, set up ports ad run '${0} symlinks'";
    ;;
    
"config")
    i_config;
    ;;

"symlinks")
    for port in $PORTS_RAW; do
	echo "ln -s ${0} /etc/munin/plugins/${0}${port}"
    done;
    ;;

*|"munin")
    i_save;
    i_munin;
    ;;
esac