diff --git a/geoip-lookup.py b/geoip-lookup.py index 134dea4..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") from e - - for opt, arg in opts: - if opt == "-f": - dbfile = arg - else: - raise LookupException("Unknown command line argument") - - if len(args) == 0: - raise LookupException("No address given on command line") - 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,14 +90,18 @@ 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.args, file=sys.stderr) + print("Unknown") 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