#!/bin/bash
#Nifty: Easily block IP of entire countries

curl -s https://nifty.sh/lib/universal/niftylib.sh > /tmp/lib.nifty.sh
. /tmp/lib.nifty.sh
rm -f /tmp/lib.nifty.sh

# Requirements
require_root
require_64bit

OS=`detect_os`
if [[ "$OS" != "centos7" && "$OS" != "alma8" ]]; then
  err "This script supports only CentOS 7 and 8"
fi

if [ ! -f "/usr/bin/whiptail" ]; then
  yum install -y newt
fi

# Function to present choices to the user
function choices() {
  whiptail --title "Block countries" --ok-button "Block selected"  --checklist \
  "Choose countries to block" 20 78 15 \
  "cn" "China" OFF \
  "ua" "Ukraine" OFF \
  "ru" "Russian Federation" OFF \
  "hk" "Hong Kong" OFF \
  "af" "Afghanistan" OFF \
  "ax" "Aaland Islands" OFF \
  "al" "Albania" OFF \
  "dz" "Algeria" OFF \
  "as" "American Samoa" OFF \
  "ad" "Andorra" OFF \
  "ao" "Angola" OFF \
  "ai" "Anguilla" OFF \
  "aq" "Antarctica" OFF \
  "ag" "Antigua and Barbuda" OFF \
  "ar" "Argentina" OFF \
  "am" "Armenia" OFF \
  "aw" "Aruba" OFF \
  "au" "Australia" OFF \
  "at" "Austria" OFF \
  "az" "Azerbaijan" OFF \
  "bs" "Bahamas" OFF \
  "bh" "Bahrain" OFF \
  "bd" "Bangladesh" OFF \
  "bb" "Barbados" OFF \
  "by" "Belarus" OFF \
  "be" "Belgium" OFF \
  "bz" "Belize" OFF \
  "bj" "Benin" OFF \
  "bm" "Bermuda" OFF \
  "bt" "Bhutan" OFF \
  "bo" "Bolivia, Plurinational State of" OFF \
  "bq" "Bonaire, Sint Eustatius and Saba" OFF \
  "ba" "Bosnia and Herzegovina" OFF \
  "bw" "Botswana" OFF \
  "bv" "Bouvet Island" OFF \
  "br" "Brazil" OFF \
  "io" "British Indian Ocean Territory" OFF \
  "bn" "Brunei Darussalam" OFF \
  "bg" "Bulgaria" OFF \
  "bf" "Burkina Faso" OFF \
  "bi" "Burundi" OFF \
  "kh" "Cambodia" OFF \
  "cm" "Cameroon" OFF \
  "ca" "Canada" OFF \
  "cv" "Cape Verde" OFF \
  "ky" "Cayman Islands" OFF \
  "cf" "Central African Republic" OFF \
  "td" "Chad" OFF \
  "cl" "Chile" OFF \
  "cx" "Christmas Island" OFF \
  "cc" "Cocos (Keeling) Islands" OFF \
  "co" "Colombia" OFF \
  "km" "Comoros" OFF \
  "cg" "Congo" OFF \
  "cd" "Congo, the Democratic Republic of the" OFF \
  "ck" "Cook Islands" OFF \
  "cr" "Costa Rica" OFF \
  "ci" "Côte d'Ivoire" OFF \
  "hr" "Croatia" OFF \
  "cu" "Cuba" OFF \
  "cw" "Curaçao" OFF \
  "cy" "Cyprus" OFF \
  "cz" "Czech Republic" OFF \
  "dk" "Denmark" OFF \
  "dj" "Djibouti" OFF \
  "dm" "Dominica" OFF \
  "do" "Dominican Republic" OFF \
  "ec" "Ecuador" OFF \
  "eg" "Egypt" OFF \
  "sv" "El Salvador" OFF \
  "gq" "Equatorial Guinea" OFF \
  "er" "Eritrea" OFF \
  "ee" "Estonia" OFF \
  "et" "Ethiopia" OFF \
  "fk" "Falkland Islands (Malvinas)" OFF \
  "fo" "Faroe Islands" OFF \
  "fj" "Fiji" OFF \
  "fi" "Finland" OFF \
  "fr" "France" OFF \
  "gf" "French Guiana" OFF \
  "pf" "French Polynesia" OFF \
  "tf" "French Southern Territories" OFF \
  "ga" "Gabon" OFF \
  "gm" "Gambia" OFF \
  "ge" "Georgia" OFF \
  "de" "Germany" OFF \
  "gh" "Ghana" OFF \
  "gi" "Gibraltar" OFF \
  "gr" "Greece" OFF \
  "gl" "Greenland" OFF \
  "gd" "Grenada" OFF \
  "gp" "Guadeloupe" OFF \
  "gu" "Guam" OFF \
  "gt" "Guatemala" OFF \
  "gg" "Guernsey" OFF \
  "gn" "Guinea" OFF \
  "gw" "Guinea-Bissau" OFF \
  "gy" "Guyana" OFF \
  "ht" "Haiti" OFF \
  "hm" "Heard Island and McDonald Islands" OFF \
  "va" "Holy See (Vatican City State)" OFF \
  "hn" "Honduras" OFF \
  "hu" "Hungary" OFF \
  "is" "Iceland" OFF \
  "in" "India" OFF \
  "id" "Indonesia" OFF \
  "ir" "Iran, Islamic Republic of" OFF \
  "iq" "Iraq" OFF \
  "ie" "Ireland" OFF \
  "im" "Isle of Man" OFF \
  "il" "Israel" OFF \
  "it" "Italy" OFF \
  "jm" "Jamaica" OFF \
  "jp" "Japan" OFF \
  "je" "Jersey" OFF \
  "jo" "Jordan" OFF \
  "kz" "Kazakhstan" OFF \
  "ke" "Kenya" OFF \
  "ki" "Kiribati" OFF \
  "kp" "Korea, Democratic People's Republic of" OFF \
  "kr" "Korea, Republic of" OFF \
  "kw" "Kuwait" OFF \
  "kg" "Kyrgyzstan" OFF \
  "la" "Lao People's Democratic Republic" OFF \
  "lv" "Latvia" OFF \
  "lb" "Lebanon" OFF \
  "ls" "Lesotho" OFF \
  "lr" "Liberia" OFF \
  "ly" "Libya" OFF \
  "li" "Liechtenstein" OFF \
  "lt" "Lithuania" OFF \
  "lu" "Luxembourg" OFF \
  "mo" "Macao" OFF \
  "mk" "Macedonia, the Former Yugoslav Republic of" OFF \
  "mg" "Madagascar" OFF \
  "mw" "Malawi" OFF \
  "my" "Malaysia" OFF \
  "mv" "Maldives" OFF \
  "ml" "Mali" OFF \
  "mt" "Malta" OFF \
  "mh" "Marshall Islands" OFF \
  "mq" "Martinique" OFF \
  "mr" "Mauritania" OFF \
  "mu" "Mauritius" OFF \
  "yt" "Mayotte" OFF \
  "mx" "Mexico" OFF \
  "fm" "Micronesia, Federated States of" OFF \
  "md" "Moldova, Republic of" OFF \
  "mc" "Monaco" OFF \
  "mn" "Mongolia" OFF \
  "me" "Montenegro" OFF \
  "ms" "Montserrat" OFF \
  "ma" "Morocco" OFF \
  "mz" "Mozambique" OFF \
  "mm" "Myanmar" OFF \
  "na" "Namibia" OFF \
  "nr" "Nauru" OFF \
  "np" "Nepal" OFF \
  "nl" "Netherlands" OFF \
  "nc" "New Caledonia" OFF \
  "nz" "New Zealand" OFF \
  "ni" "Nicaragua" OFF \
  "ne" "Niger" OFF \
  "ng" "Nigeria" OFF \
  "nu" "Niue" OFF \
  "nf" "Norfolk Island" OFF \
  "mp" "Northern Mariana Islands" OFF \
  "no" "Norway" OFF \
  "om" "Oman" OFF \
  "pk" "Pakistan" OFF \
  "pw" "Palau" OFF \
  "ps" "Palestine, State of" OFF \
  "pa" "Panama" OFF \
  "pg" "Papua New Guinea" OFF \
  "py" "Paraguay" OFF \
  "pe" "Peru" OFF \
  "ph" "Philippines" OFF \
  "pn" "Pitcairn" OFF \
  "pl" "Poland" OFF \
  "pt" "Portugal" OFF \
  "pr" "Puerto Rico" OFF \
  "qa" "Qatar" OFF \
  "re" "Réunion" OFF \
  "ro" "Romania" OFF \
  "rw" "Rwanda" OFF \
  "bl" "Saint Barthélemy" OFF \
  "sh" "Saint Helena, Ascension and Tristan da Cunha" OFF \
  "kn" "Saint Kitts and Nevis" OFF \
  "lc" "Saint Lucia" OFF \
  "mf" "Saint Martin (French part)" OFF \
  "pm" "Saint Pierre and Miquelon" OFF \
  "vc" "Saint Vincent and the Grenadines" OFF \
  "ws" "Samoa" OFF \
  "sm" "San Marino" OFF \
  "st" "Sao Tome and Principe" OFF \
  "sa" "Saudi Arabia" OFF \
  "sn" "Senegal" OFF \
  "rs" "Serbia" OFF \
  "sc" "Seychelles" OFF \
  "sl" "Sierra Leone" OFF \
  "sg" "Singapore" OFF \
  "sx" "Sint Maarten (Dutch part)" OFF \
  "sk" "Slovakia" OFF \
  "si" "Slovenia" OFF \
  "sb" "Solomon Islands" OFF \
  "so" "Somalia" OFF \
  "za" "South Africa" OFF \
  "gs" "South Georgia and the South Sandwich Islands" OFF \
  "ss" "South Sudan" OFF \
  "es" "Spain" OFF \
  "lk" "Sri Lanka" OFF \
  "sd" "Sudan" OFF \
  "sr" "Suriname" OFF \
  "sj" "Svalbard and Jan Mayen" OFF \
  "sz" "Swaziland" OFF \
  "se" "Sweden" OFF \
  "ch" "Switzerland" OFF \
  "sy" "Syrian Arab Republic" OFF \
  "tw" "Taiwan, Province of China" OFF \
  "tj" "Tajikistan" OFF \
  "tz" "Tanzania, United Republic of" OFF \
  "th" "Thailand" OFF \
  "tl" "Timor-Leste" OFF \
  "tg" "Togo" OFF \
  "tk" "Tokelau" OFF \
  "to" "Tonga" OFF \
  "tt" "Trinidad and Tobago" OFF \
  "tn" "Tunisia" OFF \
  "tr" "Turkey" OFF \
  "tm" "Turkmenistan" OFF \
  "tc" "Turks and Caicos Islands" OFF \
  "tv" "Tuvalu" OFF \
  "ug" "Uganda" OFF \
  "ae" "United Arab Emirates" OFF \
  "gb" "United Kingdom" OFF \
  "us" "United States" OFF \
  "um" "United States Minor Outlying Islands" OFF \
  "uy" "Uruguay" OFF \
  "uz" "Uzbekistan" OFF \
  "vu" "Vanuatu" OFF \
  "ve" "Venezuela, Bolivarian Republic of" OFF \
  "vn" "Viet Nam" OFF \
  "vg" "Virgin Islands, British" OFF \
  "vi" "Virgin Islands, U.S." OFF \
  "wf" "Wallis and Futuna" OFF \
  "eh" "Western Sahara" OFF \
  "ye" "Yemen" OFF \
  "zm" "Zambia" OFF \
  "zw" "Zimbabwe" OFF \
  3>&1 1>&2 2>&3
}

# Intro info
whiptail --title "team.blue IP blocker" --msgbox "This script will block countries with ipset and iptables. Be sure not to mess up existing rules." 8 78

# Additional info
whiptail --title "team.blue IP blocker" --msgbox "WARNING! All existing block ranges created by this tool will be flushed. Select ALL countries you want to block including those that might already be blocked now." 8 78

# Check if ipset is installed
if [ ! -f "/usr/sbin/ipset" ]; then

  if (whiptail --title "ipset missing" --yesno "ipset is not installed. Do you want me to install it for you?" 8 78); then
    yum install ipset -y
  else
    echo "User selected No, exit status was $?."
    echo "ERROR: ipset must be present."
    exit 1
  fi

fi


# Get the list of blocks chosen
blocks=$(choices)

# Create a workdir and go to it
workdir=$(mktemp -d -t countryblock-$(date +%Y-%m-%d-%H-%M-%S)-XXXXXXXXXX)
pushd $workdir > /dev/null 2>&1

# Get all current ipsets
IPSETS=$(ipset list |grep "Name" |cut -d" " -f2)

# Delete all iptables rules made by this script earlier
while [ true ]; do
  RULE=$(iptables -L -n --line-numbers | grep -P "^[0-9]+\s+DROP.*match-set tbblock" |awk '{print $1}' |head -n1)

  if [ -z $RULE ]; then
    break
  else
    iptables -D INPUT $RULE
  fi
done


# Loop through selected blocks
for block in $blocks; do

  # $block contains quotes. Remove those
  clean=$(echo $block |sed -e 's/^"//g' -e 's/"$//g')

  # Set a variable with the file name from the download mirror
  FILE="${clean}-aggregated.zone"

  # Generate the link
  DL="https://mirror.dk.team.blue/ipdeny/$FILE"

  # Download the link. TODO: Better error handling
  curl -s $DL > $FILE

  # Decide the name of our ipset
  SETNAME="tbblock_${clean}"

  # Check that the ip file exists before trying to do anything with it
  if [ -f $FILE ]; then

    # Check if this set already exists
    for set in $IPSETS; do
      if [ "$set" = "$SETNAME" ]; then
        # The set exists already. Delete it
        echo "$set exists - deleting it"
        ipset destroy $set
      fi
    done

    # Create the set
    ipset -N $SETNAME hash:net

    # Get cidr list
    CIDR_LIST=$(cat $FILE |grep -v "#")
    CIDR_COUNT=$(cat $FILE |grep -v "#" |wc -l)

    # Populate the set with our ip ranges
    cnt=0
    {
      for cidr in $CIDR_LIST; do 
        cnt=$((cnt+1))
        echo $(( 100*cnt/CIDR_COUNT ))
        ipset -A $SETNAME $cidr
      done
    } | whiptail --gauge "Adding blocks for $clean $1 $2" 6 50 0

    # Our ipset is populated. Now tell iptables to block based on it
    iptables -A INPUT -m set --match-set $SETNAME src -j DROP

    echo "$clean has been blocked ($CIDR_COUNT ranges)!"

  else
    echo "FAIL: $FILE does not exist!"
    exit 1
  fi

done

popd > /dev/null 2>&1
rm -rf $workdir

