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

This commit is contained in:
Manuel Friedli 2020-11-24 23:07:38 +01:00
parent 70d41673f7
commit 31f0901134

View file

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