#!/bin/sh
#
# Traffic logging tool for DD-WRT-based routers using InfluxDB
#
# Based on work from Emmanuel Brucy (e.brucy AT qut.edu.au)
# Based on work from Fredrik Erlandsson (erlis AT linux.nu)
# Based on traff_graph script by twist - http://wiki.openwrt.org/RrdTrafficWatch
# 
# This program 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; either version 2
# of the License, or (at your option) any later version.
# 
# This program 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.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

DBURL=http://192.168.100.100:8086
DBNAME=traffic
LAN_IFNAME=$(nvram get lan_ifname)
WAN_IFNAME=$(nvram get wan_ifname)

lock()
{
	while [ -f /tmp/trafficmon.lock ]; do
		if [ ! -d /proc/$(cat /tmp/trafficmon.lock) ]; then
			echo "WARNING : Lockfile detected but process $(cat /tmp/trafficmon.lock) does not exist !"
			rm -f /tmp/trafficmon.lock
		fi
		sleep 1
	done
	echo $$ > /tmp/trafficmon.lock
}

unlock()
{
	rm -f /tmp/trafficmon.lock
}

case ${1} in

"setup" )

	# Create the RRDIPT2 CHAIN (it doesn't matter if it already exists).
	# This one is for the whole LAN -> WAN and WAN -> LAN traffic measurement
	iptables -N RRDIPT2 2> /dev/null
	
	# Add the RRDIPT2 CHAIN to the FORWARD chain (if non existing).
	iptables -L FORWARD -n | grep RRDIPT2 > /dev/null
	if [ $? -ne 0 ]; then
		iptables -L FORWARD -n | grep "RRDIPT2" > /dev/null
		if [ $? -eq 0 ]; then
			echo "DEBUG : iptables chain misplaced, recreating it..."
			iptables -D FORWARD -j RRDIPT2
		fi
		iptables -I FORWARD -j RRDIPT2
	fi
	 
	# Add the LAN->WAN and WAN->LAN rules to the RRDIPT2 chain
	iptables -nvL RRDIPT2 | grep ${WAN_IFNAME}.*${LAN_IFNAME} >/dev/null
	if [ $? -ne 0 ]; then
		iptables -I RRDIPT2 -i ${WAN_IFNAME} -o ${LAN_IFNAME} -j RETURN
	fi

	iptables -nvL RRDIPT2 | grep ${LAN_IFNAME}.*${WAN_IFNAME} >/dev/null
	if [ $? -ne 0 ]; then
		iptables -I RRDIPT2 -i ${LAN_IFNAME} -o ${WAN_IFNAME} -j RETURN
	fi

	# Create the RRDIPT CHAIN (it doesn't matter if it already exists).
	iptables -N RRDIPT 2> /dev/null

	# Add the RRDIPT CHAIN to the FORWARD chain (if non existing).
	iptables -L FORWARD -n | grep RRDIPT[^2] > /dev/null
	if [ $? -ne 0 ]; then
		iptables -L FORWARD -n | grep "RRDIPT" > /dev/null
		if [ $? -eq 0 ]; then
			echo "DEBUG : iptables chain misplaced, recreating it..."
			iptables -D FORWARD -j RRDIPT
		fi
		iptables -I FORWARD -j RRDIPT
	fi

	# For each host in the ARP table
	grep ${LAN_IFNAME} /proc/net/arp | while read IP TYPE FLAGS MAC MASK IFACE
	do
		# Add iptable rules (if non existing).
		iptables -nL RRDIPT | grep "${IP} " > /dev/null
		if [ $? -ne 0 ]; then
			iptables -I RRDIPT -d ${IP} -j RETURN
			iptables -I RRDIPT -s ${IP} -j RETURN
		fi
	done	
	;;
	
"update" )
	lock

	# Read and reset counters
	iptables -L RRDIPT -vnxZ -t filter > /tmp/traffic_$$.tmp
	iptables -L RRDIPT2 -vnxZ -t filter > /tmp/global_$$.tmp
	CURDATE=`date +%s`

	grep -Ev "0x0|IP" /proc/net/arp  | while read IP TYPE FLAGS MAC MASK IFACE
	do
		IN=0
		OUT=0
		# Add new data to the graph. 
          	grep ${IP} /tmp/traffic_$$.tmp | while read PKTS BYTES TARGET PROT OPT IFIN IFOUT SRC DST
          	do
			if [ "${DST}" = "${IP}" ]; then
				IN=${BYTES}
			fi
                        
			if [ "${SRC}" = "${IP}" ]; then
				OUT=${BYTES}
			fi
			curl -is -XPOST "$DBURL/write?db=$DBNAME" --data-binary "hosts,mac=$MAC inBytes=${IN},outBytes=${OUT} ${CURDATE}000000000" >/dev/null 2>&1
		done
	done
	
	# Chain RRDIPT2 Processing
	IN=0
	OUT=0
	grep ${LAN_IFNAME} /tmp/global_$$.tmp | while read PKTS BYTES TARGET PROT OPT IFIN IFOUT SRC DST
	do
		if [ "${IFIN}" = "${LAN_IFNAME}" ]; then
			IN=${BYTES}
		fi
                        
		if [ "${IFIN}" = "${WAN_IFNAME}" ]; then
			OUT=${BYTES}
		fi
		curl -is -XPOST "$DBURL/write?db=$DBNAME" --data-binary "hosts,mac=00:00:00:00:00:00 inBytes=${IN},outBytes=${OUT} ${CURDATE}000000000" >/dev/null 2>&1
	done
	
	# Free some memory
	rm -f /tmp/*_$$.tmp
	unlock
	;;

*)
	echo "Usage : $0 {setup|update}"
	echo "Options : "
	echo "   $0 setup"
	echo "   $0 update"
	exit
	;;
esac