From 9ce8603f3378129173f3a3d3fac326428a78cb5d Mon Sep 17 00:00:00 2001
From: Alexandre Dulaunoy <a@foo.be>
Date: Sun, 6 Feb 2022 11:33:17 +0100
Subject: [PATCH] new: [multiple MMDB file support] MMDB server can now load a
 list of MMDB files + new GeoOpen file including ASN details added as example

GeoOpen files for country only -> https://cra.circl.lu/opendata/geo-open/mmdb-country/
GeoOpen files with country and ASN/ASN description -> https://cra.circl.lu/opendata/geo-open/mmdb-country-asn/
---
 bin/server.py          | 65 +++++++++++++++++++++++++-----------------
 etc/server.conf.sample |  2 +-
 2 files changed, 40 insertions(+), 27 deletions(-)

diff --git a/bin/server.py b/bin/server.py
index 7853b1c..5c607fb 100644
--- a/bin/server.py
+++ b/bin/server.py
@@ -15,7 +15,7 @@ from wsgiref.simple_server import make_server
 import falcon
 import maxminddb
 
-version = "0.1"
+version = "0.5"
 config = configparser.ConfigParser()
 config.read('../etc/server.conf')
 mmdb_file = config['global'].get('mmdb_file')
@@ -23,6 +23,8 @@ pubsub = config['global'].getboolean('lookup_pubsub')
 port = config['global'].getint('port')
 country_file = config['global'].get('country_file')
 
+mmdb_files = mmdb_file.split(",")
+
 with open(country_file) as j:
     country_info = json.load(j)
 
@@ -30,16 +32,19 @@ if pubsub:
     import redis
     rdb = redis.Redis(host='127.0.0.1')
 
-meta = {}
-q = maxminddb.open_database(mmdb_file, maxminddb.MODE_MEMORY)
-meta['description'] = q.metadata().description
-meta['build_db'] = time.strftime(
-    '%Y-%m-%d %H:%M:%S', time.localtime(q.metadata().build_epoch)
-)
-meta['db_source'] = q.metadata().database_type
-meta['nb_nodes'] = q.metadata().node_count
-
+mmdbs = []
+for mmdb_file in mmdb_files:
+    meta = {}
+    meta['reader'] = maxminddb.open_database(mmdb_file, maxminddb.MODE_MEMORY)
+    meta['description'] = meta['reader'].metadata().description
+    meta['build_db'] = time.strftime(
+        '%Y-%m-%d %H:%M:%S', time.localtime(meta['reader'].metadata().build_epoch)
+    )
+    meta['db_source'] = meta['reader'].metadata().database_type
+    meta['nb_nodes'] = meta['reader'].metadata().node_count
+    mmdbs.append(meta)
 
+print(mmdbs)
 def validIPAddress(IP: str) -> bool:
     try:
         type(ip_address(IP))
@@ -72,14 +77,18 @@ class GeoLookup:
             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}')
-        georesult = q.get(value)
-        ret.append(georesult)
-        georesult['meta'] = meta
-        georesult['ip'] = value
-        if georesult['country']['iso_code'] != 'None':
-            georesult['country_info'] = countryLookup(country=georesult['country']['iso_code'])
-        else:
-            georesult['country_info'] = {}
+        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
 
@@ -88,14 +97,18 @@ class MyGeoLookup:
     def on_get(self, req, resp):
         ret = []
         ips = req.access_route
-        georesult = q.get(ips[0])
-        ret.append(georesult)
-        georesult['meta'] = meta
-        georesult['ip'] = ips[0]
-        if georesult['country']['iso_code'] != 'None':
-            georesult['country_info'] = countryLookup(country=georesult['country']['iso_code'])
-        else:
-            georesult['country_info'] = {}
+        for mmdb in mmdbs:
+            m = {}
+            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
 
diff --git a/etc/server.conf.sample b/etc/server.conf.sample
index 894beaf..2ccdd2a 100644
--- a/etc/server.conf.sample
+++ b/etc/server.conf.sample
@@ -1,5 +1,5 @@
 [global]
-mmdb_file = /home/adulau/git/mmdb-server/db/GeoOpen-country.mmd
+mmdb_file = /home/adulau/git/mmdb-server/db/GeoOpen-Country.mmdb,/home/adulau/git/mmdb-server/db/GeoOpen-Country-ASN.mmdb
 country_file = /home/adulau/git/mmdb-server/db/country.json
 lookup_pubsub = no 
 port = 8000