From 7328bbac8fcf64f092729822b6a7ff5be6f87bb5 Mon Sep 17 00:00:00 2001
From: Manuel Friedli <manuel@fritteli.ch>
Date: Tue, 15 Sep 2020 14:06:54 +0200
Subject: [PATCH] Support a list of countries to ban.

---
 ddos-mitigator.sh | 62 +++++++++++++++++++++++++++--------------------
 1 file changed, 36 insertions(+), 26 deletions(-)

diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh
index 22fee09..3a9a88b 100755
--- a/ddos-mitigator.sh
+++ b/ddos-mitigator.sh
@@ -80,32 +80,38 @@ function print_help() {
 Usage: $(basename $0) -d FILE [OPTION...]
 
   ${bold}Mandatory options:${reset}
-  -d, --database=FILE    The path to the GeoIP2 database file (must be either
-                         country or city database).
+  -d, --database=FILE                 The path to the GeoIP2 database file (must
+                                      be either country or city database).
 
   ${bold}Optional options:${reset}
-  -a, --auto[=LIMIT]     Enable the autopilot for automatically banning IP
-                         addresses of the desired country (see also -c option).
-                         When LIMIT is given, only auto-ban IP addresses with at
-                         least LIMIT current connections.
-                         When LIMIT is omitted, assume LIMIT=1.
+  -a, --auto[=LIMIT]                  Enable the autopilot for automatically
+                                      banning IP addresses of the desired
+                                      countries (see also -c option).
+                                      When LIMIT is given, only auto-ban IP
+                                      addresses with at least LIMIT current
+                                      connections.
+                                      When LIMIT is omitted, assume LIMIT=1.
 
-  -c, --country=COUNTRY  The country-code to block; defaults to 'CN' (China).
+  -c, --country=COUNTRY[,COUNTRY...]  The country-codes to block as a list of
+                                      comma-separated values; defaults to 'CN'
+                                      (China).
 
-  -j, --jail=JAIL        Specify the JAIL to use for banning the IP addresses.
-                         Defaults to 'apache-auth'.
+  -j, --jail=JAIL                     Specify the JAIL to use for banning the IP
+                                      addresses.
+                                      Defaults to 'apache-auth'.
 
-  -n, --netmask=SIZE     SIZE defines the subnet size in bytes to be analyzed.
-                         Valid values are:
-                         - 1 or 8 for class A networks (X.0.0.0/8)
-                         - 2 or 16 for class B networks (X.X.0.0/16)
-                         - 3 or 24 for class C networks (X.X.X.0/24)
-                         - 4 or 32 for class D networks (X.X.X.X/32)
+  -n, --netmask=SIZE                  SIZE defines the subnet size in bytes to
+                                      be analyzed.
+                                      Valid values are:
+                                      - 1 or 8 for class A networks (X.0.0.0/8)
+                                      - 2 or 16 for class B networks (X.X.0.0/16)
+                                      - 3 or 24 for class C networks (X.X.X.0/24)
+                                      - 4 or 32 for class D networks (X.X.X.X/32)
 
-  -p, --port=PORT        The desired port to monitor.
-                         Defaults to 443 (https).
+  -p, --port=PORT                     The desired port to monitor.
+                                      Defaults to 443 (https).
 
-  -h, --help             Show this help message
+  -h, --help                          Show this help message
 
 Mandatory or optional arguments to long options are also mandatory or optional
 for any corresponding short options.
@@ -174,7 +180,11 @@ function parse_command_line_args() {
 				shift
 			;;
 			'-c'|'--country')
-				bancountry="$2"
+				IFS=',' read -ra bancountries <<< "$2"
+				if [[ -z ${bancountries[@]// } ]] ; then
+					echo "Invalid argument for parameter 'country': '$2'. Invoke with --help for help." >&2
+					exit 1
+				fi
 				shift
 			;;
 			'-d'|'--database')
@@ -320,16 +330,16 @@ function process_file () {
 		set_highlight_color "${count}"
 		country="$("${curdir}/geoip-lookup.py" -f "${database}" "${addronly}")"
 		if [[ autopilot -eq 0 ]] ; then
-			echo "Country: '${country}'"
+			echo "Country: '${yellow}${country}${reset}'"
 		fi
 		echo -n "Address ${bold}$((nline++)) of ${nlines}${reset}: \
 Found '${blue}${addrwithsuffix}${reset}' ${hilite}${count}${reset} times."
 
 		if [[ ${autopilot} -eq 0 ]] ; then
-			echo -n  "Ban [y/N/s=No, and skip remaining]? "
+			echo -n " Ban [y/N/s=No, and skip remaining]? "
 			read banaction
 		else
-			if [[ "${country}" == "${bancountry}" ]] ; then
+			if [[ " ${bancountries[@]} " =~ " ${country} " ]] ; then
 				if [[ $count -ge $autopilot ]] ; then
 					echo -en "\n${red}Autopilot active. ${reset}"
 					banaction=y
@@ -356,11 +366,11 @@ skipping remaining addresses."
 			;;
 			"y" | "Y" )
 				echo "Adding '${blue}${addrwithsuffix}${reset}' to \
-banlist."
+banlist (country=${yellow}${country}${reset})."
 				echo "${addrwithsuffix}" >> "${banlist}"
 			;;
 			"n" | "N" | * )
-				echo "Not banning '${blue}${addrwithsuffix}${reset}' (country='${country}')."
+				echo "Not banning '${blue}${addrwithsuffix}${reset}' (country=${yellow}${country}${reset})."
 			;;
 		esac
 	# Here goes: Pipe the file contents via filedescriptor 3.
@@ -392,7 +402,7 @@ touch "${banlist}"
 autopilot=0
 netmask=0
 jail="apache-auth"
-bancountry="CN"
+bancountries=("CN")
 database=
 port=443