linux/ddos-mitigator
1
0
Fork 0
Code Issues Pull requests Releases 5 Wiki Activity

Fix bugs and improve command line argument parsing and error handling.

Browse source
This commit is contained in:
Manuel Friedli 2020-11-24 23:07:38 +01:00
parent 70d41673f7
commit 31f0901134
1 changed files with 27 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
geoip-lookup.py
View file

@ -1,16 +1,24 @@
#!/usr/bin/python3.8 #!/usr/bin/python3.8
import getopt import argparse
import sys import sys
try: try:
import geoip2.database import geoip2.database
import geoip2.errors
except ImportError: except ImportError:
print( 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) file=sys.stderr)
exit(1) 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): class LookupException(Exception):
""" """
@ -47,6 +55,14 @@ def get_county_code(ipaddress, dbfile):
raise LookupException("Unsupported DB type: " + dbtype) raise LookupException("Unsupported DB type: " + dbtype)
return country.iso_code 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: finally:
if reader: if reader:
reader.close() reader.close()
@ -60,21 +76,12 @@ def parse_command_line(argv):
:return: :return:
""" """
dbfile = None dbfile = None
try: parser = argparse.ArgumentParser(description='Get the country code from an IP address')
opts, args = getopt.getopt(argv, "f:") parser.add_argument('-f', dest='dbfile', required=True, help="Path to the GeoIP2 database file")
except getopt.GetoptError as e: parser.add_argument('address', help="The IP address to check")
raise LookupException("Error parsing command line. Usage: geoip-lookup.py -f /path/to/geoip2-database.mmdb 192.168.1.1") from e args = parser.parse_args()
for opt, arg in opts: return args.dbfile, args.address
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]
def main(argv): def main(argv):
""" """
@ -83,12 +90,12 @@ def main(argv):
:param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress" :param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress"
:return: :return:
""" """
try try:
(dbfile, ipaddress) = parse_command_line(argv) (dbfile, ipaddress) = parse_command_line(argv)
code = get_county_code(ipaddress, dbfile) code = get_county_code(ipaddress, dbfile)
print(code) print(code)
except LookupException as e: except LookupException as e:
print(e, file=sys.stderr) print(e.args, file=sys.stderr)
print("Unknown") print("Unknown")
@ -96,5 +103,5 @@ if __name__ == '__main__':
try: try:
main(sys.argv[1:]) main(sys.argv[1:])
except BaseException as e: 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 raise e