#!/bin/sh

# Debugging function
debug() {
        if [ $DEBUG -eq 1 ]; then
                echo "DEBUG: "$1 $2 $3
        fi; 
}

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

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

EXTRA_OPTIONS="$2 $3 $4"
      
set | grep SOURCE
set | grep TARGET
set | grep ROLES

# 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=""

	echo "DNS: "$DNS;
		
	for D in $(echo $DNS);
	do
		
		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 ROLE in $(echo $ROLES); do
                               	FILTERED_BY_ROLE=$(docker inspect $UP -f '{{.Name}} {{.Config.Labels.roles}}' | uniq | grep $ROLE| awk '{print $1}')
                                if [[ "$(echo $FILTERED_BY_ROLE)" != "" ]]; then                                                                    
                                        UPS="$UPS $FILTERED_BY_ROLE";     
				fi
                        done 
                                                                                                
                        UP=$UPS                                                                                              
                fi              

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

	done;
}

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

if [[ -z "$SOURCE" ]]; then
	SOURCE_IP="0.0.0.0/0";

else
	IDX=0
	for i in $(echo $SOURCE) ; do
	
		if [[ "$i" != *"."* ]]; then
			name_resolver $i;
			debug "source ip is $APP_IP";
			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="$SOURCE_IP/24";
				eval SOURCE_IP_$IDX="$SOURCE_IP/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

if [[ -z "$TARGET" ]]; then
	TARGET_IP="0.0.0.0/0";

else
	IDX=0
	for i in $(echo $TARGET) ; do
	
		if [[ "$i" != *"."* ]]; then
			name_resolver $i;
			debug "target ip is $APP_IP";
			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="$TARGET_IP/24";
				eval TARGET_IP_$IDX="$TARGET_IP/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


prerouting() {
	LINES=$($IPTABLES -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
	if [ -n "$LINES" ] ; then
		for i in $LINES; do 
			debug "$IPTABLES -D PREROUTING $i"
			$IPTABLES -D PREROUTING $i
			sleep 0.1
		done
	fi

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

postrouting() {

	if [[ ! -z "$SOURCE_IP" ]]; then
	SOURCE_IP_FOR_POSTROUTING="$(echo $SOURCE_IP | cut -d . -f1-3).0/24";
	debug "source ip is $SOURCE_IP_FOR_POSTROUTING"
	fi
	
	if [[ ! -z "$TARGET_IP" ]]; then
	TARGET_IP_FOR_POSTROUTING="$(echo $TARGET_IP | cut -d . -f1-3).0/24";
	debug "target ip is $TARGET_IP_FOR_POSTROUTING"
	fi
	
	if  [ -n "$SOURCE_IP" ] ; then
		LINES=$($IPTABLES -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
		if [ -n "$LINES" ] ; then
			for i in $LINES; do 
				debug "$IPTABLES -D POSTROUTING $i"
				$IPTABLES -D POSTROUTING $i
				sleep 0.1
			done
		fi
	
	debug "$IPTABLES -I POSTROUTING -s $SOURCE_IP_FOR_POSTROUTING -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment "$COMMENT" -j MASQUERADE"	
	$IPTABLES -I POSTROUTING -d $SOURCE_IP_FOR_POSTROUTING -p $PROTOCOL --dport $SOURCE_PORT -m comment --comment "$COMMENT" -j MASQUERADE	

	fi 

	if  [ -n "$TARGET_IP" ] ; then
		LINES=$($IPTABLES -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
		if [ -n "$LINES" ] ; then
			for i in $LINES; do 
				debug "$IPTABLES -D POSTROUTING $i"
				$IPTABLES -D POSTROUTING $i
				sleep 0.1
			done
		fi
	debug "$IPTABLES -I POSTROUTING -s $TARGET_IP_FOR_POSTROUTING -p $PROTOCOL --dport $TARGET_PORT -m comment --comment "$COMMENT" -j MASQUERADE"	
	$IPTABLES -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

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

IPTABLES=/sbin/iptables

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

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)

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

#############################
# 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 -- $IPTABLES -t nat";
		debug "iptables: "$IPTABLES;
	fi 	
		
	if [[ "$PREROUTING" == "true" ]] ; then
		prerouting;
	fi

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

else 

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

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

	IPTABLES_OPTIONS=""
	GREP_OPTIONS=""
	if [ "$SOURCE_IP" != "" ]; then
		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
		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 --line-number -n --list $CHAIN | grep $PROTOCOL $GREP_OPTIONS | awk '{print \$1}'| tac";
	debug "$IPTABLES_COMMAND";
	LINES=$(eval $IPTABLES_COMMAND);

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


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

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