#!/bin/bash
## @author gdm85
##
## provide IP-locking of all exposed ports of docker container from a specific whitelist address
## NOTE: iptables rules will be stale once container is stopped/killed/removed
#

SCRIPTS=$(dirname $(readlink -m $0)) || exit $?

if ! type jq 2>&1 >/dev/null; then
	echo "jq command is not available" 1>&2
	exit 1
fi

if [[ $# -lt 2 ]]; then
	echo "Usage: docker-iplock container-name whitelist_ipv4_1 [whitelist_ipv4_2] [...whitelist_ipv4_n]"
	exit 1
fi

function ipt_forward_update() {
	local CID="$1"
	if [[ "$CID" == "<no value>" ]]; then
		echo "Invalid container ID" 1>&2
		return 1
	fi
	local WHITELIST4="$2"

	local CONTAINER_IPv4=$(docker inspect --format '{{ .NetworkSettings.IPAddress }}' $CID)
	if [[ -z "$CONTAINER_IPv4" ]]; then
		return 1
	fi

	test ! -z "$DEBUG" && echo "ipt_forward_update: $CID"

	local STANZA="$(docker inspect --format '{{json .HostConfig.PortBindings }}' $CID)" || return $?
	test ! -z "$DEBUG" && echo "Bindings: $STANZA"

	local I=0
	for P in $(echo "$STANZA" | jq -M 'keys' | tail -n+2 | head -n-1 | sed s/,$// | sed 's/"//g' ); do
		## port on container
		local SRCPORT="$(echo $P | awk -F/ '{ print $1 }')"
		test ! -z "$DEBUG" && echo "$CID source port: $SRCPORT"

		## get port binding (if any)
		local SNIP=$(echo $STANZA | jq -M "to_entries | .[$I].value")

		if [[ -z "$SNIP" || "$SNIP" == "null" ]]; then
			let I=I+1
			continue
		fi
		test ! -z "$DEBUG" && echo "$CID binding: $SNIP"

		local HOSTIP="$(echo $SNIP | jq -M -r '.[0].HostIp')"
		test ! -z "$DEBUG" && echo "$CID host ip: $HOSTIP"

		## match only bindings on docker host
		if [[ -z "$HOSTIP" || "$HOSTIP" == '0.0.0.0' ]]; then
			local HOSTPORT="$(echo $SNIP | jq -M -r '.[0].HostPort')"

			local TOREMOVE="-d ${CONTAINER_IPv4}/32 ! -i docker0 -o docker0 -p tcp -m tcp --dport $SRCPORT -j ACCEPT"
			iptables -D FORWARD $TOREMOVE --wait || return $?

			test ! -z "$DEBUG" && echo "$CID: iptables rule removed: $TOREMOVE"
			local IPv4
			for IPv4 in $WHITELIST4; do
				local TOADD="-s ${IPv4}/32 -d ${CONTAINER_IPv4}/32 ! -i docker0 -o docker0 -p tcp -m tcp --dport $SRCPORT -j ACCEPT"
				iptables -I FORWARD 1 $TOADD --wait || return $?
				test ! -z "$DEBUG" && echo "$CID: iptables rule added: $TOADD"
			done
		fi

		let I=I+1
	done
}

CONTAINER="$1"

shift 1

WHITELIST4="$@"

if [ -z "$WHITELIST4" ]; then
	echo "Invalid whitelist addresses specified" 1>&2
	exit 1
fi

CID=$(docker inspect --format '{{ .Id }}' $CONTAINER) || exit $?

## now run rules update
ipt_forward_update "$CID" "$WHITELIST4"