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:
+    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") from e
+    args = parser.parse_args()
-
-    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]
     
+    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:
     """
+    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