Merge pull request #4949 from sarthurdev/geoip_refactor

geoip: T7926: T8049: Refactor geoip code and add MaxMind support

1 files changed, 47 insertions, 4 deletions

diff --git a/src/helpers/geoip-update.py b/src/helpers/geoip-update.py
index 22d26e538..8a414267e 100755
--- a/src/helpers/geoip-update.py
+++ b/src/helpers/geoip-update.py
@@ -18,7 +18,13 @@ import argparse
 import sys
 
 from vyos.configquery import ConfigTreeQuery
-from vyos.firewall import geoip_update
+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):
     if config:
@@ -27,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,
@@ -35,9 +43,44 @@ def get_config(config=None):
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("--force", help="Force update", action="store_true")
+    parser.add_argument("--init", help="Initialise", action="store_true")
     args = parser.parse_args()
 
-    firewall, policy = get_config()
-    if not geoip_update(firewall=firewall, policy=policy, force=args.force):
+    if args.init:
+        db_initialise()
+        db_import_dbip_ranges(delete_file=True)
+        sys.exit(0)
+
+    options, firewall, policy = get_config()
+
+    if not db_is_initialised():
+        db_initialise()
+
+    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)
+
+    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)
+
+    if not geoip_update(firewall=firewall, policy=policy):
         sys.exit(1)