From d1e0115800ed26d687168d4149b7e7a9230e63cf Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Mon, 10 Aug 2020 00:33:01 +0200 Subject: [PATCH 01/10] Try and support blocking IPv6 addresses. Doesn't work if you choose 8, 16 or 24bit IPv4 base blocking. --- ddos-mitigator.sh | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index 8270584..b9ef911 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -61,10 +61,12 @@ suffix8="/8" suffix16="/16" suffix24="/24" suffix32="/32" +suffixipv6="" ext8=".0.0.0" ext16=".0.0" ext24=".0" ext32="" +extipv6="" # Define some constants to format the output in a colorful way. red="$(printf '\033[38;2;255;0;43m')" @@ -150,6 +152,12 @@ Usage: $(basename $0) -d FILE [OPTION...] - 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) + - 5 or 128 for IPv6 addresses (no IPv4 + addresses are considered) + + -6, --enable-ipv6 Enable banning of IPv6 addresses in + addition to the IPv4 addresses/networks + defined by -n (if SIZE is 1, 2, 3 or 4) -p, --port=PORT The desired port to monitor. Defaults to 443 (https). @@ -195,7 +203,7 @@ function filter() { } function parse_command_line_args() { - TEMP=$(getopt -o 'a::,c:,d:,e,j:,n:,p:,h' -l 'auto::,country:,database:,dependencies,jail:,netmask:,port:,help' -- "$@") + TEMP=$(getopt -o 'a::,c:,d:,e,j:,n:,p:,6,h' -l 'auto::,country:,database:,dependencies,jail:,netmask:,port:,enable-ipv6,help' -- "$@") if [ $? -ne 0 ]; then echo 'Error parsing command line options. Terminating. Invoke with --help for help.' >&2 @@ -256,6 +264,9 @@ function parse_command_line_args() { '4' | '32') netmask=32 ;; + '5'|'128') + netmask="ipv6" + ;; *) echo "Invalid argument for parameter 'netmask': '${2}'. Invoke with --help for help." >&2 exit 1 @@ -267,6 +278,9 @@ function parse_command_line_args() { port="${2}" shift ;; + '-6'|'--enable-ipv6') + enableipv6=1 + ;; '-h' | '--help') print_help exit @@ -366,6 +380,7 @@ function process_file() { 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). @@ -438,6 +453,7 @@ file8="${tmpdir}/sorted-http-8.txt" file16="${tmpdir}/sorted-http-16.txt" file24="${tmpdir}/sorted-http-24.txt" file32="${tmpdir}/sorted-http-32.txt" +fileipv6="${tmpdir}/sorted-http-ipv6.txt" # This file will contain the addresses to be banned. banlist="${tmpdir}/banlist.txt" @@ -445,6 +461,7 @@ touch "${banlist}" # Parse the command line options autopilot=0 +enableipv6=0 netmask=0 jail="apache-auth" bancountries=("CN") @@ -470,6 +487,8 @@ connections=$(ss -HOn state established "( sport = :${port} )" | tr -s '[:blank: 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}" +# IPv6: [aaaa:...]:443 +echo "${connections}" | grep '^\[[^:]' - | cut -d']' -f1 | sed -e 's/^\[//' | grep -v '^$' | sort > "${fileipv6}" # Group and sort the data into the subnet-specific files. sort "${fileraw}" >"${file32}" @@ -482,18 +501,21 @@ filter "${file32}" "${ext32}" "${suffix32}" filter "${file24}" "${ext24}" "${suffix24}" filter "${file16}" "${ext16}" "${suffix16}" filter "${file8}" "${ext8}" "${suffix8}" +filter "${fileipv6}" "" "" # 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}" +uniq -c "${fileipv6}" | sort -rn | sponge "${fileipv6}" # 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) +nlinesipv6=$(cat "${fileipv6}" | wc -l) if [ ${netmask} -eq 0 ]; then # Now let the user choose which file to process. @@ -502,6 +524,7 @@ if [ ${netmask} -eq 0 ]; then echo "[2] 16bit: ${nlines16} entries" echo "[3] 24bit: ${nlines24} entries" echo "[4] 32bit: ${nlines32} entries" + echo "[5] IPv6: ${nlinesipv6} 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 @@ -520,6 +543,9 @@ if [ ${netmask} -eq 0 ]; then "4") netmask=32 ;; + "5" ) + netmask="ipv6" + ;; "Q" | "q") echo "You chose to abort. That's fine! Have a nice day!" exit @@ -544,6 +570,15 @@ unset TEMP echo "Processing ${file}." +# Handle IPv6 option +if [[ "${enableipv6}" -eq 1 ]]; then + if [[ "${netmask}" != "ipv6" ]]; then + cat "${file}" "${fileipv6}" | sort -n | sponge "${file}" + nlines=$(( ${nlines} + ${nlinesipv6} )) + fi +fi + + # Invoke the processing function on the chosen file. process_file "${file}" From 0871a25bf76aca1ff72a8defb8f9532afd935d78 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 21:28:15 +0100 Subject: [PATCH 02/10] Add the possibility of using a config file instead of CLI parameters. --- ddos-mitigator.conf | 30 +++++++++++++++++++++++++++ ddos-mitigator.sh | 49 ++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 74 insertions(+), 5 deletions(-) create mode 100644 ddos-mitigator.conf diff --git a/ddos-mitigator.conf b/ddos-mitigator.conf new file mode 100644 index 0000000..a7ebc54 --- /dev/null +++ b/ddos-mitigator.conf @@ -0,0 +1,30 @@ +# Example configuration file for ddos-mitigator.sh. +# PLEASE TAKE CARE not to put any whitespace around the '=' signs, as this file is directly sourced by the script +# file and this needs to conform to the BASH syntax. Also, make sure to declare the COUNTRIES variable with the +# correct array syntax: COUNTRIES=("AA" "BB" "CC"), or to comment it out altogether. + +# The path to the GeoIP2 database file (must be either country or city database). This parameter is mandatory. If it is +# not specified here, it must be given on the command line (through the -d option). +DATABASE_FILE="/path/to/geoip/country-or-city-database.mmdb" + +# Enable the autopilot for automatically banning IP addresses of the desired countries (see also COUNTRIES option). +# Only ban IP addresses with at least AUTOPILOT current connections. If the value is not specified or 0, don't +# automatically ban IP addresses, but run in interactive mode. +AUTOPILOT="1" + +# Defines the subnet size in bytes to be analyzed. Valid values are: +# - 8 for class A networks (X.0.0.0/8) +# - 16 for class B networks (X.X.0.0/16) +# - 24 for class C networks (X.X.X.0/24) +# - 32 for class D networks (X.X.X.X/32) +# If not specified, run in interactive mode and prompt for the netmask size. +NETMASK="8" + +# The country-codes to block as an array. Defaults to "CN" (China). +#COUNTRIES=("CN" "HK" "TW") + +# Specify the JAIL to use for banning the IP addresses. Defaults to 'apache-auth'. +#JAIL="apache-auth" + +# The desired port to monitor. Defaults to 443 (https). +#PORT="443" diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index 8270584..01689ea 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -32,10 +32,7 @@ # # ################################################################################ -# Set the host's own IP address. So far, only an IPv4 address is supported. -MY_IP="94.199.214.20" - -# After this point, no editing is required. +# 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: @@ -139,6 +136,13 @@ Usage: $(basename $0) -d FILE [OPTION...] 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'. @@ -194,8 +198,30 @@ function filter() { mv "${filtered}" "${file}" } +function parse_config_file() { + source "${configfile}" + if [[ -z "${autopilot+x}" ]]; then + autopilot="${AUTOPILOT}" + fi + if [[ -z "${bancountries}" ]]; then + bancountries=()${COUNTRIES[@]}) + fi + if [[ -z "${database+x}" ]]; then + database="${DATABASE_FILE}" + fi + if [[ -z "${jail+x}" ]]; then + jail="${JAIL}" + fi + if [[ -z "${netmask+x}" ]]; then + netmask="${NETMASK}" + fi + if [[ -z "${port+x}" ]]; then + port="${PORT}" + fi +} + function parse_command_line_args() { - TEMP=$(getopt -o 'a::,c:,d:,e,j:,n:,p:,h' -l 'auto::,country:,database:,dependencies,jail:,netmask:,port:,help' -- "$@") + 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 @@ -238,6 +264,14 @@ function parse_command_line_args() { check_dependencies exit $? ;; + '-f' | '--config-file') + configfile="${2}" + if [[ ! -f "${configfile}" || ! -r "${configfile}" ]]; then + echo "Can not read configuration file '${2}'. Invoke with --help for help." >&2 + exit 1 + fi + shift + ;; '-j' | '--jail') jail="${2}" shift @@ -283,6 +317,11 @@ function parse_command_line_args() { shift done + # If the config file option is set, parse the config file. + if [[ ! -z ${configfile+x} ]]; then + parse_config_file + fi + if [[ -z "${database}" ]]; then echo "No GeoIP database specified. Invoke with --help for more information." >&2 exit 1 From 8e0f22da8f646df2a888b759597109b1b910d892 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 20:42:11 +0100 Subject: [PATCH 03/10] Make it actually work. To do: validate input from config file. --- ddos-mitigator.conf | 2 +- ddos-mitigator.sh | 43 ++++++++++++++++++++++++++++++++----------- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/ddos-mitigator.conf b/ddos-mitigator.conf index a7ebc54..abcd59e 100644 --- a/ddos-mitigator.conf +++ b/ddos-mitigator.conf @@ -10,7 +10,7 @@ DATABASE_FILE="/path/to/geoip/country-or-city-database.mmdb" # Enable the autopilot for automatically banning IP addresses of the desired countries (see also COUNTRIES option). # Only ban IP addresses with at least AUTOPILOT current connections. If the value is not specified or 0, don't # automatically ban IP addresses, but run in interactive mode. -AUTOPILOT="1" +AUTOPILOT="0" # Defines the subnet size in bytes to be analyzed. Valid values are: # - 8 for class A networks (X.0.0.0/8) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index 01689ea..a5b4b80 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -198,24 +198,42 @@ function filter() { mv "${filtered}" "${file}" } +function set_default_values() { + if [[ -z "${autopilot}" ]]; then + autopilot=0 + fi + if [[ -z "${netmask}" ]]; then + netmask=0 + fi + if [[ -z "${jail}" ]]; then + jail="apache-auth" + fi + if [[ -z "${bancountries}" ]]; then + bancountries=("CN") + fi + if [[ -z "${port}" ]]; then + port=443 + fi +} + function parse_config_file() { source "${configfile}" - if [[ -z "${autopilot+x}" ]]; then + if [[ -z "${autopilot}" ]]; then autopilot="${AUTOPILOT}" fi if [[ -z "${bancountries}" ]]; then - bancountries=()${COUNTRIES[@]}) + bancountries=(${COUNTRIES[@]}) fi - if [[ -z "${database+x}" ]]; then + if [[ -z "${database}" ]]; then database="${DATABASE_FILE}" fi - if [[ -z "${jail+x}" ]]; then + if [[ -z "${jail}" ]]; then jail="${JAIL}" fi - if [[ -z "${netmask+x}" ]]; then + if [[ -z "${netmask}" ]]; then netmask="${NETMASK}" fi - if [[ -z "${port+x}" ]]; then + if [[ -z "${port}" ]]; then port="${PORT}" fi } @@ -331,6 +349,9 @@ function parse_command_line_args() { echo "Database '${database}' is not accessible." >&2 exit 1 fi + + # Here, we set the default values for all options that have not been set yet. + set_default_values } ################################################################################ @@ -483,12 +504,12 @@ banlist="${tmpdir}/banlist.txt" touch "${banlist}" # Parse the command line options -autopilot=0 -netmask=0 -jail="apache-auth" -bancountries=("CN") +autopilot= +netmask= +jail= +bancountries= database= -port=443 +port= parse_command_line_args "$@" From 11ad1b471a7aecdb30b60e8390fa26328f86f7d8 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 21:28:28 +0100 Subject: [PATCH 04/10] Change the way default parameters are set. --- ddos-mitigator.sh | 76 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 54 insertions(+), 22 deletions(-) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index a5b4b80..abfdf31 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -256,10 +256,10 @@ function parse_command_line_args() { '') autopilot=1 ;; - *[!0-9]*) - echo "Invalid argument for parameter 'auto': '${2}'. Invoke with --help for help." >&2 - exit 1 - ;; + # *[!0-9]*) + # echo "Invalid argument for parameter 'auto': '${2}'. Invoke with --help for help." >&2 + # exit 1 + # ;; *) autopilot="$2" ;; @@ -268,10 +268,10 @@ function parse_command_line_args() { ;; '-c' | '--country') 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 + # if [[ -z ${bancountries[@]// /} ]]; then + # echo "Invalid argument for parameter 'country': '${2}'. Invoke with --help for help." >&2 + # exit 1 + # fi shift ;; '-d' | '--database') @@ -284,10 +284,6 @@ function parse_command_line_args() { ;; '-f' | '--config-file') configfile="${2}" - if [[ ! -f "${configfile}" || ! -r "${configfile}" ]]; then - echo "Can not read configuration file '${2}'. Invoke with --help for help." >&2 - exit 1 - fi shift ;; '-j' | '--jail') @@ -296,21 +292,20 @@ function parse_command_line_args() { ;; '-n' | '--netmask') case "${2}" in - '1' | '8') + '1') netmask=8 ;; - '2' | '16') + '2') netmask=16 ;; - '3' | '24') + '3') netmask=24 ;; - '4' | '32') + '4') netmask=32 ;; *) - echo "Invalid argument for parameter 'netmask': '${2}'. Invoke with --help for help." >&2 - exit 1 + netmask="${2}" ;; esac shift @@ -337,21 +332,57 @@ function parse_command_line_args() { # 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() { + # Autopilot + case "${autopilot}" in + *[!0-9]*) + echo "Invalid value argument for parameter 'auto' / 'AUTOPILOT': '${autopilot}'." >&2 + echo "Invoke with --help for help." >&2 + exit 1 + ;; + esac + + # Countries + if [[ -z ${bancountries[@]// /} ]]; then + echo "Invalid argument for parameter 'country' / 'COUNTRIES': '${bancountries[*]}'." >&2 + echo "Invoke with --help for help." >&2 + exit 1 + fi + + # Netmask + case "${netmask}" in + '8' | '16' | '24' | '32') + # Everything OK. + ;; + *) + echo "Invalid argument for parameter 'netmask': '${2}'." >&2 + echo "Invoke with --help for help." >&2 + exit 1 + ;; + esac + + # GeoIP-Database if [[ -z "${database}" ]]; then echo "No GeoIP database specified. Invoke with --help for more information." >&2 exit 1 fi - if [[ ! -r "${database}" ]]; then + if [[ ! -f "${database}" || ! -r "${database}" ]]; then echo "Database '${database}' is not accessible." >&2 exit 1 fi - - # Here, we set the default values for all options that have not been set yet. - set_default_values } ################################################################################ @@ -512,6 +543,7 @@ database= port= parse_command_line_args "$@" +validate_parameter_values check_dependencies dependencies_ok=$? From f5180a5e57f36c3b2687f7e3c18a52df03870c57 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 21:15:37 +0100 Subject: [PATCH 05/10] Improve parameter validation. --- ddos-mitigator.sh | 87 +++++++++++++++++++++++++++-------------------- 1 file changed, 51 insertions(+), 36 deletions(-) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index abfdf31..adcf837 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -256,10 +256,6 @@ function parse_command_line_args() { '') autopilot=1 ;; - # *[!0-9]*) - # echo "Invalid argument for parameter 'auto': '${2}'. Invoke with --help for help." >&2 - # exit 1 - # ;; *) autopilot="$2" ;; @@ -268,10 +264,6 @@ function parse_command_line_args() { ;; '-c' | '--country') 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') @@ -345,34 +337,6 @@ function parse_command_line_args() { } function validate_parameter_values() { - # Autopilot - case "${autopilot}" in - *[!0-9]*) - echo "Invalid value argument for parameter 'auto' / 'AUTOPILOT': '${autopilot}'." >&2 - echo "Invoke with --help for help." >&2 - exit 1 - ;; - esac - - # Countries - if [[ -z ${bancountries[@]// /} ]]; then - echo "Invalid argument for parameter 'country' / 'COUNTRIES': '${bancountries[*]}'." >&2 - echo "Invoke with --help for help." >&2 - exit 1 - fi - - # Netmask - case "${netmask}" in - '8' | '16' | '24' | '32') - # Everything OK. - ;; - *) - echo "Invalid argument for parameter 'netmask': '${2}'." >&2 - echo "Invoke with --help for help." >&2 - exit 1 - ;; - esac - # GeoIP-Database if [[ -z "${database}" ]]; then echo "No GeoIP database specified. Invoke with --help for more information." >&2 @@ -383,6 +347,57 @@ function validate_parameter_values() { 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 } ################################################################################ From 5078363348241ede1a9b75d1cf0e67d8313cac28 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 21:21:44 +0100 Subject: [PATCH 06/10] Harmonize parameter processing order. --- ddos-mitigator.sh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index adcf837..f90c0e1 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -202,14 +202,14 @@ function set_default_values() { if [[ -z "${autopilot}" ]]; then autopilot=0 fi - if [[ -z "${netmask}" ]]; then - netmask=0 + if [[ -z "${bancountries}" ]]; then + bancountries=("CN") fi if [[ -z "${jail}" ]]; then jail="apache-auth" fi - if [[ -z "${bancountries}" ]]; then - bancountries=("CN") + if [[ -z "${netmask}" ]]; then + netmask=0 fi if [[ -z "${port}" ]]; then port=443 From b3afda4932e12af2aa9b693cfc6f0f505198259c Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 22:06:59 +0100 Subject: [PATCH 07/10] Fix silly, SILLY syntax error. --- ddos-mitigator.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/ddos-mitigator.sh b/ddos-mitigator.sh index f90c0e1..5dbc2b7 100755 --- a/ddos-mitigator.sh +++ b/ddos-mitigator.sh @@ -396,7 +396,6 @@ function validate_parameter_values() { echo "Value must be between 0 ... 65535 (inclusive)." >&2 echo "Invoke with --help for help." >&2 exit 1 - ;; fi } From 70d41673f70b05dc1b5ecfcf36c189d52a2ef503 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 22:04:38 +0100 Subject: [PATCH 08/10] Try and handle python exceptions more gracefully. --- geoip-lookup.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/geoip-lookup.py b/geoip-lookup.py index 134dea4..4a83521 100755 --- a/geoip-lookup.py +++ b/geoip-lookup.py @@ -63,16 +63,16 @@ def parse_command_line(argv): try: opts, args = getopt.getopt(argv, "f:") except getopt.GetoptError as e: - raise LookupException("Error parsing command line") from e + raise LookupException("Error parsing command line. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") from e for opt, arg in opts: if opt == "-f": dbfile = arg else: - raise LookupException("Unknown command line argument") + raise LookupException("Unknown command line argument. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") if len(args) == 0: - raise LookupException("No address given on command line") + raise LookupException("No address given on command line. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") return dbfile, args[0] @@ -83,9 +83,13 @@ def main(argv): :param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress" :return: """ - (dbfile, ipaddress) = parse_command_line(argv) - code = get_county_code(ipaddress, dbfile) - print(code) + try + (dbfile, ipaddress) = parse_command_line(argv) + code = get_county_code(ipaddress, dbfile) + print(code) + except LookupException as e: + print(e, file=sys.stderr) + print("Unknown") if __name__ == '__main__': From 31f09011341dec17a63cd643bda133dc2913cde7 Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 24 Nov 2020 23:07:38 +0100 Subject: [PATCH 09/10] Fix bugs and improve command line argument parsing and error handling. --- geoip-lookup.py | 47 +++++++++++++++++++++++++++-------------------- 1 file changed, 27 insertions(+), 20 deletions(-) diff --git a/geoip-lookup.py b/geoip-lookup.py index 4a83521..55d3515 100755 --- a/geoip-lookup.py +++ b/geoip-lookup.py @@ -1,16 +1,24 @@ #!/usr/bin/python3.8 -import getopt +import argparse import sys try: import geoip2.database + import geoip2.errors except ImportError: print( - "Required module geoip2.database not found. On Gentoo Linux, please install dev-python/geoip2 from the 'fritteli' overlay.", + "Required modules geoip2.database and geoip2.errors not found. On Gentoo Linux, please install dev-python/geoip2 from the 'fritteli' overlay.", file=sys.stderr) exit(1) +try: + import maxminddb.errors +except ImportError: + print( + "Required module maxminddb.errors not found. On Gentoo Linux, please install dev-python/maxminddb from the 'fritteli' overlay.", + file=sys.stderr) + exit(1) class LookupException(Exception): """ @@ -47,6 +55,14 @@ def get_county_code(ipaddress, dbfile): raise LookupException("Unsupported DB type: " + dbtype) return country.iso_code + except FileNotFoundError as e: + raise LookupException(e.args) + except maxminddb.errors.InvalidDatabaseError as e: + raise LookupException(e.args) + except geoip2.errors.AddressNotFoundError as e: + raise LookupException(e.args) + except ValueError as e: + raise LookupException(e.args) finally: if reader: reader.close() @@ -60,21 +76,12 @@ def parse_command_line(argv): :return: """ dbfile = None - try: - opts, args = getopt.getopt(argv, "f:") - except getopt.GetoptError as e: - raise LookupException("Error parsing command line. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") from e - - for opt, arg in opts: - if opt == "-f": - dbfile = arg - else: - raise LookupException("Unknown command line argument. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") - - if len(args) == 0: - raise LookupException("No address given on command line. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") - return dbfile, args[0] - + parser = argparse.ArgumentParser(description='Get the country code from an IP address') + parser.add_argument('-f', dest='dbfile', required=True, help="Path to the GeoIP2 database file") + parser.add_argument('address', help="The IP address to check") + args = parser.parse_args() + + return args.dbfile, args.address def main(argv): """ @@ -83,12 +90,12 @@ def main(argv): :param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress" :return: """ - try + try: (dbfile, ipaddress) = parse_command_line(argv) code = get_county_code(ipaddress, dbfile) print(code) except LookupException as e: - print(e, file=sys.stderr) + print(e.args, file=sys.stderr) print("Unknown") @@ -96,5 +103,5 @@ if __name__ == '__main__': try: main(sys.argv[1:]) except BaseException as e: - print("Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1", file=sys.stderr) + print("Unknown") raise e From 64918328c9b01b50af11d28def4dadfffeed979f Mon Sep 17 00:00:00 2001 From: Manuel Friedli Date: Tue, 8 Jun 2021 22:45:25 +0200 Subject: [PATCH 10/10] Don't call python3.8, but python. --- geoip-lookup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/geoip-lookup.py b/geoip-lookup.py index 55d3515..ccf4fc5 100755 --- a/geoip-lookup.py +++ b/geoip-lookup.py @@ -1,4 +1,4 @@ -#!/usr/bin/python3.8 +#!/usr/bin/python import argparse import sys