#!/bin/bash

COMMAND=$1

CHAIN=perfSONAR

# Determine BWCTL and OWAMP ports
BWCTLPARAMS=("iperf_port" "nuttcp_port" "owamp_port" "peer_port")
BWCTLDEFAULTS=("5001:5300" "5301:5600" "5601:5900" "6001:6200")
OWAMPPARAMS=("testports")
OWAMPDEFAULTS=("8760:9960")

BWCTLFILE="/etc/bwctl-server/bwctl-server.conf"
OWAMPFILE="/etc/owamp-server/owamp-server.conf"
FIREWALL_RULES="/etc/perfsonar/toolkit/perfsonar_firewall_settings.conf"
OLD_FIREWALL_RULES="/etc/perfsonar/toolkit/old_firewall_settings.conf"

if [ ! -f $FIREWALL_RULES ]; then
    echo "Couldn't find firewall rules"
    exit -1
fi

main() {
    if [ "${COMMAND}" == "install" ]; then   
        deleteOldFirewallRules
        echo "Adding perfSONAR firewall rules"
        addFirewallRules
        saveFirewallRules
    elif [ "${COMMAND}" == "uninstall" ]; then
        deleteFirewallRules
    fi
}

saveFirewallRules() {
    for iptables_cmd in "iptables" "ip6tables"; do
        service $iptables_cmd save
    done
}

deleteOldFirewallRules(){
    for iptables_cmd in "iptables" "ip6tables"; do
        readarray firewall_rules < ${OLD_FIREWALL_RULES}
        for rule in "${firewall_rules[@]}"; do
            case ${rule} in
                "-D INPUT"*)
                    $iptables_cmd $rule &> /dev/null
                    ;;
            esac
        done

        p=`echo ${BWCTLPARAMS[@]}`;
        deletePortRuleFromFile "$iptables_cmd" "$BWCTLFILE" "$p";
        p=`echo ${OWAMPPARAMS[@]}`;
        deletePortRuleFromFile "$iptables_cmd" "$OWAMPFILE" "$p";
    done
}

deleteFirewallRules(){
    for iptables_cmd in "iptables" "ip6tables"; do
        $iptables_cmd -F $CHAIN
        if [ $? != 0 ]; then
            echo "Problem removing $CHAIN firewall rules";
            exit -1;
        fi
    done
}

addFirewallRules(){
    for iptables_cmd in "iptables" "ip6tables"; do
        # Create the chain if it doesn't currently exit
        $iptables_cmd -n -L $CHAIN &> /dev/null
        if [ $? != 0 ]; then
            $iptables_cmd -N $CHAIN
            if [ $? != 0 ]; then
                echo "Problem creating firewall chain $CHAIN";
                exit -1;
            fi
        fi

        # Flush all existing rules
        $iptables_cmd -F $CHAIN
        if [ $? != 0 ]; then
            echo "Problem removing existing firewall rules from $CHAIN";
            exit -1;
        fi

        # The BWCTL and OWAMP rules come first for performance reasons
        p=`echo ${BWCTLPARAMS[@]}`;
        d=`echo ${BWCTLDEFAULTS[@]}`
        addPortRuleFromFile "$iptables_cmd" "$BWCTLFILE" "$p" "$d";
        p=`echo ${OWAMPPARAMS[@]}`;
        d=`echo ${OWAMPDEFAULTS[@]}`
        addPortRuleFromFile "$iptables_cmd" "$OWAMPFILE" "$p" "$d";
        
        #ICMP rules based on v6 or v4
        if [ "$iptables_cmd" = "ip6tables" ]; then 
            $iptables_cmd -A perfSONAR -p icmpv6 -j ACCEPT
        else
             $iptables_cmd -A perfSONAR -p icmp --icmp-type any -j ACCEPT
        fi
        
        
        # Then all the rules from the static file
        readarray firewall_rules < ${FIREWALL_RULES}
        for rule in "${firewall_rules[@]}"; do
            case ${rule} in
                "-A ${CHAIN}"*)
                    $iptables_cmd $rule &> /dev/null
                    ;;
                "-I ${CHAIN}"*)
                    $iptables_cmd $rule &> /dev/null
                    ;;
            esac
        done

        # Add a rule at the end to jump us back to the main process chain
        $iptables_cmd -A ${CHAIN} -j RETURN

        # Make sure that there's an INPUT rule to point us into the 
        $iptables_cmd -n -L INPUT | grep -q "${CHAIN}[ \t]"
        if [ $? != 0 ]; then
            $iptables_cmd -I INPUT -j ${CHAIN}
        fi
        
        $iptables_cmd -n -L INPUT | grep -q REJECT
        if [ $? != 0 ]; then 
            # Make sure we reject traffic
            $iptables_cmd -A INPUT -j REJECT
        fi
        
        $iptables_cmd -n -L OUTPUT | grep -q "all"
        if [ $? != 0 ]; then 
            # Allow outgoing from loopback
            $iptables_cmd -A OUTPUT -o lo -j ACCEPT
        fi
        
        $iptables_cmd -n -L FORWARD | grep -q REJECT
        if [ $? != 0 ]; then 
            # Make sure we reject FORWARD traffic
            $iptables_cmd -A FORWARD -j REJECT
        fi
    done
}

addPortRuleFromFile(){
	iptables_cmd=$1
	file=$2
	local params=( `echo "$3"` )
	local defaults=( `echo "$4"` )

	total=${#params[*]}
	for (( i=0; i<=$(( $total - 1 )); i++ ))
	do
		param_name="${params[$i]}"
		param_default="${defaults[$i]}"
		portrange=`awk -v search="^$param_name" '$0 !~ /^#/ && $0 ~ search {$1=""; print $0}' $file | sed s/-/:/g`

		if [ -z "$portrange" ]; then
			portrange=$param_default
			echo "Unable to find definition for \"$param_name\" in \"$file\" for \"$iptables_cmd\", using default of \"$portrange\""
		fi

		if [ ! -z "$portrange" ]; then
			$iptables_cmd -A $CHAIN -m udp -p udp --dport $portrange -j ACCEPT
			$iptables_cmd -A $CHAIN -m state --state NEW -m tcp -p tcp --dport $portrange -j ACCEPT
		fi
	done 
}

deletePortRuleFromFile(){
	iptables_cmd=$1
	file=$2
	local params=( `echo "$3"` )
	for i in "${params[@]}"
	do
		portrange=`awk -v search="^$i" '$0 !~ /^#/ && $0 ~ search {$1=""; print $0}' $file | sed s/-/:/g`
		if [ ! -z "$portrange" ]; then
			$iptables_cmd -D INPUT -m udp -p udp --dport $portrange -j ACCEPT &> /dev/null
			$iptables_cmd -D INPUT -m state --state NEW,ESTABLISHED -m tcp -p tcp --dport $portrange -j ACCEPT &> /dev/null
		fi
	done 
}

main

exit 0
