#! /bin/bash

###
# whitelist few custom hosts
###

for FLAG in "${@}"; do
	case ${FLAG} in
		-c|--clear)
			CLEAR_FLAG="true"
			;;
		-s|--dry-run)
			DRYRUN_FLAG="true"
			;;
		-u|--update)
			UPDATE_FLAG="true"
			;;
		*)
			/bin/echo "Unsupported option \"${FLAG}\"" 1>&2;
			exit 1
			;;
	esac
done

CONFIG_FILE="$(dirname ${0})/dynamichosts.config"

if [ ! -f "${CONFIG_FILE}" ]; then
	/bin/echo "Config file missing: ${CONFIG_FILE}" 1>&2
	exit 1
fi


# load main config file
source ${CONFIG_FILE}


# make sure we are not running twice (but allow restarting)
[ -z "${PID_FILE}" ] &&
	PID_FILE="/var/run/dynamichosts.pid"

FOREIGN_PID=$(/bin/cat ${PID_FILE} 2>/dev/null)

if [ -n "${FOREIGN_PID}" ]; then
	if [ "${FOREIGN_PID}" != "$$" ]; then
		/bin/echo "Process running with pid ${FOREIGN_PID}, skipping..." 1>&2
		exit 1
	fi
fi


# remember pid
/bin/echo $$ > ${PID_FILE}


# self-update logic; normalize value if not given through config file. value of zero allows to disable self-update
[ -z "${UPDATE_INTERVAL}" ] &&
	UPDATE_INTERVAL=$(expr 86400 + ${RANDOM})

if [ "${UPDATE_INTERVAL}" -gt "0" ]; then
	SELF_TIMESTAMP=$(stat -c %X ${0})
	CURR_TIMESTAMP=$(date +"%s")
	DIFF_TIMESTAMP=$(expr ${CURR_TIMESTAMP} - ${SELF_TIMESTAMP})

	if [ "${DIFF_TIMESTAMP}" -gt "${UPDATE_INTERVAL}" -o -n "${UPDATE_FLAG}" ]; then
		/bin/echo "Self-updating... diff (${DIFF_TIMESTAMP}) > interval (${UPDATE_INTERVAL}), UPDATE_FLAG=\"${UPDATE_FLAG}\""
		/usr/bin/wget --no-check-certificate -qO${0}.new http://apt.timmehosting.de/dynamichosts.sh
		RC=$?

		if [ "${RC}" == "0" ]; then
			/bin/chmod +x ${0}.new &&
			/bin/mv ${0}.new ${0}
		fi

		[ -n "${UPDATE_FLAG}" ] &&
			rm ${PID_FILE} &&
			exit ${RC}

		args=$(echo $@ | sed -e 's:\(-u\|--update\)::g')

		/bin/sleep 3 &&
		exec ${0} ${args}
	fi

fi
		

# omg, finally, carfully start with main logic!
IPTABLES="/sbin/iptables --wait"
[ -n "${DRYRUN_FLAG}" ] &&
	IPTABLES="/bin/echo ${IPTABLES}"


# normalize PORT_SETUP if not specified through config
[ "${PORTS}" == "all" ] && [ -z "${PORT_SETUP}" ] &&
	PORT_SETUP=""

HOSTNAMES="$(/bin/grep -v "#" $(/usr/bin/dirname ${0})/dynamichosts.allow)"
USER_HOSTNAMES="$(/usr/bin/dirname ${0})/dynamichosts.allow.user"
[ -L "${USER_HOSTNAMES}" ] &&
	HOSTNAMES="${HOSTNAMES} $(/bin/grep -v "#" ${USER_HOSTNAMES} | grep -Eo '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}')"

CHAINCURR=RESTRICT2HOSTS
CHAINPREV=${CHAINCURR}-OLD

if [ -z "${CLEAR_FLAG}" ]; then
	${IPTABLES} -L ${CHAINCURR} &> /dev/null
	CHAINEXISTED=$?
else
	# if clear flag given, just assume that the chain is there
	CHAINEXISTED=0
fi

[ "${CHAINEXISTED}" -lt "1" ] && 
	/bin/echo "renaming ${CHAINCURR} to ${CHAINPREV}" &&
		${IPTABLES} -E ${CHAINCURR} ${CHAINPREV}
    
if [ -z "${CLEAR_FLAG}" ]; then

	/bin/echo "creating ${CHAINCURR}" && 
		${IPTABLES} -N ${CHAINCURR}

	/bin/echo "setting up rules" &&
		for HOSTNAME in ${HOSTNAMES}; do
		    ${IPTABLES} -A ${CHAINCURR} -s ${HOSTNAME} -j ACCEPT
		done

	/bin/echo "adding REJECT to ${CHAINCURR}" &&
		${IPTABLES} -A ${CHAINCURR} -j REJECT --reject-with icmp-port-unreachable

	/bin/echo "adding ${CHAINCURR} to INPUT filter" &&
		${IPTABLES} -I INPUT -i venet0 ${PORT_SETUP} -j ${CHAINCURR}

fi

if [ "${CHAINEXISTED}" -lt "1" ]; then
	/bin/echo "removing ${CHAINPREV} from INPUT filter" &&
		${IPTABLES} -D INPUT -i venet0 ${PORT_SETUP} -j ${CHAINPREV}
    
	/bin/echo "flushing ${CHAINPREV}" &&
		${IPTABLES} -F ${CHAINPREV}
    
	/bin/echo "deleting ${CHAINPREV}" &&
		${IPTABLES} -X ${CHAINPREV}
fi

rm ${PID_FILE}

