#!/bin/bash
#-------------------------------------------------------------------
#~ @(#) Name : ux_ipset_update.sh
#~ @(#) Desc : Update the ipset that iptables references for allowing based on country.
# Takes 2 parameters: ipset name (no spaces), country name e.g. 'Australia', 'france' ...
# iptables should have an existing '--match-set' rule e.g
#  $ iptables -I INPUT -p tcp --dport 22 -m set --match-set australia4 src -j ACCEPT
#  $ ip6tables -I INPUT -p tcp --dport 22 -m set --match-set australia6 src -j ACCEPT
#~ @(#) version : 1.1
# Auteur : hello@di-marco.net
# Date : 2024-05-08
#-------------------------------------------------------------------
# Version history
# V1.0 [2020-02-01]
# V1.1 [2024-05-08]
#-------------------------------------------------------------------
#~ Usage : ux_ipset_update.sh whitelist <country>
#-------------------------------------------------------------------

SRC=$(dirname $0)

CHANGELOG="/var/log/ipset_changes.log"
IPSET="/usr/sbin/ipset"

IPSET_UP_RULES="/etc/ipset.up.rules"
IPTABLES_UP_RULES="/etc/iptables.up.rules"

accountID="<accountID>"
licenseKEY="<LicenseKEY>"

if [ $# -ne 2 ]; then
    echo "Usage $0 <ipset_name> <country>"
    exit
fi

SETNAME=$1
COUNTRY=$2
TMPDIR=/tmp/$$_tmp
GEOIP_FILE=${TMPDIR}/$$_geoip.zip
CIDR_FILE=${TMPDIR}/$$_${SETNAME}_CIDR.zone
mkdir -p $TMPDIR

function fail {
    echo "FAILED: "$1
        cleanup
    exit 1
}

function cleanup {
	rm -rf $TMPDIR
}

cd $TMPDIR
[ $? -ne 0 ] && fail "Problem using tmp directory $TMPDIR"

curlout=`curl -J -L -u "${accountID}":"${licenseKEY}" 'https://download.maxmind.com/geoip/databases/GeoLite2-Country-CSV/download?suffix=zip' -o $GEOIP_FILE 2>&1`
[ $? -ne 0 ] && fail "Problem downloading GeoIP database: $curlout"

unzip -q $GEOIP_FILE
cd GeoLite2-Country-CSV*

GEOID=`egrep -i "${COUNTRY}" GeoLite2-Country-Locations-en.csv | cut -f1 -d,`
[ $? -ne 0 ] && fail "$COUNTRY not found in GeoIP file"

for IPvX in 4 6; do

	SET=${SETNAME}${IPvX}
	grep $GEOID GeoLite2-Country-Blocks-IPv${IPvX}.csv | cut -f1 -d, > $CIDR_FILE

	if [ "$IPvX" -eq "4" ]; then
		$IPSET -exist create ${SET} hash:net family inet
		$IPSET create ${SET}-tmp hash:net family inet
	else
		$IPSET -exist create ${SET} hash:net family inet6
		$IPSET create ${SET}-tmp hash:net family inet6
	fi
	[ $? -ne 0 ] && fail "Problem creating tmp ipset"

	while read line; do
		echo "$IPSET add ${SET}-tmp $line;"
		$IPSET add ${SET}-tmp $line;
		[ $? -ne 0 ] && fail "Problem updating tmp ipset"
	done < $CIDR_FILE

	# Check for changes
	if [ -n $CHANGELOG ]; then
		$IPSET save ${SET} > ${TMPDIR}/$$_${SET}.ipset
		$IPSET save ${SET}-tmp > ${TMPDIR}/$$_${SET}-tmp.ipset
		echo `date`": Processing $COUNTRY..." >> $CHANGELOG
		echo `wc -l $CIDR_FILE | cut -f1 -d' '`" blocks loaded" >> $CHANGELOG
		diffout=`diff ${TMPDIR}/$$_${SET}.ipset ${TMPDIR}/$$_${SET}-tmp.ipset`
		if [ $? -ne 0 ]; then 
			echo $diffout >> $CHANGELOG
		else
			echo "No changes" >> $CHANGELOG
		fi
		echo >> $CHANGELOG
	fi

	$IPSET swap ${SET}-tmp $SET
	[ $? -ne 0 ] && fail "Problem updating live ipset"

	$IPSET destroy ${SET}-tmp
	[ $? -ne 0 ] && fail "Problem destroying tmp ipset"
done

cleanup

# Add local IPv4 and IPv6 address
/usr/sbin/ipset add whitelist4 127.0.0.0/8
/usr/sbin/ipset add whitelist4 10.0.0.0/8
/usr/sbin/ipset add whitelist4 172.16.0.0/12
/usr/sbin/ipset add whitelist4 192.168.0.0/16
/usr/sbin/ipset add whitelist6 ::/128
/usr/sbin/ipset add whitelist6 fc00::/7

# Update rules and persit at reboot
if [ -f ${IPSET_UP_RULES} ];then
  cp ${IPSET_UP_RULES} ${IPSET_UP_RULES}.save
  $IPSET save > ${IPSET_UP_RULES}
else 
  $IPSET save > ${IPSET_UP_RULES}
fi

if [ -f ${IPTABLES_UP_RULES} ];then
  cp ${IPTABLES_UP_RULES} ${IPTABLES_UP_RULES}.save
else 
  touch ${IPTABLES_UP_RULES}
fi
cat <<EOF >${IPTABLES_UP_RULES}
# Generated by $0 on $(date)
*filter
:INPUT ACCEPT [0:0]
:FORWARD ACCEPT [0:0]
:OUTPUT ACCEPT [0:0]
:LOGGING - [0:0]
-A INPUT -m state --state RELATED,ESTABLISHED -j ACCEPT
-A INPUT -m set ! --match-set ${1}4 src -j DROP
-A INPUT -m set ! --match-set ${1}6 src -j DROP
COMMIT
# Completed on $(date)
EOF

exit 0