#!/bin/sh
################################################################################
#                                                                              #
# Try and prevent apache overloads by banning IP addresses that have (too)     #
# many open connections.                                                       #
# This script uses ss to determine the connections to a configurable port      #
# on the host machine and provides automated GeoIP information retrieval based #
# the address or the /24-, /16- or /8-subnet thereof. A GeoIP city- or country #
# database must be installed separately and is provided to the script via a    #
# command line parameter.                                                      #
# Addresses (or subnets) are presented to the user in order of descending      #
# connection count. For each address (or subnet), the user can choose to ban   #
# or ignore it. Addresses (or subnets) chosen to be banned will be blocked by  #
# a configurable jail of fail2ban.                                             #
# Author: Manuel Friedli, <manuel@fritteli.ch>                                 #
# This script is licenced under the GNU General Public Licence, version 3 or   #
# later.                                                                       #
#                                                                              #
################################################################################

################################################################################
#                                                                              #
# Prerequisites:                                                               #
#   - app-admin/sudo (`sudo`)                                                  #
#   - dev-lang/python:3.8 (`python`)                                           #
#   - net-analyzer/fail2ban (`fail2ban-client`)                                #
#   - sys-apps/coreutils (`cut`, `id`, `sort`, `touch`, `tr`, `uniq`)          #
#   - sys-apps/grep (`grep`)                                                   #
#   - sys-apps/iproute2 (`ss`)                                                 #
#   - sys-apps/moreutils (`sponge`)                                            #
#   - sys-apps/util-linux (`getopt`)                                           #
#                                                                              #
################################################################################

# Store the start time; this enables us to output the total runtime at the end.
start="$(date +%s)"

# Dependencies of this script. Simple array with the following structure:
# (command package [...])
dependencies=(
	"sudo" "app-admin/sudo"
	"python" "dev-lang/python:3.8"
	"fail2ban-client" "net-analyzer/fail2ban"
	"cut" "sys-apps/coreutils"
	"id" "sys-apps/coreutils"
	"sort" "sys-apps/coreutils"
	"touch" "sys-apps/coreutils"
	"tr" "sys-apps/coreutils"
	"uniq" "sys-apps/coreutils"
	"grep" "sys-apps/grep"
	"ss" "sys-apps/iproute2"
	"sponge" "sys-apps/moreutils"
	"getopt" "sys-apps/util-linux"
)

# These suffixes must be appended to the respective addresses and subnets.
suffix8="/8"
suffix16="/16"
suffix24="/24"
suffix32="/32"
ext8=".0.0.0"
ext16=".0.0"
ext24=".0"
ext32=""

# Define some constants to format the output in a colorful way.
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 2>/dev/null >&2; rm -r ${tmpdir}' EXIT

function is_installed() {
	which "${1}" 2>/dev/null >&2
	return $?
}

function print_missing_dependency() {
	local command="${1}"
	local package="${2}"

	echo "${red}Command ${bold}${command}${reset}${red} not found.${reset} Please install package ${blue}${package}${reset}." >&2
}

function check_dependencies() {
	local arraylength=${#dependencies[@]}
	local res=
	local command=
	local package=
	# 0: true, all installed; 1: false, at least one command/package missing
	local all_installed=0

	for ((i = 0; i < ${arraylength}; i += 2)); do
		command="${dependencies[$i]}"
		package="${dependencies[$i + 1]}"
		is_installed "${command}" "${package}"
		res=$?
		if [[ ${res} -ne 0 ]]; then
			print_missing_dependency "${command}" "${package}"
			all_installed=1
		fi
	done

	return ${all_installed}
}

function print_help() {
	cat <<ENDOFHELP
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
                                      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[,COUNTRY...]  The country-codes to block as a list of
                                      comma-separated values; defaults to 'CN'
                                      (China).

  -e, --dependencies                  Check if all required dependencies are
                                      installed. If all dependencies are found,
                                      the program terminates with exit code 0.
                                      Otherwise, missing dependencies are
                                      printed to stderr and the program
                                      terminates with exit code 1.

  -f, --config-file=FILENAME          Specify the full path to the configuration
                                      file. If omitted, all options are read
                                      from the command line.
                                      Any parameter specified on the command
                                      line takes precedence over configuration
                                      option read from the file.

  -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)

  -p, --port=PORT                     The desired port to monitor.
                                      Defaults to 443 (https).

  -h, --help                          Show this help message

Mandatory or optional arguments to long options are also mandatory or optional
for any corresponding short options.

When invoked without optional options, the autopilot is disabled and the
netmask SIZE is inquired interactively.
ENDOFHELP
}

function exec_as_root() {
	if [[ $(id -un) == "root" ]]; then
		"$@"
	else
		sudo "$@"
	fi
}

function filter() {
	# list of current connections
	file="${1}"
	# subnet extension, e.g. ".0.0"
	ext="${2}"
	# subnet suffix, e.g. "/16"
	suffix="${3}"
	rm -f "${filtered}"
	touch "${filtered}"

	# Reject already banned addresses
	while read -r -u3 address; do
		if [[ "${banned}" != *"${address}${ext}${suffix}"* ]]; then
			echo "${address}" >>"${filtered}"
		else
			echo "IGNORING ${address}${ext}${suffix}, already banned."
		fi
	done 3<"${file}"

	mv "${filtered}" "${file}"
}

function set_default_values() {
	if [[ -z "${autopilot}" ]]; then
		autopilot=0
	fi
	if [[ -z "${bancountries}" ]]; then
		bancountries=("CN")
	fi
	if [[ -z "${jail}" ]]; then
		jail="apache-auth"
	fi
	if [[ -z "${netmask}" ]]; then
		netmask=0
	fi
	if [[ -z "${port}" ]]; then
		port=443
	fi
}

function parse_config_file() {
	source "${configfile}"
	if [[ -z "${autopilot}" ]]; then
		autopilot="${AUTOPILOT}"
	fi
	if [[ -z "${bancountries}" ]]; then
		bancountries=(${COUNTRIES[@]})
	fi
	if [[ -z "${database}" ]]; then
		database="${DATABASE_FILE}"
	fi
	if [[ -z "${jail}" ]]; then
		jail="${JAIL}"
	fi
	if [[ -z "${netmask}" ]]; then
		netmask="${NETMASK}"
	fi
	if [[ -z "${port}" ]]; then
		port="${PORT}"
	fi
}

function parse_command_line_args() {
	TEMP=$(getopt -o 'a::,c:,d:,e,f:,j:,n:,p:,h' -l 'auto::,country:,database:,dependencies,config-file:,jail:,netmask:,port:,help' -- "$@")

	if [ $? -ne 0 ]; then
		echo 'Error parsing command line options. Terminating. Invoke with --help for help.' >&2
		exit 1
	fi

	eval set -- "${TEMP}"
	unset TEMP

	while true; do
		case "${1}" in
		'-a' | '--auto')
			case "${2}" in
			'')
				autopilot=1
				;;
			*)
				autopilot="$2"
				;;
			esac
			shift
			;;
		'-c' | '--country')
			IFS=',' read -ra bancountries <<<"${2}"
			shift
			;;
		'-d' | '--database')
			database="${2}"
			shift
			;;
		'-e' | '--dependencies')
			check_dependencies
			exit $?
			;;
		'-f' | '--config-file')
			configfile="${2}"
			shift
			;;
		'-j' | '--jail')
			jail="${2}"
			shift
			;;
		'-n' | '--netmask')
			case "${2}" in
			'1')
				netmask=8
				;;
			'2')
				netmask=16
				;;
			'3')
				netmask=24
				;;
			'4')
				netmask=32
				;;
			*)
				netmask="${2}"
				;;
			esac
			shift
			;;
		'-p' | '--port')
			port="${2}"
			shift
			;;
		'-h' | '--help')
			print_help
			exit
			;;
		'--')
			shift
			break
			;;
		*)
			echo "Unknown error on command line argument '${1}'. Terminating." >&2
			exit 1
			;;
		esac
		shift
	done

	# If the config file option is set, parse the config file.
	if [[ ! -z ${configfile+x} ]]; then
		if [[ ! -f "${configfile}" || ! -r "${configfile}" ]]; then
			echo "Can not read configuration file '${2}'. Invoke with --help for help." >&2
			exit 1
		fi

		parse_config_file
	fi

	# Here, we set the default values for all options that have not been set yet.
	set_default_values
}

function validate_parameter_values() {
	# GeoIP-Database
	if [[ -z "${database}" ]]; then
		echo "No GeoIP database specified. Invoke with --help for more information." >&2
		exit 1
	fi

	if [[ ! -f "${database}" || ! -r "${database}" ]]; then
		echo "Database '${database}' is not accessible." >&2
		exit 1
	fi

	# Autopilot
	case "${autopilot}" in
	*[!0-9]*)
		echo "Invalid value for parameter 'auto' / 'AUTOPILOT': '${autopilot}'." >&2
		echo "Invoke with --help for help." >&2
		exit 1
		;;
	esac

	# Countries
	if [[ -z ${bancountries[@]// /} ]]; then
		echo "Invalid value for parameter 'country' / 'COUNTRIES': '${bancountries[*]}'." >&2
		echo "Invoke with --help for help." >&2
		exit 1
	fi

	# Jail
	if [[ -z "${jail}" ]]; then
		echo "Invalid value for parameter 'jail' / 'JAIL': '${jail}'." >&2
		echo "Invoke with --help for help." >&2
		exit 1
	fi

	# Netmask
	case "${netmask}" in
	'0' | '8' | '16' | '24' | '32')
		# Everything OK.
		;;
	*)
		echo "Invalid value for parameter 'netmask': '${2}'." >&2
		echo "Invoke with --help for help." >&2
		exit 1
		;;
	esac

	# Port
	case "${port}" in
	*[!0-9]*)
		echo "Invalid value for parameter 'port' / 'PORT': '${autopilot}'." >&2
		echo "Invoke with --help for help." >&2
		exit 1
		;;
	esac
	if [[ ${port} -lt 0 || ${port} -gt 65535 ]]; then
		echo "Invalid value for parameter 'port' / 'PORT': '${autopilot}'." >&2
		echo "Value must be between 0 ... 65535 (inclusive)." >&2
		echo "Invoke with --help for help." >&2
		exit 1
	fi
}

################################################################################
# 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 lookup for the IP addresses country is made in the
# GeoIP database. 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 addronly=''
	local addrwithsuffix=''
	local banaction=''
	local nline=1
	local country=
	# 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)"
		addronly="$(echo "${line}" | cut -d' ' -f3-)${ext}"
		addrwithsuffix="${addronly}${suffix}"
		set_highlight_color "${count}"
		country="$("${curdir}/geoip-lookup.py" -f "${database}" "${addronly}")"
		if [[ ${autopilot} -eq 0 ]]; then
			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]? "
			read banaction
		else
			if [[ " ${bancountries[@]} " =~ " ${country} " ]]; 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 "Not banning '${blue}${addrwithsuffix}${reset}', \
skipping remaining addresses."
			return
			;;
		"y" | "Y")
			echo "Adding '${blue}${addrwithsuffix}${reset}' to \
banlist (country=${yellow}${country}${reset})."
			echo "${addrwithsuffix}" >>"${banlist}"
			;;
		"n" | "N" | *)
			echo "Not banning '${blue}${addrwithsuffix}${reset}' (country=${yellow}${country}${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)"

# Set up all file paths
curdir="$(dirname "$0")"
# 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"

touch "${banlist}"

# Parse the command line options
autopilot=
netmask=
jail=
bancountries=
database=
port=

parse_command_line_args "$@"
validate_parameter_values

check_dependencies
dependencies_ok=$?
if [[ ${dependencies_ok} -ne 0 ]]; then
	exit ${dependencies_ok}
fi

# List already banned addresses in the chosen jail
banned="$(exec_as_root fail2ban-client get "${jail}" banip)"

# Determine the current connections to the desired port; store the raw data in
# $fileraw.
connections=$(ss -HOn state established "( sport = :${port} )" | tr -s '[:blank:]' | cut -d' ' -f5)

# IPv6-mapped-IPv4: [::ffff:192.168.0.1]:443
echo "${connections}" | grep '^\[::ffff:' - | cut -d: -f4 | cut -d] -f1 | grep -v '^$' >"${fileraw}"
# Pure IPv4: 192.168.0.1:443
echo "${connections}" | grep -v '^\[' - | cut -d: -f1 | grep -v '^$' >>"${fileraw}"

# Group and sort the data into the subnet-specific files.
sort "${fileraw}" >"${file32}"
cut -d. -f1-3 "${fileraw}" | sort >"${file24}"
cut -d. -f1-2 "${fileraw}" | sort >"${file16}"
cut -d. -f1 "${fileraw}" | sort >"${file8}"

# Filter already banned addresses
filter "${file32}" "${ext32}" "${suffix32}"
filter "${file24}" "${ext24}" "${suffix24}"
filter "${file16}" "${ext16}" "${suffix16}"
filter "${file8}" "${ext8}" "${suffix8}"

# Determine the number of connections per address
uniq -c "${file32}" | sort -rn | sponge "${file32}"
uniq -c "${file24}" | sort -rn | sponge "${file24}"
uniq -c "${file16}" | sort -rn | sponge "${file16}"
uniq -c "${file8}" | sort -rn | sponge "${file8}"

# Determine the number of entries per file.
nlines32=$(cat "${file32}" | wc -l)
nlines24=$(cat "${file24}" | wc -l)
nlines16=$(cat "${file16}" | wc -l)
nlines8=$(cat "${file8}" | wc -l)

if [ ${netmask} -eq 0 ]; then
	# Now let the user choose which file to process.
	echo "We've got:"
	echo "[1] 8bit: ${nlines8} entries"
	echo "[2] 16bit: ${nlines16} entries"
	echo "[3] 24bit: ${nlines24} entries"
	echo "[4] 32bit: ${nlines32} entries"
	read -p 'Which one do you want to work with (q=Quit) [1-4]? ' choice

	# Based on the user's choice, initialize the variables $file, $ext and
	# $nlines, which will be used after this point. Also, $choice will be
	# used to color the output based on subnet-type.
	case "${choice}" in
	"1")
		netmask=8
		;;
	"2")
		netmask=16
		;;
	"3")
		netmask=24
		;;
	"4")
		netmask=32
		;;
	"Q" | "q")
		echo "You chose to abort. That's fine! Have a nice day!"
		exit
		;;
	*)
		echo "Invalid input: ${choice}. I'm out of here."
		exit 1
		;;
	esac
fi

# Now initialize the variables $file, $ext and $nlines based on the chosen $netmask
TEMP="file${netmask}"
file="${!TEMP}"
TEMP="ext${netmask}"
ext="${!TEMP}"
TEMP="suffix${netmask}"
suffix="${!TEMP}"
TEMP="nlines${netmask}"
nlines="${!TEMP}"
unset TEMP

echo "Processing ${file}."

# Invoke the processing function on the chosen file.
process_file "${file}"

echo "These are the addresses to be banned:"
cat "${banlist}"

# Make sure the user has to (re)-identify him- or herself before actually
# banning anything.
sudo -k

# Iterate over all addresses in $banlist and invoke fail2ban-client on each
# one of them.
while read -r addrwithsuffix; do
	echo "Banning ${addrwithsuffix} ..."
	exec_as_root fail2ban-client set "${jail}" banip "${addrwithsuffix}"
done <"${banlist}"

end="$(date +%s)"

echo "${green}All done in $((end - start)) seconds!${reset}"