#!/bin/bash
# This script downloads a list of foreign countries' IP address ranges and it marks these IP ranges via ipset command.
# Author: bisco <bisco__AT__autistici__DOT_org>
# ipset solution has been suggested by Matt Wilcox script: https://mattwilcox.net/web-development/unexpected-ddos-blocking-china-with-ipset-and-iptables

set -e

### Variables
URL="http://www.ipdeny.com/ipblocks/data/countries"
ZONES="cn hk ru tw kp kr sg"
SCRIPTDIR="/root/firewall"
ZONEDIR="${SCRIPTDIR}/blocked_zones/"
RULESFILE="/tmp/iptables_rules-$(date +'%s')"
IPSET=$(which ipset)
IPTR=$(which iptables-restore)
IPTS=$(which iptables-save)
RM=$(which rm)
WGET=$(which wget)
MKDIR=$(which mkdir)
### End Variables


### Code.
### Please don't edit below unless if you know what you're doing.


### check_error function
### Check if things went all done when exit
check_error()
{
    if [ ! $? = 0 ]; then
        echo "====> ERROR!"
        exit
    fi
}


### trap built-in command
### Run "check_error" function on exit process status
trap check_error EXIT


### usage function
### Show help and exits
usage()
{
    echo "Usage: $0 [option]"
    echo ""
    echo "-c | --check     : check all binaries and system are ok"
    echo "-d | --download  : download only the IP range lists"
    echo "-b | --block     : set the IP range lists into the ${ZONES} chain"
    echo "-h | --help      : show this help and exits"
    echo ""
}


### check_binary function
### This function checks if all binaries are properly installed on target system
check_binary()
{
    if [ "${IPSET}" == "" ] || [ "${IPTR}" == "" ] || [ "${RM}" == "" ] || [ "${WGET}" == "" ] || [ "${MKDIR}" == "" ]; then
	echo "Some binaries not found. Please install: ipset, iptables-restore, rm, wget"
	exit 1;
    else
	echo "All binaries are properly installed. Let's go on!"
    fi
}


### prepare_system function
### This function handles all system operation needed.
prepare_system()
{
    # Create the $ZONEDIR directory. The "-p" flag to create the parents directories if needed and it doesn't show error if path exists.
    ${MKDIR} -p "${ZONEDIR}"

    ### Remove old country zone files.
    ### The ${VARIABLE:?} syntax is useful to avoid wrong file deletion, if the variable content is empty.
    ### Examples:
    ### rm -rf "${EMPTY_VAR}"/* executes rm -rf /*
    ### rm -rf "${EMPTY_VAR:?}"/* raises an error
    ${RM} -rf "${ZONEDIR:?}"/*

    ### Save current iptables rules on temporary file
    ${IPTS} > "${RULESFILE}"
}    

### download_zones function
### This function downloads all the selected zones
download_zones()
{
    # Here there's a for cycle on an array of elements
    for zone in $(echo "${ZONES[*]}");
    do
	echo "Downloading $zone.zone file";
	${WGET} --quiet -P "${ZONEDIR}" "${URL}/$zone.zone";
    done
}


### create_chains function
### Create one set per IP range, specifying the hash type to "net"
create_chains()
{
    for chain in $(echo "${ZONES[*]}");
    do
	echo "Creating $chain chain";
	${IPSET} create -exist ${chain} hash:net;
    done
}


### populate_chains
### This function insert the IP range lists into relative chain
populate_chains()
{
    for chain in $(echo "${ZONES[*]}");
    do
	echo "Populating $chain";
	for ip in $(cat $ZONEDIR/$chain.zone);
	do
	    ${IPSET} add -exist $chain $ip;
	done;
    done
}


### reload_rules function
### Reload iptables rules saved during the prepare_system function
### Without reloading iptables rules, new IP ranges won't loaded
reload_rules()
{
    echo "Reloading iptables rules"
    ${IPTR} < "${RULESFILE}"
}


### Main function
case $1 in
    -c | --check)
	check_binary;;
    -d | --download)
	check_binary;
	prepare_system;
	download_zones;;
    -b | --block)
	check_binary;
	prepare_system;
	download_zones;
	create_chains;
	populate_chains;
	reload_rules;;
    -h | --help | *)
	usage;;
esac