#!/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

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

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)

# 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