#!/bin/sh
# Script to manage iptables blocking by country.
# This requires the xtables-addons:
# 	sudo aptitude install xtables-addons-common libtext-csv-xs-perl

# Block Russia, Belarus, China (a random selection for testing)
COUNTRIES="ru by cn" 


GEOIP_DIR=/usr/share/xt_geoip
GEOIP_CHAIN=geoip_restrictions

if [ $(id -u) -ne 0 ]; then
	echo "Please run this script as root."
	exit
fi

if [ ! -d ${GEOIP_DIR} ]; then
	mkdir ${GEOIP_DIR}
fi

do_start() {
	if [ $(iptables -n -L ${GEOIP_CHAIN} --line-numbers 2>/dev/null | grep ^[0-9] | wc -l) -ne 0 ]; then
		echo "iptables chain ${GEOIP_CHAIN} already exists. Restart if you wish to change it. Exiting..."
		return
	fi
	iptables -N ${GEOIP_CHAIN}
	# Always allow web and ssh
	iptables -A ${GEOIP_CHAIN} -p tcp -m multiport --dports http,https -j ACCEPT
	for country in ${COUNTRIES}; do
		iptables -A ${GEOIP_CHAIN} -m geoip --dst-cc ${country} -j REJECT
	done
	# Jump back to OUTPUT
	iptables -A ${GEOIP_CHAIN} -j RETURN
	# Jump to it
	iptables -t filter -A OUTPUT -j ${GEOIP_CHAIN}
	return
}

do_stop() {
	if [ $(iptables -n -L ${GEOIP_CHAIN} --line-numbers 2>/dev/null | grep ^[0-9] | wc -l) -eq 0 ]; then
		echo "iptables chain ${GEOIP_CHAIN} doesn't exist. Exiting..."
		return
	fi
	# Remove jump
	iptables -t filter -D OUTPUT -j ${GEOIP_CHAIN}
	# Flush and delete chain
	iptables -F ${GEOIP_CHAIN}
	iptables -X ${GEOIP_CHAIN}
	return
}

do_update() {
	TEMP_DIR=$(mktemp -d)
	# Download
	mkdir -p ${TEMP_DIR}
	cd ${TEMP_DIR}
	/usr/lib/xtables-addons/xt_geoip_dl 1>/dev/null 2>&1
	# Compile
	if [ $(ls -1 *.csv 2>/dev/null | wc -l) -ne 0 ]; then
		/usr/lib/xtables-addons/xt_geoip_build -D ${GEOIP_DIR} *.csv >/dev/null
	fi
	# Cleanup
	cd /root
	rm -rf ${TEMP_DIR}
	return
}

case "$1" in
start)
	do_start
	echo "GeoIP restrictions enabled."
	;;
stop)
	do_stop
	echo "GeoIP restrictions removed."
	;;
restart|reload)
	do_stop
	do_start
	echo "GeoIP restrictions reloaded."
	;;
update)
	do_update
	# Restart iptables
	do_stop
	do_start
	echo "GeoIP database updated."
	;;
*)
	echo "Usage: $0 {start|stop|restart|update}"
	exit 1
	;;
esac