From e46a378460df9eb0deedf35a4ac113b176b0a3b6 Mon Sep 17 00:00:00 2001
From: Manuel Friedli <manuel@fritteli.ch>
Date: Mon, 3 Aug 2020 13:32:32 +0200
Subject: [PATCH] "Cherry-pick" (manually) some of the refactoring.

---
 ddos-mitigator.sh | 391 ++++++++++++++++++++++++----------------------
 1 file changed, 204 insertions(+), 187 deletions(-)

diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh
index cffead1..d0eb2ce 100755
--- a/ddos-mitigator.sh
+++ b/ddos-mitigator.sh
@@ -72,46 +72,40 @@ reset="\033[0m"
 trap 'sudo -k; popd >/dev/null; rm -r ${tmpdir}' EXIT
 
 function check_installed() {
-       local command="$1"
-       local package="$2"
-       which "${command}" 2>/dev/null >&2
-       local result=$?
+	local command="$1"
+	local package="$2"
+	which "${command}" 2>/dev/null >&2
+	local result=$?
 
-       if [[ "${result}" -ne 0 ]] ; then
-               echo -e "${red}Command ${bold}${command}${reset}${red} not found.${reset} Please install package ${blue}${package}${reset}."
-               exit 1
-       fi
+	if [[ "${result}" -ne 0 ]] ; then
+		echo -e "${red}Command ${bold}${command}${reset}${red} not found.${reset} Please install package ${blue}${package}${reset}."
+		exit 1
+	fi
 }
 
-# Create a temp directory, chdir into it and create the (initially empty)
-# banlist file.
-tmpdir=$(mktemp -d)
-# Suppress output of dir stack.
-pushd "${tmpdir}" > /dev/null
-touch "${banlist}"
-
 function print_help() {
 	cat <<ENDOFHELP
 Usage: $(basename $0) [OPTION...]
 
-  -a, --auto[=LIMIT]  Enable the autopilot for automatically banning chinese IP
-                      addresses (whois output must contain "source: APNIC" and
-                      "country: CN").
-                      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 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.
 
-  -j, --jail=JAIL     Specify the JAIL to use for banning the IP addresses. If
-                      not set, uses 'apache-auth'.
+  -c, --country=COUNTRY  The country-code to block; defaults to 'CN' (China).
 
-  -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)
+  -j, --jail=JAIL        Specify the JAIL to use for banning the IP addresses.
+                         Defaults to 'apache-auth'.
 
-  -h, --help          Show this help message
+  -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)
+
+  -h, --help             Show this help message
 
 Mandatory or optional arguments to long options are also mandatory or optional
 for any corresponding short options.
@@ -129,8 +123,28 @@ function exec_as_root() {
 	fi
 }
 
+function filter() {
+	# list of current connections
+	file="$1"
+	# subnet extension, e.g. ".0.0/16"
+	ext="$2"
+	rm -f "${filtered}"
+
+	# Reject already banned addresses
+	while read -r -u3 address ; do
+		if [[ "${banned}" != *"${address}${ext}"* ]] ; then
+			echo "Considering ${address}."
+			echo "${address}" >> "${filtered}"
+		else
+			echo "IGNORING ${address}, already banned."
+		fi
+	done 3< "${file}"
+
+	mv "${filtered}" "${file}"
+}
+
 function parse_command_line_args() {
-	TEMP=$(getopt -o 'a::,j:,n:,h' -l 'auto::,jail:,netmask:,help' -- "$@")
+	TEMP=$(getopt -o 'a::,c:,j:,n:,h' -l 'auto::,country:,jail:,netmask:,help' -- "$@")
 
 	if [ $? -ne 0 ] ; then
 		echo 'Error parsing command line options. Terminating. Invoke with --help for help.' >&2
@@ -157,6 +171,10 @@ function parse_command_line_args() {
 				esac
 				shift
 			;;
+			'-c'|'--country')
+				bancountry="$2"
+				shift
+			;;
 			'-j'|'--jail')
 				jail="$2"
 				shift
@@ -199,10 +217,164 @@ function parse_command_line_args() {
 	done
 }
 
+################################################################################
+# Set the highlighting color for the address count. The color goes from green
+# (low count) over yellow (intermediate count) to red (high count). Depending
+# on the size of the subnet (/8, /16, /24 or /32), the transitions from one
+# color to the next happen at different values.
+################################################################################
+function set_highlight_color() {
+	local count=$1
+	case "${choice}" in
+		"1" )
+		# /32: 0 <= green < 3 <= yellow < 5 <= red
+			if [ $count -ge 5 ] ; then
+				hilite="${red}"
+			elif [ $count -ge 3 ] ; then
+				hilite="${yellow}"
+			else
+				hilite="${green}"
+			fi
+		;;
+		"2" )
+		# /24: 0 <= green < 7 <= yellow < 13 <= red
+			if [ $count -ge 13 ] ; then
+				hilite="${red}"
+			elif [ $count -ge 7 ] ; then
+				hilite="${yellow}"
+			else
+				hilite="${green}"
+			fi
+		;;
+		"3" )
+		# /16: 0 <= green < 13 <= yellow < 25 <= red
+			if [ $count -ge 25 ] ; then
+				hilite="${red}"
+			elif [ $count -ge 13 ] ; then
+				hilite="${yellow}"
+			else
+				hilite="${green}"
+			fi
+		;;
+		"4" )
+		# /8: 0 <= green < 21 <= yellow < 49 <= red
+			if [ $count -ge 49 ] ; then
+				hilite="${red}"
+			elif [ $count -ge 21 ] ; then
+				hilite="${yellow}"
+			else
+				hilite="${green}"
+			fi
+		;;
+		* )
+		# ???: We should never get here. As a fall-back, just use no
+		# highlighting.
+			hilite=""
+		;;
+	esac
+}
+
+################################################################################
+# Process the file denoted by $1. For each line in the file, the count and the
+# address are displayed and a whois request is made. The user can then choose to
+# ban or ignore the address. Addresses chosen to be banned are appended to the
+# $banlist.
+################################################################################
+function process_file () {
+	local file="${1}"
+	local line=''
+	local count=0
+	local addr=''
+	local banaction=''
+	local nline=1
+	local country_cn=1
+	# Read the contents from filedescriptor 3 (important: Don's use the
+	# standard filedescriptor because we need to handle user input from
+	# within the loop).
+	while IFS= read -r -u3 line ; do
+		line="$(echo "${line}" | tr -s '[:blank:]')"
+		count="$(echo "${line}" | cut -d' ' -f2)"
+		addr="$(echo "${line}" | cut -d' ' -f3-)${ext}"
+		set_highlight_color "${count}"
+		if [[ autopilot -eq 0 ]] ; then
+			whois "${addr}" | tee "${whoisoutput}"
+		else
+			whois "${addr}" > "${whoisoutput}"
+		fi
+		grep -iq "^country: *${bancountry}$" "${whoisoutput}"
+		country_cn=$?
+		echo -en "Address ${bold}$((nline++)) of ${nlines}${reset}: \
+Found '${blue}${addr}${reset}' ${hilite}${count}${reset} times."
+
+		if [[ ${autopilot} -eq 0 ]] ; then
+			echo -en  "Ban [y/N/s=No, and skip remaining]? "
+			read banaction
+		else
+			if [[ ${country_cn} -eq 0 ]] ; then
+				if [[ $count -ge $autopilot ]] ; then
+					echo -en "\n${red}Autopilot active. ${reset}"
+					banaction=y
+				else
+					echo -e "\n${yellow}Autopilot active. Ignoring remaining addresses due to limit of ${autopilot}.${reset}"
+					return
+				fi
+			else
+				if [[ $count -ge $autopilot ]] ; then
+					echo -en "\n${green}Autopilot active. ${reset}"
+					banaction=n
+				else
+					echo -e "\n${green}Autopilot active.${reset} ${yellow}Ignoring remaining addresses due to limit of ${autopilot}.${reset}"
+					return
+				fi
+			fi
+		fi
+
+		case "${banaction}" in
+			"s" | "S" )
+				echo -e "Not banning '${blue}${addr}${reset}', \
+skipping remaining addresses."
+				return
+			;;
+			"y" | "Y" )
+				echo -e "Adding '${blue}${addr}${reset}' to \
+banlist."
+				echo "${addr}" >> "${banlist}"
+			;;
+			"n" | "N" | * )
+				echo -e "Not banning '${blue}${addr}${reset}'."
+			;;
+		esac
+	# Here goes: Pipe the file contents via filedescriptor 3.
+	done 3< "${file}"
+	echo "Processed all entries in ${file}."
+}
+
+# Create a temp directory, chdir into it and create the (initially empty)
+# banlist file.
+tmpdir=$(mktemp -d)
+# Suppress output of dir stack.
+pushd "${tmpdir}" > /dev/null
+touch "${banlist}"
+
+check_installed "sudo" "app-admin/sudo"
+check_installed "fail2ban-client" "net-analyzer/fail2ban"
+check_installed "whois" "net-misc/whois"
+check_installed "cut" "sys-apps/coreutils"
+check_installed "id" "sys-apps/coreutils"
+check_installed "sort" "sys-apps/coreutils"
+check_installed "touch" "sys-apps/coreutils"
+check_installed "tr" "sys-apps/coreutils"
+check_installed "uniq" "sys-apps/coreutils"
+check_installed "grep" "sys-apps/grep"
+check_installed "sponge" "sys-apps/moreutils"
+check_installed "netstat" "sys-apps/net_tools"
+check_installed "getopt" "sys-apps/util-linux"
+
 # Parse the command line options
 autopilot=0
 netmask=0
 jail="apache-auth"
+bancountry="CN"
 
 parse_command_line_args "$@"
 
@@ -220,26 +392,6 @@ cut -d. -f1-3 "${fileraw}" | sort > "${file24}"
 cut -d. -f1-2 "${fileraw}" | sort > "${file16}"
 cut -d. -f1 "${fileraw}" | sort > "${file8}"
 
-function filter() {
-	# list of current connections
-	file="$1"
-	# subnet extension, e.g. ".0.0/16"
-	ext="$2"
-	rm -f "${filtered}"
-
-	# Reject already banned addresses
-	while read -r -u3 address ; do
-		if [[ "${banned}" != *"${address}${ext}"* ]] ; then
-			echo "Considering ${address}."
-			echo "${address}" >> "${filtered}"
-		else
-			echo "IGNORING ${address}, already banned."
-		fi
-	done 3< "${file}"
-
-	mv "${filtered}" "${file}"
-}
-
 # Filter already banned addresses
 filter "${file32}" "${ext32}"
 filter "${file24}" "${ext24}"
@@ -305,143 +457,8 @@ unset TEMP
 
 echo "Processing ${file}."
 
-################################################################################
-# Set the highlighting color for the address count. The color goes from green
-# (low count) over yellow (intermediate count) to red (high count). Depending
-# on the size of the subnet (/8, /16, /24 or /32), the transitions from one
-# color to the next happen at different values.
-################################################################################
-function setHilite() {
-	local count=$1
-	case "${choice}" in
-		"1" )
-		# /32: 0 <= green < 3 <= yellow < 5 <= red
-			if [ $count -ge 5 ] ; then
-				hilite="${red}"
-			elif [ $count -ge 3 ] ; then
-				hilite="${yellow}"
-			else
-				hilite="${green}"
-			fi
-		;;
-		"2" )
-		# /24: 0 <= green < 7 <= yellow < 13 <= red
-			if [ $count -ge 13 ] ; then
-				hilite="${red}"
-			elif [ $count -ge 7 ] ; then
-				hilite="${yellow}"
-			else
-				hilite="${green}"
-			fi
-		;;
-		"3" )
-		# /16: 0 <= green < 13 <= yellow < 25 <= red
-			if [ $count -ge 25 ] ; then
-				hilite="${red}"
-			elif [ $count -ge 13 ] ; then
-				hilite="${yellow}"
-			else
-				hilite="${green}"
-			fi
-		;;
-		"4" )
-		# /8: 0 <= green < 21 <= yellow < 49 <= red
-			if [ $count -ge 49 ] ; then
-				hilite="${red}"
-			elif [ $count -ge 21 ] ; then
-				hilite="${yellow}"
-			else
-				hilite="${green}"
-			fi
-		;;
-		* )
-		# ???: We should never get here. As a fall-back, just use no
-		# highlighting.
-			hilite=""
-		;;
-	esac
-}
-
-################################################################################
-# Process the file denoted by $1. For each line in the file, the count and the
-# address are displayed and a whois request is made. The user can then choose to
-# ban or ignore the address. Addresses chosen to be banned are appended to the
-# $banlist.
-################################################################################
-function processFile () {
-	local file="${1}"
-	local line=''
-	local count=0
-	local addr=''
-	local banaction=''
-	local nline=1
-	local country_cn=1
-	local source_apnic=1
-	# Read the contents from filedescriptor 3 (important: Don's use the
-	# standard filedescriptor because we need to handle user input from
-	# within the loop).
-	while IFS= read -r -u3 line ; do
-		line="$(echo "${line}" | tr -s '[:blank:]')"
-		count="$(echo "${line}" | cut -d' ' -f2)"
-		addr="$(echo "${line}" | cut -d' ' -f3-)${ext}"
-		setHilite "${count}"
-		if [[ autopilot -eq 0 ]] ; then
-			whois "${addr}" | tee "${whoisoutput}"
-		else
-			whois "${addr}" > "${whoisoutput}"
-		fi
-		grep -iq "^country: *cn$" "${whoisoutput}"
-		country_cn=$?
-		grep -iq "^source: *apnic$" "${whoisoutput}"
-		source_apnic=$?
-		echo -en "Address ${bold}$((nline++)) of ${nlines}${reset}: \
-Found '${blue}${addr}${reset}' ${hilite}${count}${reset} times."
-
-		if [[ ${autopilot} -eq 0 ]] ; then
-			echo -en  "Ban [y/N/s=No, and skip remaining]? "
-			read banaction
-		else
-			if [[ ${country_cn} -eq 0 && ${source_apnic} -eq 0 ]] ; then
-				if [[ $count -ge $autopilot ]] ; then
-					echo -en "\n${red}Autopilot active. ${reset}"
-					banaction=y
-				else
-					echo -e "\n${yellow}Autopilot active. Ignoring remaining addresses due to limit of ${autopilot}.${reset}"
-					return
-				fi
-			else
-				if [[ $count -ge $autopilot ]] ; then
-					echo -en "\n${green}Autopilot active. ${reset}"
-					banaction=n
-				else
-					echo -e "\n${green}Autopilot active.${reset} ${yellow}Ignoring remaining addresses due to limit of ${autopilot}.${reset}"
-					return
-				fi
-			fi
-		fi
-
-		case "${banaction}" in
-			"s" | "S" )
-				echo -e "Not banning '${blue}${addr}${reset}', \
-skipping remaining addresses."
-				return
-			;;
-			"y" | "Y" )
-				echo -e "Adding '${blue}${addr}${reset}' to \
-banlist."
-				echo "${addr}" >> "${banlist}"
-			;;
-			"n" | "N" | * )
-				echo -e "Not banning '${blue}${addr}${reset}'."
-			;;
-		esac
-	# Here goes: Pipe the file contents via filedescriptor 3.
-	done 3< "${file}"
-	echo "Processed all entries in ${file}."
-}
-
 # Invoke the processing function on the chosen file.
-processFile "${file}"
+process_file "${file}"
 
 echo "These are the addresses to be banned:"
 cat "${banlist}"