geoip-lookup.py

@@ -26,12 +26,13 @@ class LookupException(Exception):
     """
     pass
 
-def get_geoip_result(ipaddress, dbfile):
+
+def get_county_code(ipaddress, dbfile):
     """
-    Perform a lookup of the given ipaddress in the database.
+    Determine the country code that the given ipaddress comes from.
     :param ipaddress: The IP address to look up
     :param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database)
-    :return: The lookup result
+    :return: The ISO country code (2 letters)
     """
     if not ipaddress:
         raise LookupException("No address given")
@@ -42,17 +43,18 @@ def get_geoip_result(ipaddress, dbfile):
     try:
         reader = geoip2.database.Reader(dbfile)
         dbtype = reader.metadata().database_type
-        result = None
+        country = None
         if dbtype == 'GeoLite2-City' or dbtype == 'GeoIP2-City':
-            result = reader.city(ipaddress)
-        elif dbtype == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
-            result = reader.country(ipaddress)
+            country = reader.city(ipaddress).country
+        elif dbfile == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
+            country = reader.country(ipaddress).country
         # ASN is not supported
         # elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':
-        else:
+
+        if not country:
             raise LookupException("Unsupported DB type: " + dbtype)
 
-        return result
+        return country.iso_code
     except FileNotFoundError as e:
         raise LookupException(e.args)
     except maxminddb.errors.InvalidDatabaseError as e:
@@ -65,30 +67,6 @@ def get_geoip_result(ipaddress, dbfile):
         if reader:
             reader.close()
 
-def get_county_code(ipaddress, dbfile):
-    """
-    Determine the country code that the given ipaddress comes from.
-    :param ipaddress: The IP address to look up
-    :param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database)
-    :return: The ISO country code (2 letters)
-    """
-    result = get_geoip_result(ipaddress, dbfile)
-
-    return result.country.iso_code
-
-def get_details(ipaddress, dbfile):
-    """
-    Determine the country code, continent code and the network that the given ipaddress comes from.
-    :param ipaddress: The IP address to look up
-    :param dbfile: The path to the GeoIP2/GeoLite2 database file (Country or City database)
-    :return: A string consisting of the ISO country code (2 letters), ISO continent code (2 letters) and network (CIDR notation), concatenated by ","
-    """
-    result = get_geoip_result(ipaddress, dbfile)
-    country = result.country.iso_code
-    continent = result.continent.code
-    network = result.traits.network
-
-    return "%s,%s,%s" % (country, continent, network)
 
 def parse_command_line(argv):
     """
@@ -100,34 +78,26 @@ def parse_command_line(argv):
     dbfile = None
     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('-d', dest='detail', action='store_const', const=True, default=False, help="Verbose output: Print continent code and network along with country code")
-    parser.add_argument('--detail', dest='detail', action='store_const', const=True, default=False, help="Verbose output: Print continent code and network along with country code")
     parser.add_argument('address', help="The IP address to check")
     args = parser.parse_args()
-
-    return args.dbfile, args.address, args.detail
+    
+    return args.dbfile, args.address
 
 def main(argv):
     """
     Read the database file and the IP address from the command line and print the corresponding ISO country code on
     stdout.
-    :param argv: Format: "-f /path/to/database.mmdb [--detail|-d] ip.v4.add.ress"
+    :param argv: Format: "-f /path/to/database.mmdb ip.v4.add.ress"
     :return:
     """
     try:
-        (dbfile, ipaddress, detail) = parse_command_line(argv)
-        if detail:
-            code = get_details(ipaddress, dbfile)
-        else:
-            code = get_county_code(ipaddress, dbfile)
-
+        (dbfile, ipaddress) = parse_command_line(argv)
+        code = get_county_code(ipaddress, dbfile)
         print(code)
     except LookupException as e:
         print(e.args, file=sys.stderr)
-        if detail:
-            print("Unknown,,")
-        else:
-            print("Unknown")
+        print("Unknown")
+
 
 if __name__ == '__main__':
     try: