# Copyright 2014 Jan Huwald <jh@sotun.de>
# 
# 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 3 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, see <http://www.gnu.org/licenses/>.

# abort on error semantic is required to produce a terse firewall
# description
set -e
set -o pipefail

# params: command [rulesets...]
#   ruleset are implemented as functions the call the various
#   callbacks defined below
function huwfwall() {
    local cmd="$1"
    shift
    # install callback function to wrap *table commands
    case "$cmd" in
	install) ;&
	check)
	    function target_callback() {
		if ! $*; then
		    echo -e "while executing\t $*";
		    # print backtrace
		    local i=2
		    while caller $i; do i=$((i+1)); done \
			| awk '{ print "in " $3 ":" $2 ":" $1 }'
		    return 1;
		fi
	    }
	    ;;
	print)
	    function target_callback() { echo $*; }
    esac
    # execute given subcommand
    case "$cmd" in
	install)
	    v=$(current_version)
	    while [ $# -gt 0 ]; do
		v=$((v + 1))
		setup_root_chains
		set +e
		(set -e; $1)
		local ret=$?
		set -e
		if [ $ret -eq 0 ] ; then break; fi
		echo "in ruleset $1\nTrying next ruleset"
		gc_rules
		shift
	    done
	    if [ $# -gt 0 ]; then
		commit_rules
		gc_rules
		return 0
	    else
		echo "All rulesets exhausted. Firewall setup gailed."
		gc_rules
		return 1
	    fi
	    ;;
	check)
	    while [ $# -gt 0 ]; do
		v=test$RANDOM
		setup_root_chains
		set +e
		(set -e; $1)
		local ret=$?
		set -e
		if [ $ret -eq 0 ]; then
		    echo "Ruleset '$1' is ok"
		else
		    echo "Ruleset '$1' failed"
		fi
		gc_rules
		shift
	    done
	    ;;
	print)
	    function target_callback() { echo $*; }
	    v="\${v}"
	    while [ $# -gt 0 ]; do
		echo "# ruleset $1 will install"
		echo "# builtin:"
		setup_root_chains
		echo "# user-defined:"
		$1
		[ $# -eq 1 ] || echo
		shift
	    done
	    ;;
	gc)
	    gc_rules
	    ;;
	current-version)
	    current_version
	    ;;
	*)
	    echo -e \
"Usage: $0 <install|check|print|gc>
\tinstall\tSetup the firewall
\tcheck\tCheck rules by installing, but not using them
\tprint\tShow the iptable commands to be executed
\tgc\tGarbage collect unused rules"
    esac
}

function rule() {
    local targets="$1"
    local chain="$2"
    shift 2
    iterate_targets "$targets" -A ${v}_"$chain" $*
}

function chain() {
    local targets="$1"
    local chain="$2"
    iterate_targets "$targets" -N ${v}_"$chain"
}

function iterate_targets() {
    local targets="$1"
    shift
    local target
    for i in $(echo "$targets" | sed 's/\(.\)/\1 /g'); do
	case $i in
	    4) target=iptables;;
	    6) target=ip6tables;;
	    e) target=ebtables;;
	    *) echo "invalid target '$target'"; return 1
	esac
	target_callback $target $*
    done
}

function setup_root_chains() {
    iterate_root_chains setup_root_chain_callback;
}
function setup_root_chain_callback() {
    local target="$1"; local table="$2"; local chain="$3"
    target_callback $target -t $table -N ${v}_${chain}
}

function commit_rules() {
    # first insert all heads to minimize incoherent rule time; then
    # remove all previous rules in the chain
    iterate_root_chains commit_rules_insert_head
    iterate_root_chains commit_rules_remove_tail
}
function commit_rules_insert_head() {
    local target="$1"; local table="$2"; local chain="$3"
    local chain_upper=$(echo "$chain" | tr a-z A-Z)
    target_callback $target -t $table -I $chain_upper -j ${v}_$chain
}
function commit_rules_remove_tail() {
    # call targets directly, not via target_callback: we need the
    # return status to determine loop termination
    local target="$1"; local table="$2"; local chain="$3"
    local chain_upper=$(echo "$chain" | tr a-z A-Z)
    while $target -t $table -D $chain_upper 2 2>/dev/null; do :; done
}

function current_version() {
    (echo 1; iterate_root_chains current_version_callback) | sort -n | tail -n1
}
function current_version_callback() {
    local target="$1"; local table="$2"; local chain="$3"
    local chain_upper=$(echo "$chain" | tr a-z A-Z)
    local opt
    if grep -q ip <<<$target; then opt="-n"; fi
    "$target" -t "$table" $opt -L "$chain_upper" \
	| (egrep "^[0-9]+_$chain" || true) \
	| cut -f1 -d_ | head -n1
}

function iterate_root_chains() {
    local callback="$1"; local i
    for i in {iptables,ip6tables}/filter/{input,forward,output} \
	     iptables/nat/{prerouting,input,output,postrouting} \
	     ebtables/filter/{input,forward,output}; do
	local target table chain
	read target table chain <<<$(echo $i | tr / ' ')
	$callback $target $table $chain
    done
}

function gc_rules() {
    local i
    for i in {iptables,ip6tables}/filter iptables/nat; do
	local target table
	read target table <<<$(echo $i | tr / ' ')
	local succ=true
	while $succ; do
	    succ=false
	    for i in $($target -t $table -L \
		| grep ^Chain | grep '(0 references)' \
		| cut -f2 -d' '); do
		succ=true
		$target  -t $table -F $i
		$target  -t $table -X $i
	    done
	done
    done
    ebtables -X
}