Add IPInfo mmdb files support

.gitignore

@@ -140,3 +140,4 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
+poetry.lock

etc/server.conf.sample

@@ -1,5 +1,6 @@
 [global]
 mmdb_file = db/GeoOpen-Country.mmdb,db/GeoOpen-Country-ASN.mmdb
+#mmdb_file = db/ipinfo_country_asn.mmdb,db/ipinfo_country.mmdb
 country_file = db/country.json
 lookup_pubsub = no 
 port = 8000

mmdb_server/mmdb_server.py

@@ -15,7 +15,7 @@
 import falcon
 import maxminddb
 
-version = "0.5"
+version = "0.6"
 config = configparser.ConfigParser()
 config.read('etc/server.conf')
 mmdb_file = config['global'].get('mmdb_file')
@@ -30,6 +30,7 @@
 
 if pubsub:
     import redis
+
     rdb = redis.Redis(host='127.0.0.1')
 
 mmdbs = []
@@ -70,62 +71,57 @@ def countryLookup(country: str) -> dict:
         return {}
 
 
+def _process_lookup(ip):
+    """Helper function to process IP lookup and format response."""
+    ret = []
+    for mmdb in mmdbs:
+        georesult = mmdb['reader'].get(ip) or {}  # Ensure dictionary, prevent NoneType errors
+        m = mmdb.copy()
+        del m['reader']
+        georesult['meta'] = m
+        georesult['ip'] = ip
+
+        # Determine country code from old or new format
+        country_code = None
+        if isinstance(georesult.get('country'), dict):
+            country_code = georesult['country'].get('iso_code')
+        elif isinstance(georesult.get('country'), str):
+            country_code = georesult['country']
+
+        georesult['country_info'] = countryLookup(country_code) if country_code else {}
+        ret.append(georesult)
+    return ret
+
+
 class GeoLookup:
     def on_get(self, req, resp, value):
-        ret = []
         ua = req.get_header('User-Agent')
         ips = req.access_route
         if not validIPAddress(value):
             resp.status = falcon.HTTP_422
             resp.media = "IPv4 or IPv6 address is in an incorrect format. Dotted decimal for IPv4 or textual representation for IPv6 are required."
             return
         pubLookup(value=f'{value} via {ips} using {ua}')
-        for mmdb in mmdbs:
-            m = {}
-            georesult = mmdb['reader'].get(value)
-            m = mmdb.copy()
-            del m['reader']
-            georesult['meta'] = m
-            georesult['ip'] = value
-            if georesult['country']['iso_code'] != 'None':
-                georesult['country_info'] = countryLookup(country=georesult['country']['iso_code'])
-            else:
-                georesult['country_info'] = {}
-            ret.append(georesult)
-        resp.media = ret
-        return
+        resp.media = _process_lookup(value)
 
 
 class MyGeoLookup:
     def on_get(self, req, resp):
-        ret = []
         ips = req.access_route
-        for mmdb in mmdbs:
-            georesult = mmdb['reader'].get(ips[0])
-            m = mmdb.copy()
-            del m['reader']
-            georesult['meta'] = m
-            georesult['ip'] = ips[0]
-            if georesult['country']['iso_code'] != 'None':
-                georesult['country_info'] = countryLookup(country=georesult['country']['iso_code'])
-            else:
-                georesult['country_info'] = {}
-            ret.append(georesult)
-        resp.media = ret
-        return
+        resp.media = _process_lookup(ips[0])
 
 
 app = falcon.App()
 
 app.add_route('/geolookup/{value}', GeoLookup())
 app.add_route('/', MyGeoLookup())
 
+
 def main():
     with make_server('', port, app) as httpd:
         print(f'Serving on port {port}...')
         httpd.serve_forever()
 
 
 if __name__ == '__main__':
-    main()
-
+    main()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mmdb-server"
-version = "0.5.0"
+version = "0.6.0"
 description = ""
 authors = ["Alexandre Dulaunoy"]
 readme = "README.md"