geoip: T8049: Add MaxMind database support

4 files changed, 168 insertions, 15 deletions

diff --git a/interface-definitions/include/firewall/global-options.xml.i b/interface-definitions/include/firewall/global-options.xml.i
index e19f3a7c5..7ec07100d 100644
--- a/interface-definitions/include/firewall/global-options.xml.i
+++ b/interface-definitions/include/firewall/global-options.xml.i
@@ -130,6 +130,49 @@
       </properties>
       <defaultValue>enable</defaultValue>
     </leafNode>
+    <node name="geoip">
+      <properties>
+        <help>GeoIP options</help>
+      </properties>
+      <children>
+        <leafNode name="provider">
+          <properties>
+            <help>GeoIP database provider</help>
+            <completionHelp>
+              <list>db-ip maxmind</list>
+            </completionHelp>
+            <valueHelp>
+              <format>db-ip</format>
+              <description>Use GeoIP database by DB-IP.com</description>
+            </valueHelp>
+            <valueHelp>
+              <format>maxmind</format>
+              <description>Use GeoIP database by MaxMind (Requires API key)</description>
+            </valueHelp>
+            <constraint>
+              <regex>(db-ip|maxmind)</regex>
+            </constraint>
+          </properties>
+          <defaultValue>db-ip</defaultValue>
+        </leafNode>
+        <leafNode name="maxmind-account-id">
+          <properties>
+            <help>Account ID for MaxMind GeoIP database</help>
+          </properties>
+        </leafNode>
+        <leafNode name="maxmind-license-key">
+          <properties>
+            <help>License key for MaxMind GeoIP database</help>
+          </properties>
+        </leafNode>
+        <leafNode name="maxmind-lite">
+          <properties>
+            <help>Use MaxMind GeoLite2 database</help>
+            <valueless/>
+          </properties>
+        </leafNode>
+      </children>
+    </node>
     <leafNode name="ip-src-route">
       <properties>
         <help>Policy for handling IPv4 packets with source route option</help>
diff --git a/python/vyos/geoip.py b/python/vyos/geoip.py
index a942b888b..9a2b40899 100644
--- a/python/vyos/geoip.py
+++ b/python/vyos/geoip.py
@@ -3,7 +3,9 @@ import csv
 import gzip
 import os
 import sqlite3
+import zipfile
 
+from io import TextIOWrapper
 from pathlib import Path
 from time import strftime
 
@@ -13,7 +15,8 @@ from vyos.utils.dict import dict_search_recursive
 from vyos.utils.process import run
 
 nftables_geoip_conf = '/run/nftables-geoip.conf'
-geoip_database_raw = '/usr/share/vyos-geoip/dbip-country-lite.csv.gz'
+dbip_database_raw = '/usr/share/vyos-geoip/dbip-country-lite.csv.gz'
+mm_database_raw = '/usr/share/vyos-geoip/maxmind-country.zip'
 geoip_database_path = '/var/cache/vyos/geoip-lookup.db'
 geoip_lock_file = '/var/lock/vyos-geoip.lock'
 
@@ -22,11 +25,24 @@ geoip_lock_file = '/var/lock/vyos-geoip.lock'
 def geoip_download_dbip():
     url = 'https://download.db-ip.com/free/dbip-country-lite-{}.csv.gz'.format(strftime("%Y-%m"))
     try:
-        dirname = os.path.dirname(geoip_database_raw)
+        dirname = os.path.dirname(dbip_database_raw)
         if not os.path.exists(dirname):
             os.mkdir(dirname)
 
-        download(geoip_database_raw, url)
+        download(dbip_database_raw, url)
+        return True
+    except:
+        return False
+
+def geoip_download_maxmind(account_id : str, license_key: str, lite : bool) -> bool:
+    db_str = 'GeoLite2' if lite else 'GeoIP2'
+    url = f'https://{account_id}:{license_key}@download.maxmind.com/geoip/databases/{db_str}-Country-CSV/download?suffix=zip'
+    try:
+        dirname = os.path.dirname(mm_database_raw)
+        if not os.path.exists(dirname):
+            os.mkdir(dirname)
+
+        download(mm_database_raw, url)
         return True
     except:
         return False
@@ -61,14 +77,14 @@ def db_initialise():
         conn.commit()
 
 def db_import_dbip_ranges(replace=True, delete_file=False):
-    if not os.path.exists(geoip_database_raw):
+    if not os.path.exists(dbip_database_raw):
         return False
 
     if not os.path.exists(geoip_database_path):
         return False
 
     try:
-        with gzip.open(geoip_database_raw, mode='rt') as csv_fh:
+        with gzip.open(dbip_database_raw, mode='rt') as csv_fh:
             reader = csv.reader(csv_fh)
 
             with sqlite3.connect(geoip_database_path) as conn:
@@ -83,7 +99,74 @@ def db_import_dbip_ranges(replace=True, delete_file=False):
                 conn.commit()
 
         if delete_file:
-            os.unlink(geoip_database_raw)
+            os.unlink(dbip_database_raw)
+
+        return True
+    except:
+        return False
+
+def db_import_maxmind_ranges(replace=True, delete_file=False):
+    if not os.path.exists(mm_database_raw):
+        return False
+
+    if not zipfile.is_zipfile(mm_database_raw):
+        return False
+
+    if not os.path.exists(geoip_database_path):
+        return False
+
+    try:
+        with zipfile.ZipFile(mm_database_raw, mode='r') as zip_fh:
+            directory = os.path.dirname(zip_fh.namelist()[0])
+            prefix = 'GeoLite2' if any(f.startswith('GeoLite2') for f in zip_fh.namelist()) else 'GeoIP2'
+
+            ipv4_file = f'{directory}/{prefix}-Country-Blocks-IPv4.csv'
+            ipv6_file = f'{directory}/{prefix}-Country-Blocks-IPv6.csv'
+            locations_file = f'{directory}/{prefix}-Country-Locations-en.csv'
+            locations_map = {}
+
+            with zip_fh.open(locations_file) as raw_csv_fh:
+                with TextIOWrapper(raw_csv_fh, encoding='utf-8') as csv_fh:
+                    reader = csv.DictReader(csv_fh)
+
+                    for row in reader:
+                        id = row['geoname_id']
+                        locations_map[id] = row['country_iso_code']
+
+            with sqlite3.connect(geoip_database_path) as conn:
+                cur = conn.cursor()
+
+                if replace:
+                    cur.execute('DELETE FROM geoip_ranges')
+
+                with zip_fh.open(ipv4_file) as raw_csv_fh:
+                    with TextIOWrapper(raw_csv_fh, encoding='utf-8') as csv_fh:
+                        reader = csv.DictReader(csv_fh)
+                        for row in reader:
+                            id = row['geoname_id']
+
+                            if not id or id not in locations_map:
+                                continue
+
+                            code = locations_map[id]
+                            cur.execute('INSERT INTO geoip_ranges (country_code, range, version) VALUES (?, ?, 4)', (code.lower(), row['network']))
+
+                with zip_fh.open(ipv6_file) as raw_csv_fh:
+                    with TextIOWrapper(raw_csv_fh, encoding='utf-8') as csv_fh:
+                        reader = csv.DictReader(csv_fh)
+                        for row in reader:
+                            id = row['geoname_id']
+
+                            if not id or id not in locations_map:
+                                continue
+
+                            code = locations_map[id]
+                            cur.execute('INSERT INTO geoip_ranges (country_code, range, version) VALUES (?, ?, 6)', (code.lower(), row['network']))
+
+                conn.commit()
+
+        if delete_file:
+            os.unlink(mm_database_raw)
 
         return True
     except:
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 18c250d08..02c9b2e3e 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -491,6 +491,12 @@ def verify(firewall):
                 for ifname in interfaces:
                     verify_hardware_offload(ifname)
 
+    if dict_search_args(firewall, 'global_options', 'geoip', 'provider') == 'maxmind':
+        geoip_options = dict_search_args(firewall, 'global_options', 'geoip')
+        required_keys = ['maxmind_account_id', 'maxmind_license_key']
+        if not all(key in geoip_options for key in required_keys):
+            raise ConfigError('MaxMind GeoIP provider requires maxmind-account-id and maxmind-license-key')
+
     if dict_search('global_options.state_policy', firewall) is not None:
         # Generate list of chains where conntrack is disabled
         conntrack_disabled_list = []
diff --git a/src/helpers/geoip-update.py b/src/helpers/geoip-update.py
index a6838c62d..8a414267e 100755
--- a/src/helpers/geoip-update.py
+++ b/src/helpers/geoip-update.py
@@ -19,9 +19,11 @@ import sys
 
 from vyos.configquery import ConfigTreeQuery
 from vyos.geoip import geoip_download_dbip
+from vyos.geoip import geoip_download_maxmind
 from vyos.geoip import db_initialise
 from vyos.geoip import db_is_initialised
 from vyos.geoip import db_import_dbip_ranges
+from vyos.geoip import db_import_maxmind_ranges
 from vyos.geoip import geoip_update
 
 def get_config(config=None):
@@ -31,6 +33,8 @@ def get_config(config=None):
         conf = ConfigTreeQuery()
 
     return (
+        conf.get_config_dict(['firewall', 'global-options', 'geoip'], key_mangling=('-', '_'), get_first_key=True,
+                                    no_tag_node_value_mangle=True, with_defaults=True),
         conf.get_config_dict(['firewall'], key_mangling=('-', '_'), get_first_key=True,
                                     no_tag_node_value_mangle=True) if conf.exists(['firewall']) else None,
         conf.get_config_dict(['policy'], key_mangling=('-', '_'), get_first_key=True,
@@ -47,19 +51,36 @@ if __name__ == '__main__':
         db_import_dbip_ranges(delete_file=True)
         sys.exit(0)
 
+    options, firewall, policy = get_config()
+
     if not db_is_initialised():
         db_initialise()
 
-    print('Dowloading latest DB-IP database...')
-    if not geoip_download_dbip():
-        print('Failed to download, aborting.')
-        sys.exit(1)
+    if options['provider'] == 'db-ip':
+        print('Dowloading latest DB-IP database...')
+        if not geoip_download_dbip():
+            print('Failed to download, aborting.')
+            sys.exit(1)
 
-    print('Extracting database...')
-    if not db_import_dbip_ranges(delete_file=True):
-        print('Failed to extract, aborting.')
-        sys.exit(1)
+        print('Extracting database...')
+        if not db_import_dbip_ranges(delete_file=True):
+            print('Failed to extract, aborting.')
+            sys.exit(1)
+
+    elif options['provider'] == 'maxmind':
+        account_id = options['maxmind_account_id']
+        license_key = options['maxmind_license_key']
+        lite = 'maxmind_lite' in options
+
+        print('Dowloading latest MaxMind database...')
+        if not geoip_download_maxmind(account_id, license_key, lite):
+            print('Failed to download, aborting.')
+            sys.exit(1)
+
+        print('Extracting database...')
+        if not db_import_maxmind_ranges(delete_file=True):
+            print('Failed to extract, aborting.')
+            sys.exit(1)
 
-    firewall, policy = get_config()
     if not geoip_update(firewall=firewall, policy=policy):
         sys.exit(1)