#!/bin/bash

# config
ports="80 9090 22 443 25 3306"
log="/var/log/traffic_mon.log"
chain="traffic_mon"
time_date=`date +'%Y%m%d %H:%M:%S'`

# check if iptables is installed
iptables=`which iptables`
if [ ! -x $iptables ]; then
	echo "No iptables installed!"
	exit 1
fi

init () {
	x=0
	# create chain
	$iptables -N $chain 
	for i in $ports; do
		$iptables -A INPUT -p tcp --dport $i -m state --state NEW,ESTABLISHED,RELATED -j $chain
		$iptables -A OUTPUT -p tcp --dport $i -m state --state NEW,ESTABLISHED,RELATED -j $chain
	done
	
	# create log file
	for i in $ports; do
		if [ $x -eq 0 ]; then
			echo -ne "\t\t     port $i" >> $log
		else
			echo -ne "\t   port $i" >> $log
		fi
		let x+=1
	done
	echo "" >> $log && x=0
	for i in $ports; do
		if [ $x -eq 0 ]; then
			echo -ne "\t\t    in    out\t" >> $log
		else
			echo -ne "   in  out\t" >> $log
		fi
		let x+=1
	done
	echo "" >> $log
}

check () {
	if [ ! -f $log ]; then
		echo "You need to run \"init\" before!"
		exit 1
	fi
}

clearcounter () {
	# clear the counter
	$iptables -Z INPUT
	$iptables -Z OUTPUT
}

collect () {
	# check if init was run before
	check
	echo -n "$time_date" >> $log
	for i in $ports; do
		in=`/sbin/iptables -L INPUT -v -n | /bin/egrep "${chain}.*${i}" | awk '{ print $2 }'`
		out=`/sbin/iptables -L OUTPUT -v -n | /bin/egrep "${chain}.*${i}" | awk '{ print $2 }'`
		echo -ne "   "$in"  "$out"\t" >> $log
	done
	echo "" >> $log
	clearcounter
}

clearrules () {
	check
	echo "Do you really want to clear all iptable rules"
	echo -n "and delete the log file? (y/n): "
	read yesno
	if [ "x$yesno" = "xy" ]; then
		$iptables -F
		$iptables -P INPUT ACCEPT
		$iptables -P OUTPUT ACCEPT
		$iptables -P FORWARD ACCEPT
		$iptables -X $chain
		rm -f $log
	else
		echo "Nothing cleared!"
	fi
}

print_help () {
	echo "Usage: $0 <option>"
	echo -e "\t   help: this help message"
	echo -e "\t   init: initialize the iptable rules"
	echo -e "\tcollect: collect the traffic"
	echo -e "\t  clear: clear the iptable rules and log file"
}


case "$1" in
	init)
		init
	;;
	collect)
		collect
	;;
	clear)
		clearrules
	;;
	help)
		print_help
	;;
	*)
		print_help
esac

exit 0