#!/bin/sh
F_RATELIMIT="/opt/trs/iptables_rcounter"
F_DROP="/opt/trs/iptables_dcounter"
TMP_CRU="/opt/trs/iptables_tmpcount"
umask 002

function lock {
    if [ "$#" -ne 1 ]; then
        echo 'usage: lock [LOCKFILENAME]' 1>&2
        return 1
    fi

    LOCKFILE="$1"

    # make a file with our PID
    # easy way to show who's waiting for a lock
    if ! touch "$LOCKFILE.$$" ; then
        echo "failed to create PID lockfile: $1 pid:$$ parent:$PPID" 1>&2
        return 1
    fi

    # try to symlink it
    while ! ln "$LOCKFILE.$$" "$LOCKFILE" 2>/dev/null
    do
        # if the symlink failed, wait for the current lock holder to exit
        echo "waiting for lock pid:$$ parent:$PPID"
        sleep 1
    done

    # symlink was created successfully, lock acquired

    # if the locking process exits with unlocking, delete our lock
    trap 'rm -f "$LOCKFILE" "$LOCKFILE.$$"' EXIT

    return 0
}

lock /opt/trs/mutex_firewall

if [[ "$2" == "filter4" ]]; then
    FILTER_OUT=`iptables -nxvL`
    RATELIMIT_COUNT=`iptables -nxvL RATELIMIT_CHAIN| awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
    DISCARD_COUNT=`iptables -nxvL DISCARD_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
elif [[ "$2" == "filter6" ]]; then
    FILTER6_OUT=`ip6tables -nxvL`
    RATELIMIT6_COUNT=`ip6tables -nxvL RATELIMIT_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
elif [[ "$2" == "mangle4" ]]; then
    MANGLE_OUT=`iptables -t mangle -nxvL`
elif [[ "$2" == "mangle6" ]]; then
    MANGLE6_OUT=`ip6tables -t mangle -nxvL`
else
    FILTER_OUT=`iptables -nxvL`
    FILTER6_OUT=`ip6tables -nxvL`
    MANGLE_OUT=`iptables -t mangle -nxvL`
    MANGLE6_OUT=`ip6tables -t mangle -nxvL`
    RATELIMIT_COUNT=`iptables -nxvL RATELIMIT_CHAIN| awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
    DISCARD_COUNT=`iptables -nxvL DISCARD_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
    RATELIMIT6_COUNT=`ip6tables -nxvL RATELIMIT_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
fi

IPTABLES_OUT=`echo -e "$FILTER_OUT\n$FILTER6_OUT\n$MANGLE_OUT\n$MANGLE6_OUT"`
MODE="BOTH"
if [[ "x$1" != "x" ]]; then
    MODE=$1
fi

if [[ -f $F_DROP ]]; then
     rm -rf $TMP_CRU
else

     if [[ -f $TMP_CRU ]]; then
        . $TMP_CRU
     fi
fi

if [[ "$MODE" == "RATELIMIT" ]] || [[ "$MODE" == "BOTH" ]]; then
    REJECT_COUNT=`echo "$IPTABLES_OUT"| awk -F ' ' '$3=="REJECT" {sum += $1} END {print sum}'`
    RATELIMIT_MANGLE=`iptables -t mangle -nxvL RATELIMIT_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
    if [[ -f $F_RATELIMIT ]]; then
        . $F_RATELIMIT
        rm -rf $F_RATELIMIT
    fi
    if [[ "x$REJECT_COUNT" == "x" ]];then
            REJECT_COUNT=0
    fi

    if [[ "x$RATELIMIT_COUNT" == "x" ]];then
            RATELIMIT_COUNT=0
    fi

    if [[ "x$RATELIMIT6_COUNT" == "x" ]];then
            RATELIMIT6_COUNT=0
    fi

    if [[ "x$RATELIMIT_MANGLE" == "x" ]];then
            RATELIMIT_MANGLE=0
    fi

    if [[ "x$RATE_COUNT" == "x" ]];then
            RATE_COUNT=0
    fi
    if [[ "x$CURRENT_RATE" == "x" ]];then
        CURRENT_RATE=0
    fi

    CURRENT_RATE1=$(($REJECT_COUNT+$RATELIMIT_COUNT+$RATELIMIT6_COUNT+$RATELIMIT_MANGLE))
    R_COUNT=$(($CURRENT_RATE1+$RATE_COUNT-$CURRENT_RATE))

    cat <<EOT >> $F_RATELIMIT
RATE_COUNT=$R_COUNT
EOT
echo "RATE_COUNT=$R_COUNT"
#echo "CURRENT_RATE=$CURRENT_RATE"
fi
if [[ "$1" == "DROP" ]] || [[ "$MODE" == "BOTH" ]]; then
    DISCARD6_COUNT=`ip6tables -nxvL DISCARD_CHAIN | awk -F ' ' '$3=="DROP" {sum += $1} END {print sum}'`
    FORWARD_COUNT=`echo "$IPTABLES_OUT" |awk -F ' ' '$1=="Chain" && $2=="FORWARD" && $4=="DROP" {sum += $5} END {print sum}'`
    INPUT_DROP=`echo "$IPTABLES_OUT" |awk -F ' ' '$1=="Chain" && $2=="INPUT" && $4=="DROP" {sum += $5} END {print sum}'`
    if [[ -f $F_DROP ]]; then
        . $F_DROP
        rm -rf $F_DROP
    fi
    if [[ "x$DISCARD_COUNT" == "x" ]];then
            DISCARD_COUNT=0
    fi

    if [[ "x$DISCARD6_COUNT" == "x" ]];then
            DISCARD6_COUNT=0
    fi
    if [[ "x$FORWARD_COUNT" == "x" ]];then
            FORWARD_COUNT=0
    fi

    if [[ "x$FORWARD6_COUNT" == "x" ]];then
            FORWARD6_COUNT=0
    fi

    if [[ "x$INPUT_DROP" == "x" ]];then
            INPUT_DROP=0
    fi

    if [[ "x$INPUT6_DROP" == "x" ]];then
            INPUT6_DROP=0
    fi

    if [[ "x$DROP_COUNT" == "x" ]];then
            DROP_COUNT=0
    fi
    if [[ "x$CURRENT_DROP" == "x" ]]; then
	  CURRENT_DROP=0
    fi

    CURRENT_DROP1=$(($DISCARD_COUNT+$DISCARD6_COUNT+$FORWARD_COUNT+$FORWARD6_COUNT+$INPUT_DROP+$INPUT6_DROP))
    D_COUNT=$(($CURRENT_DROP1+$DROP_COUNT-$CURRENT_DROP))

    cat <<EOT >> $F_DROP
DROP_COUNT=$D_COUNT
EOT
echo "DROP_COUNT=$D_COUNT"
#echo "CURRENT_DROP=$CURRENT_DROP"
fi

if [[ "$3" == "PM_FETCH" ]]; then
    rm -rf $TMP_CRU
	cat <<EOT >> $TMP_CRU
CURRENT_DROP=$CURRENT_DROP1
CURRENT_RATE=$CURRENT_RATE1
EOT
fi
