#!/usr/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")
TWAMPPARAMS=("testports")
TWAMPDEFAULTS=("18760:19960")

BWCTLFILE="/etc/bwctl-server/bwctl-server.conf"
OWAMPFILE="/etc/owamp-server/owamp-server.conf"
TWAMPFILE="/etc/twamp-server/twamp-server.conf"
FIREWALL_RULES="/etc/perfsonar/toolkit/perfsonar_firewall_settings.conf"
FIREWALLD_RULES="/etc/perfsonar/toolkit/perfsonar_firewalld_settings.conf"

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

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

saveFirewallRules() {
    if type firewall-cmd &>/dev/null; then
        firewall-cmd --quiet --runtime-to-permanent
    else
        for iptables_cmd in "iptables" "ip6tables"; do
            service $iptables_cmd save
        done
    fi
}

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
}

addFirewalldRules(){
        # Load new XML service files
        firewall-cmd --quiet --reload &> /dev/null

        # 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";
        p=`echo ${TWAMPPARAMS[@]}`;
        d=`echo ${TWAMPDEFAULTS[@]}`
        addPortRuleFromFile "$iptables_cmd" "$TWAMPFILE" "$p" "$d";

        # Then all the rules from the static file
        readarray firewalld_rules < ${FIREWALLD_RULES}
        for rule in "${firewalld_rules[@]}"; do
            case ${rule} in
                "--add-"*)
                    firewall-cmd --quiet $rule &> /dev/null
                    ;;
                "--remove-"*)
                    firewall-cmd --quiet $rule &> /dev/null
                    ;;
            esac
        done
}

addFirewallRules(){
    if type firewall-cmd &>/dev/null; then
        addFirewalldRules
        return
    fi

    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";
        p=`echo ${TWAMPPARAMS[@]}`;
        d=`echo ${TWAMPDEFAULTS[@]}`
        addPortRuleFromFile "$iptables_cmd" "$TWAMPFILE" "$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"` )
    
    if [ -f "$file" ]; then
        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
                if type firewall-cmd &>/dev/null; then
                    portrange=${portrange/:/-}
                    portrange=${portrange// /}
                    firewall-cmd --quiet --add-port=$portrange/udp
                    firewall-cmd --quiet --add-port=$portrange/tcp
                else
                    $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
            fi
        done
    fi
}

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
}

#enable firewalld/iptables
if type systemctl &>/dev/null; then
    systemctl enable firewalld 2>/dev/null
    systemctl start firewalld 2>/dev/null
else
    /sbin/chkconfig iptables on 2>/dev/null
    /etc/init.d/iptables start 2>/dev/null
fi

#configure firewall rules
main

exit 0
