#!/bin/bash
# Record all active iptables rule counters into RRD datasets.
# Optionally, graph the result. Intended for use with cron.
#
# requires bash 4.0+ regexes, substring manipulation

# make sure we're running flocked against the PID file
if ! grep -q /usr/bin/lckdo /proc/$PPID/cmdline; then
	exec /usr/bin/lckdo /run/iptables-rrdtool.pid "$0" "$@";
fi

[ -s /etc/default/iptables-rrdtool ] && . /etc/default/iptables-rrdtool
RRDDIR="${RRDDIR:-/var/lib/iptables-rrdtool}"
PNGDIR="${PNGDIR:-$RRDDIR}"

set -e
#set -x

# "Internal Field Separator" for composing arguments from command substitution, among other things
IFS=$'\n'

update () {

  name="${1}"
  counter="${2}"

  if ! [ -s "${RRDDIR}/${name}.rrd" ]; then
    rrdtool create "${RRDDIR}/${name}.rrd" \
    DS:rule:DERIVE:600:0:U \
    RRA:AVERAGE:0.5:1:576  \
    RRA:AVERAGE:0.5:6:720  \
    RRA:AVERAGE:0.5:24:720 \
    RRA:AVERAGE:0.5:288:730
  fi

  rrdtool update "${RRDDIR}/${name}.rrd" "N:${counter}"

  if [ "${GENPNG}" = "1" ]; then
    for interval in d w m y; do
      rrdtool graph "${PNGDIR}/${name}:${interval}.png" --full-size-mode -w 640 -h 480 --units=si --logarithmic --start -1${interval} DEF:rule="${RRDDIR}/${name}.rrd":rule:AVERAGE LINE1:rule#0080ff:rule VDEF:total=rule,TOTAL GPRINT:total:Total\\\:%8.3lf\ %s
    done
  fi
}

for table in $(cat /proc/net/ip_tables_names); do
  for rule in $(/sbin/iptables-save -c -t ${table}); do
# iptables-save has inconsistent output for no sane reason; we grab rematch
# strings both before and after the counters, then hamfistedly combine them
# under the generally safe assumption that one is blank
    if [[ "${rule}" =~ ^(.*)\[[0-9]+:([0-9]+)\](.*)$ ]]; then
      name="${BASH_REMATCH[1]}${BASH_REMATCH[3]}"
      name="${name//\/32/}"    # remove extraneous POSIX-unsafe '/32' string
      name="${name//\//slash}" # replace all other instances of POSIX-unsafe '/'
      name="${name/#:/ -P }"   # normalize iptables-save's dumb policy rule output
      name="${name//:/colon}"  # remove colon because rrdtool can't deal with POSIX filenames
      name="${name% }"         # remove trailing whitespace
      name="iptables -t ${table}${name}"
      counter="${BASH_REMATCH[2]}"

      update "${name}" "${counter}"
    fi
  done
done

for interface in $(ls /sys/class/net/); do
  for gress in tx rx; do
    read counter < /sys/class/net/"${interface}"/statistics/"${gress}"_bytes
    # what the hell are allowed characters for interface names, anyway?
    interface="${interface//\//slash}" # replace instances of POSIX-unsafe '/'
    interface="${interface//:/colon}"  # remove colon because rrdtool can't deal with POSIX fileinterfaces
    name="${interface} ${gress}"
    update "${name}" "${counter}"
  done
done