#!/bin/sh

# Debugging function
debug() {
    if [ $DEBUG -eq 1 ]; then
        echo $(date)" DEBUG: "$1 $2 $3
    else
        echo $(date)" DEBUG: "$1 $2 $3 >>/var/log/iptables.log
    fi
}

# Task type variables
DEBUG=0
ROUTE=$ROUTE
HOST=$HOST
PREROUTING=$PREROUTING
POSTROUTING=$POSTROUTING

# Mandatory task variables
CHAIN=$CHAIN
NAME=$NAME
COMMENT="$COMMENT"
NAME="$NAME-$COMMENT"
PROTOCOL=$TYPE
OPERATION=$OPERATION

EXTRA_OPTIONS="$2 $3 $4"

set | grep SOURCE
set | grep TARGET
set | grep ROLES

SERVICE_FILES=$SERVICE_FILES
HOST_FILE=$HOST_FILE
if [ "$HOST_FILE" == "" ]; then
    HOST_FILE="/etc/dns/hosts.local"
fi

RETRIES_NUMBER=$RETRIES_NUMBER

if [ -z "$RETRIES_NUMBER" ]; then
    RETRIES_NUMBER=2
fi

# turn on debug mode by extra option "debug"
if [[ "$(echo "$EXTRA_OPTIONS" | grep debug)" != "" ]]; then
    DEBUG=1
fi

# finding IPv4 addresses from application names.
name_resolver() {

    local DNS_IP
    local DNS=$1
    APP_IP=""
    UP_COUNT=0
    SRV_COUNT=0

    echo "DNS: "$DNS

    for D in $(echo $DNS); do

        if [ -z "$STRICK_CHECK" ]; then
            # find $D as SELECTOR in hosts file
            EXISTS=$(grep -w $D $HOST_FILE)
            #EXISTS=$(grep -w "$D-" $HOST_FILE); # TODO?
            if [ -n "$EXISTS" ]; then # selector exists in hosts file
                # remove all matching selectors and all selctors followed by "-"
                #APP_IP=$(echo $EXISTS | sed s/$D-.//g | sed s/$D//g);
                APP_IP=$(echo "$EXISTS" | awk '{print $1}')
                debug "APP_IP: "$APP_IP
            else
                debug "no matching APPLICATION NAME found in $HOST_FILE"
            fi
        else
            D=$(echo $D | cut -d "-" -f1)
            UP=$(docker ps --format '{{.Names}}\t{{.Status}}' | grep Up | awk '{print $1}' | grep $D"-")
            # filtering for ROLES variables if exists.
            if [[ "$ROLES" != "null" && ! -z "$ROLES" ]]; then
                UPS=""
                for U in $(echo $UP); do
                    for ROLE in $(echo $ROLES); do
                        FILTERED_BY_ROLE=$(docker inspect $U -f '{{.Config.Labels.roles}}' | uniq | grep $ROLE)
                        if [[ "$(echo $FILTERED_BY_ROLE)" != "" ]]; then
                            UPS="$UPS $U"
                        fi
                    done
                done
                UP=$UPS
            fi

            UP_COUNT=$((UP_COUNT + $(echo $UP | wc -w)))

            for SRV_FILE in $(echo $SERVICE_FILES); do
                CONTAINER_NAMES=$(jq -r .containers[].NAME $SRV_FILE)
                for NAME in $(echo $CONTAINER_NAMES); do
                    NEWNAME=$(echo $NAME | cut -d "-" -f1)
                    if [ "$D" == "$NEWNAME" ]; then
                        if [[ "$ROLES" != "null" && ! -z "$ROLES" ]]; then
                            C_ROLES=$(jq -r --arg NAME "$NAME" '.containers[] | select(.NAME==$NAME)' $SRV_FILE | jq -r .ROLES)
                            for ROLE in $(echo $ROLES); do
                                # TODO, ha C_ROLES tobb erteket tartalmaz
                                if [ "$ROLE" == "$C_ROLES" ]; then
                                    SRV_COUNT=$((SRV_COUNT + 1))
                                fi
                            done
                        else
                            SRV_COUNT=$((SRV_COUNT + 1))
                        fi
                    fi

                done
            done

            if [ ! -z "$UP" ]; then
                for D_IP in $(echo $UP); do
                    DNS_IP=$(docker inspect $D_IP -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}')
                    if [ "$APP_IP" == "" ]; then
                        APP_IP="$DNS_IP"
                    else
                        APP_IP="$APP_IP $DNS_IP"
                    fi
                    echo "APP_IP: "$APP_IP
                done
            else
                debug "no matching running process found"
            fi
        fi
    done

    if [[ ! -z "$STRICK_CHECK" && $UP_COUNT -lt $SRV_COUNT ]]; then
        if [ "$2" == "" ]; then
            RETRIES=0
        else
            RETRIES=$2
        fi
        if [ $RETRIES -le $RETRIES_NUMBER ]; then
            debug "Try to reread container names at $RETRIES"
            sleep 1
            RETRIES=$((RETRIES + 1))
            name_resolver $DNS $RETRIES
        else
            debug "Not enough running process found for executing firewall rules, exiting"
            exit
        fi
    fi
}

if [[ -z "$TYPE" ]]; then
    TYPE="tcp"
fi

if [[ -z "$SOURCE_IP" ]]; then
    if [[ -z "$SOURCE" ]]; then
        #SOURCE_IP="0.0.0.0/0";
        echo "No source IP added"
    elif [ "$(set | grep -w SOURCE_IFACE)" != "" ]; then
        SOURCE=$SOURCE_IFACE
        echo "VPN interface added instead of IP or domain name"
    else
        IDX=0
        for i in $(echo $SOURCE); do

            if [[ "$i" != [0-9]*"."[0-9]*"."[0-9]*"."[0-9]* ]]; then
                name_resolver $i
                debug "source ip is $APP_IP"
                if [[ -z "$APP_IP" ]]; then
                    debug "No any IP address found for SOURCE: $SOURCE, try again to resolv"
                    name_resolver $i
                    debug "source ip is $APP_IP"

                    if [[ -z "$APP_IP" ]]; then
                        debug "No any IP address found for SOURCE: $SOURCE, giving up"
                    fi
                fi

                for IP in $(echo $APP_IP); do
                    IDX=$(expr 1 + $IDX)
                    eval SOURCE_IP_$IDX=$IP
                done

            else
                IDX=$(expr 1 + $IDX)
                if [[ "$(echo $i | cut -d . -f4)" == "0" ]]; then
                    SOURCE_IP="$i/24"
                    eval SOURCE_IP_$IDX="$i/24"
                    debug "source ip is $SOURCE_IP"
                else
                    eval SOURCE_IP_$IDX=$i
                    IP=$i
                    debug "source ip is $IP"
                fi
            fi
        done

        if [ $IDX = 1 ]; then
            SOURCE_IP=$IP
        fi

    fi
fi

if [[ -z "$TARGET_IP" ]]; then
    if [[ -z "$TARGET" ]]; then
        #TARGET_IP="0.0.0.0/0";
        echo "No target IP added"
    else
        IDX=0
        for i in $(echo $TARGET); do

            if [[ "$i" != [0-9]*"."[0-9]*"."[0-9]*"."[0-9]* ]]; then
                name_resolver $i
                debug "target ip is $APP_IP"

                if [[ -z "$APP_IP" ]]; then
                    debug "No any IP address found for TARGET: $TARGET, try again to resolv"
                    name_resolver $i
                    debug "target ip is $APP_IP"

                    if [[ -z "$APP_IP" ]]; then
                        debug "No any IP address found for TARGET: $TARGET, giving up"
                    fi
                fi

                for IP in $(echo $APP_IP); do
                    IDX=$(expr 1 + $IDX)
                    eval TARGET_IP_$IDX=$IP
                done

            else
                IDX=$(expr 1 + $IDX)
                if [[ "$(echo $i | cut -d . -f4)" == "0" ]]; then
                    TARGET_IP="$i/24"
                    eval TARGET_IP_$IDX="$i/24"
                    debug "target ip is $TARGET_IP"
                else
                    eval TARGET_IP_$IDX=$i
                    IP=$i
                    debug "target ip is $IP"
                fi
            fi
        done

        if [ $IDX = 1 ]; then
            TARGET_IP=$IP
        fi

    fi
fi

delete_lines() {

    if [ "$1" != "" ]; then
        CHAIN=$1
    fi

    if [ -n "$LINES" ]; then
        for i in $LINES; do
            debug "$IPTABLES -D $CHAIN $i"
            $IPTABLES -w -D $CHAIN $i
            sleep 0.1
        done
    fi
}

prerouting() {
    if [ "$(set | grep -w SOURCE_IFACE)" != "" ]; then
        if [ "$TARGET_IP" != "" ] && [ "$TARGET_PORT" != "" ]; then
            debug "$IPTABLES -I PREROUTING -i $SOURCE_IFACE -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment $COMMENT -j DNAT --to $TARGET_IP:$TARGET_PORT"
            $IPTABLES -w -I PREROUTING -i $SOURCE_IFACE -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment "$COMMENT" -j DNAT --to $TARGET_IP:$TARGET_PORT
        fi
    elif [ "$SOURCE_IP" != "" ]; then
        if [ "$SOURCE_PORT" != "" ]; then
            if [ "$TARGET_IP" != "" ]; then
                if [ "$TARGET_PORT" != "" ]; then

                    LINES=$($IPTABLES -w -L --line-number -n | grep DNAT | grep $SOURCE_PORT | grep $TARGET_IP | grep $TARGET_PORT | grep $COMMENT | awk '{print $1}' | tac)
                    debug "Previous prerouting lines: "$LINES
                    # DELETE UNECESSARY LINES FROM PREVIOUS RULES
                    delete_lines "PREROUTING"

                    debug "$IPTABLES -I PREROUTING -d $SOURCE_IP -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment $COMMENT -j DNAT --to $TARGET_IP:$TARGET_PORT"
                    $IPTABLES -w -I PREROUTING -d $SOURCE_IP -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment "$COMMENT" -j DNAT --to $TARGET_IP:$TARGET_PORT
                fi
            fi
        fi
    fi
}

postrouting() {

    if [ -n "$SOURCE_IP" ]; then
        SOURCE_IP_FOR_POSTROUTING="$(echo $SOURCE_IP | cut -d . -f1-3).0/24"
        debug "source ip is $SOURCE_IP_FOR_POSTROUTING"

        LINES=$($IPTABLES -w -L --line-number -n | grep MASQUERADE | grep $COMMENT | grep $SOURCE_IP_FOR_POSTROUTING | grep $SOURCE_PORT | awk '{print $1}' | tac)
        debug "Previous postrouting lines: "$LINES

        # DELETE UNECESSARY LINES FROM PREVIOUS RULES
        delete_lines "POSTROUTING"

        debug "$IPTABLES -I POSTROUTING -s $SOURCE_IP_FOR_POSTROUTING -p $PROTOCOL --sport $SOURCE_PORT -m comment --comment "$COMMENT" -j MASQUERADE"
        $IPTABLES -w -I POSTROUTING -s $SOURCE_IP_FOR_POSTROUTING -p $PROTOCOL --sport $SOURCE_PORT -m comment --comment "$COMMENT" -j MASQUERADE
    fi

    if [ -n "$TARGET_IP" ]; then
        TARGET_IP_FOR_POSTROUTING="$(echo $TARGET_IP | cut -d . -f1-3).0/24"
        debug "target ip is $TARGET_IP_FOR_POSTROUTING"

        LINES=$($IPTABLES -w -L --line-number -n | grep $COMMENT | grep $TARGET_IP_FOR_POSTROUTING | grep $TARGET_PORT | awk '{print $1}' | tac)
        debug "Previous postrouting lines: "$LINES

        # DELETE UNECESSARY LINES FROM PREVIOUS RULES
        delete_lines "POSTROUTING"

        debug "$IPTABLES -I POSTROUTING -s $TARGET_IP_FOR_POSTROUTING -p $PROTOCOL --dport $TARGET_PORT -m comment --comment "$COMMENT" -j MASQUERADE"
        $IPTABLES -w -I POSTROUTING -d $TARGET_IP_FOR_POSTROUTING -p $PROTOCOL --dport $TARGET_PORT -m comment --comment "$COMMENT" -j MASQUERADE
    fi
}

ip_route() {

    COUNT_NETWORK=$(set | grep NETWORK | wc -l)

    for network_index in $(seq 1 $COUNT_NETWORK); do
        if set | grep NETWORK_; then
            NETWORK=$(eval "echo \${"NETWORK_$network_index"}")
            GATEWAY=$(eval "echo \${"GATEWAY_$network_index"}")
        fi
        debug "ip route add "$NETWORK"/24 via "$GATEWAY
        $IP_ROUTE add $NETWORK/24 via $GATEWAY
    done
}

if [[ "$ROUTE" == "true" ]]; then
    IP_ROUTE="nsenter -t $(docker inspect --format {{.State.Pid}} $NAME) -n -- ip route"

    ip_route
    exit
fi

##############################

echo 1 >/proc/sys/net/ipv4/ip_forward

##############################

if /usr/sbin/iptables-legacy -L | grep DOCKER-USER; then

    IPTABLES="/usr/sbin/iptables-legacy"
else

    IPTABLES="/usr/sbin/iptables"
fi

###############################

COUNT_SOURCE_IP=$(set | grep SOURCE_IP | wc -l)
COUNT_SOURCE_PORT=$(set | grep SOURCE_PORT | wc -l)
COUNT_TARGET_IP=$(set | grep TARGET_IP | wc -l)
COUNT_TARGET_PORT=$(set | grep TARGET_PORT | wc -l)

# SOURCE AND TARGET PORTS ARE IN PAIRS
if [ "$COUNT_SOURCE_PORT" == "$COUNT_TARGET_PORT" ]; then
    PAIRS="1"
else
    PAIRS="0"
fi

if [ "$COUNT_SOURCE_IP" == 0 ]; then COUNT_SOURCE_IP=1; fi
for source_ip_index in $(seq 1 $COUNT_SOURCE_IP); do
    if set | grep SOURCE_IP_; then
        SOURCE_IP=$(eval "echo \${"SOURCE_IP_$source_ip_index"}")
    fi

    if [ "$COUNT_SOURCE_PORT" == 0 ]; then COUNT_SOURCE_PORT=1; fi
    for source_port_index in $(seq 1 $COUNT_SOURCE_PORT); do
        if set | grep SOURCE_PORT_; then
            SOURCE_PORT=$(eval "echo \${"SOURCE_PORT_$source_port_index"}")
        fi

        if [ "$COUNT_TARGET_IP" == 0 ]; then COUNT_TARGET_IP=1; fi
        for target_ip_index in $(seq 1 $COUNT_TARGET_IP); do
            if set | grep TARGET_IP_; then
                TARGET_IP=$(eval "echo \${"TARGET_IP_$target_ip_index"}")
            fi

            if [ "$COUNT_TARGET_PORT" == 0 ]; then COUNT_TARGET_PORT=1; fi
            for target_port_index in $(seq 1 $COUNT_TARGET_PORT); do
                if set | grep TARGET_PORT_; then
                    TARGET_PORT=$(eval "echo \${"TARGET_PORT_$target_port_index"}")
                fi

                debug "PAIRS: $PAIRS"
                debug "source_port_index: $source_port_index"
                debug "target_port_index: $target_port_index"

                # if case of pairs if indexes doesn't match then omit routing
                if [ "$PAIRS" == "1" ] && [ "$source_port_index" != "$target_port_index" ]; then
                    debug "OMIT ROUTING"
                    continue
                fi

                #############################
                # NSENTER Specific settings #

                if [[ "$PREROUTING" == "true" ]] || [[ "$POSTROUTING" == "true" ]] || [[ "$HOST" == "true" ]]; then

                    if [ "$HOST" == "true" ]; then
                        IPTABLES="/sbin/iptables -t nat"
                        debug "iptables: "$IPTABLES
                    else
                        IPTABLES="nsenter -t $(docker inspect --format {{.State.Pid}} $NAME) -n -- /sbin/iptables-legacy -t nat"
                        debug "iptables: "$IPTABLES
                    fi

                    if [[ "$PREROUTING" == "true" ]]; then
                        prerouting
                    fi

                    if [[ "$POSTROUTING" == "true" ]]; then
                        postrouting
                    fi

                else

                    ############################
                    # Host firewall settings ###

                    if $IPTABLES -w -n --list $CHAIN | grep ESTABLISHED | grep RELATED | grep ACCEPT; then
                        echo "nothing to do"
                    else
                        $IPTABLES -w -I $CHAIN -m state --state established,related -j ACCEPT
                    fi

                    IPTABLES_OPTIONS=""
                    GREP_OPTIONS=""
                    if [ "$SOURCE_IP" != "" ]; then

                        if [ "$(echo $SOURCE_IP | cut -d . -f4)" == "0" ]; then
                            SOURCE_IP="$(echo $SOURCE_IP | cut -d . -f1-3).0/24"
                        fi

                        IPTABLES_OPTIONS=$IPTABLES_OPTIONS" -s $SOURCE_IP"
                        GREP_OPTIONS=$GREP_OPTIONS"|grep -e $SOURCE_IP"

                        if [ "$SOURCE_PORT" != "" ]; then
                            IPTABLES_OPTIONS=$IPTABLES_OPTIONS" --sport $SOURCE_PORT"
                            GREP_OPTIONS=$GREP_OPTIONS"|grep -e $SOURCE_PORT"
                        fi
                    fi

                    if [ "$TARGET_IP" != "" ]; then

                        if [ "$(echo $TARGET_IP | cut -d . -f4)" == "0" ]; then
                            TARGET_IP="$(echo $TARGET_IP | cut -d . -f1-3).0/24"
                        fi

                        IPTABLES_OPTIONS=$IPTABLES_OPTIONS" -d $TARGET_IP"
                        GREP_OPTIONS=$GREP_OPTIONS"|grep -e $TARGET_IP"

                        if [ "$TARGET_PORT" != "" ]; then
                            IPTABLES_OPTIONS=$IPTABLES_OPTIONS" --dport $TARGET_PORT"
                            GREP_OPTIONS=$GREP_OPTIONS"|grep -e $TARGET_PORT"
                        fi
                    fi

                    if [[ "$SOURCE_IP" != "" && "$TARGET_IP" != "" ]]; then

                        #
                        # DELETE UNECESSARY LINES FROM PREVIOUS RULES
                        IPTABLES_COMMAND="$IPTABLES -w --line-number -n --list $CHAIN | grep $PROTOCOL $GREP_OPTIONS | awk '{print \$1}'| tac"
                        debug "$IPTABLES_COMMAND"
                        LINES=$(eval $IPTABLES_COMMAND)
                        delete_lines

                        if [ "$OPERATION" == "DELETE" ]; then

                            IPTABLES_COMMAND="$IPTABLES -w --line-number -n --list $CHAIN | grep -w "$COMMENT" | awk '{print \$1}'| tac"
                            debug "$IPTABLES_COMMAND"
                            LINES=$(eval $IPTABLES_COMMAND)
                            delete_lines

                        else
                            debug "$IPTABLES -I $CHAIN -p $PROTOCOL $IPTABLES_OPTIONS -m comment --comment "$COMMENT" -j ACCEPT"
                            $IPTABLES -w -I $CHAIN -p $PROTOCOL $IPTABLES_OPTIONS -m comment --comment "$COMMENT" -j ACCEPT
                        fi
                    fi
                    #############################
                fi

            done # target_port
        done     # target_ip
    done         # source_port
done             # source_ip