From 30e53268f55edbf2c189d8760aff036124d0bd2a Mon Sep 17 00:00:00 2001
From: Manuel Friedli <manuel@fritteli.ch>
Date: Tue, 4 Aug 2020 03:49:21 +0200
Subject: [PATCH] Prepare reading the database path from the command line.

---
 ddos-mitigator.sh | 100 +++++++++++++++++++++++++++-------------------
 1 file changed, 60 insertions(+), 40 deletions(-)

diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh
index d0eb2ce..c764708 100755
--- a/ddos-mitigator.sh
+++ b/ddos-mitigator.sh
@@ -42,18 +42,6 @@ MY_IP="94.199.214.20"
 MY_PORT="443"
 
 # After this point, no editing is required.
-# Define the files that will contain the addresses an subnets.
-fileraw="raw-http.txt"
-filtered="filtered-http.txt"
-file8="sorted-http-8.txt"
-file16="sorted-http-16.txt"
-file24="sorted-http-24.txt"
-file32="sorted-http-32.txt"
-# This file will contain the addresses to be banned.
-banlist="banlist.txt"
-# This file contains the output of the last invocation of whois
-whoisoutput="whois.txt"
-
 # These suffixes must be appended to the respective addresses and subnets.
 ext8=".0.0.0/8"
 ext16=".0.0/16"
@@ -61,15 +49,15 @@ ext24=".0/24"
 ext32="/32"
 
 # Define some constants to format the output in a colorful way.
-red="\033[38;2;255;0;43m"
-yellow="\033[38;2;255;204;0m"
-green="\033[38;2;0;179;89m"
-blue="\033[38;2;0;85;255m"
-bold="\033[1m"
-reset="\033[0m"
+red="$(printf '\033[38;2;255;0;43m')"
+yellow="$(printf '\033[38;2;255;204;0m')"
+green="$(printf '\033[38;2;0;179;89m')"
+blue="$(printf '\033[38;2;0;85;255m')"
+bold="$(printf '\033[1m')"
+reset="$(printf '\033[0m')"
 
 # Clean up when the script exits.
-trap 'sudo -k; popd >/dev/null; rm -r ${tmpdir}' EXIT
+trap 'sudo -k; rm -r ${tmpdir}' EXIT
 
 function check_installed() {
 	local command="$1"
@@ -78,15 +66,20 @@ function check_installed() {
 	local result=$?
 
 	if [[ "${result}" -ne 0 ]] ; then
-		echo -e "${red}Command ${bold}${command}${reset}${red} not found.${reset} Please install package ${blue}${package}${reset}."
+		echo "${red}Command ${bold}${command}${reset}${red} not found.${reset} Please install package ${blue}${package}${reset}."
 		exit 1
 	fi
 }
 
 function print_help() {
 	cat <<ENDOFHELP
-Usage: $(basename $0) [OPTION...]
+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).
+
+  ${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
@@ -110,8 +103,8 @@ Usage: $(basename $0) [OPTION...]
 Mandatory or optional arguments to long options are also mandatory or optional
 for any corresponding short options.
 
-When invoked without options, the autopilot is disabled and the netmask SIZE is
-inquired interactively.
+When invoked without optional options, the autopilot is disabled and the
+netmask SIZE is inquired interactively.
 ENDOFHELP
 }
 
@@ -144,7 +137,7 @@ function filter() {
 }
 
 function parse_command_line_args() {
-	TEMP=$(getopt -o 'a::,c:,j:,n:,h' -l 'auto::,country:,jail:,netmask:,help' -- "$@")
+	TEMP=$(getopt -o 'a::,c:,d:,j:,n:,h' -l 'auto::,country:,database:,jail:,netmask:,help' -- "$@")
 
 	if [ $? -ne 0 ] ; then
 		echo 'Error parsing command line options. Terminating. Invoke with --help for help.' >&2
@@ -175,6 +168,10 @@ function parse_command_line_args() {
 				bancountry="$2"
 				shift
 			;;
+			'-d'|'--database')
+				database="$2"
+				shift
+			;;
 			'-j'|'--jail')
 				jail="$2"
 				shift
@@ -215,6 +212,16 @@ function parse_command_line_args() {
 		esac
 		shift
 	done
+
+	if [[ -z "${database}" ]] ; then
+	  echo "No GeoIP database specified. Invoke with --help for more information." >&2
+		exit 1
+	fi
+
+	if [[ ! -r "${database}" ]] ; then
+	  echo "Database '${database}' is not accessible." >&2
+		exit 1
+	fi
 }
 
 ################################################################################
@@ -303,11 +310,11 @@ function process_file () {
 		fi
 		grep -iq "^country: *${bancountry}$" "${whoisoutput}"
 		country_cn=$?
-		echo -en "Address ${bold}$((nline++)) of ${nlines}${reset}: \
+		echo -n "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]? "
+			echo -n  "Ban [y/N/s=No, and skip remaining]? "
 			read banaction
 		else
 			if [[ ${country_cn} -eq 0 ]] ; then
@@ -331,17 +338,17 @@ Found '${blue}${addr}${reset}' ${hilite}${count}${reset} times."
 
 		case "${banaction}" in
 			"s" | "S" )
-				echo -e "Not banning '${blue}${addr}${reset}', \
+				echo "Not banning '${blue}${addr}${reset}', \
 skipping remaining addresses."
 				return
 			;;
 			"y" | "Y" )
-				echo -e "Adding '${blue}${addr}${reset}' to \
+				echo "Adding '${blue}${addr}${reset}' to \
 banlist."
 				echo "${addr}" >> "${banlist}"
 			;;
 			"n" | "N" | * )
-				echo -e "Not banning '${blue}${addr}${reset}'."
+				echo "Not banning '${blue}${addr}${reset}'."
 			;;
 		esac
 	# Here goes: Pipe the file contents via filedescriptor 3.
@@ -352,10 +359,31 @@ banlist."
 # 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
+
+# Set up all file paths
+# Define the files that will contain the addresses an subnets.
+fileraw="${tmpdir}/raw-http.txt"
+filtered="${tmpdir}/filtered-http.txt"
+file8="${tmpdir}/sorted-http-8.txt"
+file16="${tmpdir}/sorted-http-16.txt"
+file24="${tmpdir}/sorted-http-24.txt"
+file32="${tmpdir}/sorted-http-32.txt"
+# This file will contain the addresses to be banned.
+banlist="${tmpdir}/banlist.txt"
+# This file contains the output of the last invocation of whois
+whoisoutput="${tmpdir}/whois.txt"
+
 touch "${banlist}"
 
+# Parse the command line options
+autopilot=0
+netmask=0
+jail="apache-auth"
+bancountry="CN"
+database=
+
+parse_command_line_args "$@"
+
 check_installed "sudo" "app-admin/sudo"
 check_installed "fail2ban-client" "net-analyzer/fail2ban"
 check_installed "whois" "net-misc/whois"
@@ -370,14 +398,6 @@ 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 "$@"
-
 # List already banned addresses in the chosen jail
 banned="$(exec_as_root fail2ban-client get "${jail}" banip)"
 
@@ -474,4 +494,4 @@ while read -r addr ; do
 	exec_as_root fail2ban-client set "${jail}" banip "${addr}"
 done < "${banlist}"
 
-echo -e "${green}All done!${reset}"
+echo "${green}All done!${reset}"