From 75b8e2fb6cfbf0165a01ce0b8d60153910d9add3 Mon Sep 17 00:00:00 2001
From: Manuel Friedli <manuel@fritteli.ch>
Date: Thu, 15 Jul 2021 22:02:57 +0200
Subject: [PATCH] Refactoring.

---
 geoip-lookup.py | 80 ++++++++++++++++++++-----------------------------
 1 file changed, 33 insertions(+), 47 deletions(-)

diff --git a/geoip-lookup.py b/geoip-lookup.py
index 3bf3474..28ff122 100755
--- a/geoip-lookup.py
+++ b/geoip-lookup.py
@@ -26,51 +26,12 @@ class LookupException(Exception):
     """
     pass
 
-def get_county_code(ipaddress, dbfile):
+def get_geoip_result(ipaddress, dbfile):
     """
-    Determine the country code that the given ipaddress comes from.
+    Perform a lookup of the given ipaddress in the database.
     :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)
-    """
-    if not ipaddress:
-        raise LookupException("No address given")
-    if not dbfile:
-        raise LookupException("No db file given")
-
-    reader = None
-    try:
-        reader = geoip2.database.Reader(dbfile)
-        dbtype = reader.metadata().database_type
-        country = None
-        if dbtype == 'GeoLite2-City' or dbtype == 'GeoIP2-City':
-            country = reader.city(ipaddress).country
-        elif dbtype == 'GeoLite2-Country' or dbtype == 'GeoIP2-Country':
-            country = reader.country(ipaddress).country
-        # ASN is not supported
-        # elif dbfile == 'GeoLite2-ASN' or dbtype == 'GeoIP2-ASN':
-        else:
-            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()
-
-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 ","
+    :return: The lookup result
     """
     if not ipaddress:
         raise LookupException("No address given")
@@ -91,10 +52,7 @@ def get_details(ipaddress, dbfile):
         else:
             raise LookupException("Unsupported DB type: " + dbtype)
 
-        country = result.country.iso_code
-        continent = result.continent.code
-        network = result.traits.network
-        return "%s,%s,%s" % (country, continent, network)
+        return result
     except FileNotFoundError as e:
         raise LookupException(e.args)
     except maxminddb.errors.InvalidDatabaseError as e:
@@ -107,6 +65,31 @@ def get_details(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):
     """
     Parse the command line. First, the database file must be specified ("-f /path/to/db/file.mmdb"), then the IP address
@@ -141,7 +124,10 @@ def main(argv):
         print(code)
     except LookupException as e:
         print(e.args, file=sys.stderr)
-        print("Unknown")
+        if detail:
+            print("Unknown,,")
+        else:
+            print("Unknown")
 
 if __name__ == '__main__':
     try: