#!/bin/bash
#
# Dependencies: ipset, wget, grep
# Requirements: Run as root or with sudo
# https://snork.ca/gitsucks/#wankerblock

export PATH='/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
MYAGENT="Mozilla/5.0 (Windows NT 10.0; rv:119.0) Gecko/20100101 Firefox/119.0"
BASE1="https://github.com/herrbischoff/country-ip-blocks/raw/master/"
BASE2="https://snork.ca/gitsucks/projects/wankerblock/"
WORKDIR="/tmp/wankerblock"
TMPSET="wankerblocktempset"
RESTOREFILE="restorefile.txt"

# Make sure there are the proper number of arguments
if [ ${#} -ne 3 ]; then
  echo "Abort: Incorrect number of arguments."
  echo "Syntax: wankerblock <ipv4|ipv6> <ipsetname> <countrylist>"
  echo "See https://snork.ca/gitsucks/#wankerblock for instructions."
  exit 1
fi

# Make sure user has required privilege
if [ ${EUID} -ne 0 ]; then
   echo "This script must be run as root or with sudo to manage ipsets."
   exit 1
fi

# Does user want ipv4 or ipv6?
case ${1} in
  ipv4)
    IPFAM="inet"
    ;;
  ipv6)
    IPFAM="inet6"
    ;;
  *)
    echo "Abort: The first argument must be either ipv4 or ipv6!"
    echo "Syntax: wankerblock <ipv4|ipv6> <ipsetname> <countrylist>"
    echo "See https://snork.ca/gitsucks/#wankerblock for instructions."
    exit 1
    ;;
esac

# Does requested set already exist?
ipset list | grep "^Name: ${2}$" > /dev/null
if [ ${?} -eq 0 ]; then
  # ipset already exists
  ipset list ${2} | grep inet6 > /dev/null
  if [ ${?} -eq 0 ]; then
    # Existing ipset is IPv6
    OLDSET="inet6"
  else
    # Existing ipset is IPv4
    OLDSET="inet"
  fi
else
  # ipset does not exist yet, create it
  OLDSET=${IPFAM}
  ipset create ${2} hash:net family ${IPFAM} -exist
fi

# What if destination ipset exists and is the wrong inet family???
if [ ${IPFAM} != ${OLDSET} ]; then
  # Requested ipset is wrong family
  echo "Abort: The requested ipset already exists but is not ${1}"
  echo "Syntax: wankerblock <ipv4|ipv6> <ipsetname> <countrylist>"
  echo "See https://snork.ca/gitsucks/#wankerblock for instructions."
  exit 1
fi

# Looks like this is going to happen...
# Create working directory and temporary ipset
if [ -d ${WORKDIR} ]; then
  rm -R ${WORKDIR}
fi
mkdir -p ${WORKDIR}
ipset list | grep "^Name: ${TMPSET}$" > /dev/null
if [ ${?} -eq 0 ]; then
  ipset destroy ${TMPSET}
fi
ipset create ${TMPSET} hash:net family ${IPFAM} -exist
ipset flush ${TMPSET}

# run through the downloads and fill the temp set
TAGS=( ${3} )
for a in ${TAGS[@]}
do
  if [ ${#a} -ne 2 ]; then
    # This array element is NOT a country code (get from me)
    wget -O ${WORKDIR}/${a}.txt -U "${MYAGENT}" ${BASE2}${a}-${1}.txt
  else
    # This array element is a country code (get from herrbischoff)
    wget -O ${WORKDIR}/${a}.txt -U "${MYAGENT}" ${BASE1}${1}/${a}.cidr
  fi
  # Convert and restore
  sed -i '/^#/d' ${WORKDIR}/${a}.txt
  sed -i '/^$/d' ${WORKDIR}/${a}.txt
  sed "s/^/add ${TMPSET} /" ${WORKDIR}/${a}.txt >> ${WORKDIR}/${RESTOREFILE}
  ipset flush ${TMPSET}
done

sort ${WORKDIR}/${RESTOREFILE} | uniq > ${WORKDIR}/${RESTOREFILE}.ipset
ipset restore < ${WORKDIR}/${RESTOREFILE}.ipset
ipset swap ${TMPSET} ${2}

# ipset complete, now do a little cleanup
ipset destroy ${TMPSET}
cp ${WORKDIR}/${RESTOREFILE}.ipset /tmp
rm -R ${WORKDIR}