Merge pull request #4949 from sarthurdev/geoip_refactor

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

3 files changed, 116 insertions, 76 deletions

diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index a609e305c..02c9b2e3e 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -22,13 +22,13 @@ from glob import glob
 from sys import exit
 from vyos.base import Warning
 from vyos.config import Config
-from vyos.configdict import is_node_changed
-from vyos.configdiff import get_config_diff, Diff
+from vyos.configdict import is_node_changed, node_changed
+from vyos.configdiff import Diff
 from vyos.configdep import set_dependents, call_dependents
 from vyos.configverify import verify_interface_exists
 from vyos.ethtool import Ethtool
 from vyos.firewall import fqdn_config_parse
-from vyos.firewall import geoip_update
+from vyos.geoip import geoip_refresh, geoip_update
 from vyos.template import render
 from vyos.utils.dict import dict_search
 from vyos.utils.dict import dict_search_args
@@ -80,42 +80,29 @@ snmp_event_source = 1
 snmp_trap_mib = 'VYATTA-TRAP-MIB'
 snmp_trap_name = 'mgmtEventTrap'
 
-def geoip_updated(conf, firewall):
-    diff = get_config_diff(conf)
-    node_diff = diff.get_child_nodes_diff(['firewall'], expand_nodes=Diff.DELETE, recursive=True)
+def geoip_sets(firewall):
+    out = {'name': [], 'ipv6_name': []}
 
-    out = {
-        'name': [],
-        'ipv6_name': [],
-        'deleted_name': [],
-        'deleted_ipv6_name': []
-    }
-    updated = False
-
-    for key, path in dict_search_recursive(firewall, 'geoip'):
-        set_name = f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}'
+    for _, path in dict_search_recursive(firewall, 'geoip'):
         if (path[0] == 'ipv4'):
-            out['name'].append(set_name)
+            out['name'].append(f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}')
         elif (path[0] == 'ipv6'):
-            set_name = f'GEOIP_CC6_{path[1]}_{path[2]}_{path[4]}'
-            out['ipv6_name'].append(set_name)
+            out['ipv6_name'].append(f'GEOIP_CC6_{path[1]}_{path[2]}_{path[4]}')
 
-        updated = True
+    return out
 
-    if 'delete' in node_diff:
-        for key, path in dict_search_recursive(node_diff['delete'], 'geoip'):
-            set_name = f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}'
-            if (path[0] == 'ipv4'):
-                out['deleted_name'].append(set_name)
-            elif (path[0] == 'ipv6'):
-                set_name = f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}'
-                out['deleted_ipv6_name'].append(set_name)
-            updated = True
+def geoip_updated(conf):
+    changes = node_changed(conf, ['firewall'],
+                                 key_mangling=('-', '_'),
+                                 recursive=True,
+                                 expand_nodes=Diff.ADD | Diff.DELETE)
+    updated = False
 
-    if updated:
-        return out
+    for _, path in dict_search_recursive(changes, 'geoip'):
+        updated = True
+        break
 
-    return False
+    return updated
 
 def get_config(config=None):
     if config:
@@ -135,7 +122,8 @@ def get_config(config=None):
         # Update nat and policy-route as firewall groups were updated
         set_dependents('group_resync', conf)
 
-    firewall['geoip_updated'] = geoip_updated(conf, firewall)
+    firewall['geoip_sets'] = geoip_sets(firewall)
+    firewall['geoip_updated'] = geoip_updated(conf)
 
     fqdn_config_parse(firewall, 'firewall')
 
@@ -503,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 = []
@@ -725,11 +719,12 @@ def apply(firewall):
             domain_action = 'stop'
     call(f'systemctl {domain_action} vyos-domain-resolver.service')
 
-    if firewall['geoip_updated']:
+    if firewall['geoip_sets']:
         # Call helper script to Update set contents
-        if 'name' in firewall['geoip_updated'] or 'ipv6_name' in firewall['geoip_updated']:
-            print('Updating GeoIP. Please wait...')
-            geoip_update(firewall=firewall)
+        if 'name' in firewall['geoip_sets'] or 'ipv6_name' in firewall['geoip_sets']:
+            if firewall['geoip_updated'] or not geoip_refresh():
+                print('Updating GeoIP. Please wait...')
+                geoip_update(firewall)
 
     return None
 
diff --git a/src/conf_mode/policy_route.py b/src/conf_mode/policy_route.py
index 15d735f75..17f8e6bf9 100755
--- a/src/conf_mode/policy_route.py
+++ b/src/conf_mode/policy_route.py
@@ -21,7 +21,8 @@ from sys import exit
 
 from vyos.base import Warning
 from vyos.config import Config
-from vyos.configdiff import get_config_diff, Diff
+from vyos.configdict import node_changed
+from vyos.configdiff import Diff
 from vyos.template import render
 from vyos.utils.dict import dict_search_args
 from vyos.utils.dict import dict_search_recursive
@@ -30,7 +31,7 @@ from vyos.utils.process import run
 from vyos.utils.network import get_vrf_tableid
 from vyos.defaults import rt_global_table
 from vyos.defaults import rt_global_vrf
-from vyos.firewall import geoip_update
+from vyos.geoip import  geoip_refresh, geoip_update
 from vyos import ConfigError
 from vyos import airbag
 airbag.enable()
@@ -46,42 +47,40 @@ valid_groups = [
     'interface_group'
 ]
 
-def geoip_updated(conf, policy):
-    diff = get_config_diff(conf)
-    node_diff = diff.get_child_nodes_diff(['policy'], expand_nodes=Diff.DELETE, recursive=True)
-
-    out = {
-        'name': [],
-        'ipv6_name': [],
-        'deleted_name': [],
-        'deleted_ipv6_name': []
-    }
+def geoip_updated(conf):
     updated = False
 
-    for key, path in dict_search_recursive(policy, 'geoip'):
-        set_name = f'GEOIP_CC_{path[0]}_{path[1]}_{path[3]}'
-        if (path[0] == 'route'):
-            out['name'].append(set_name)
-        elif (path[0] == 'route6'):
-            set_name = f'GEOIP_CC6_{path[0]}_{path[1]}_{path[3]}'
-            out['ipv6_name'].append(set_name)
+    changes_v4 = node_changed(conf, ['policy', 'route'],
+                                 key_mangling=('-', '_'),
+                                 recursive=True,
+                                 expand_nodes=Diff.ADD | Diff.DELETE)
 
+    for _, path in dict_search_recursive(changes_v4, 'geoip'):
         updated = True
+        break
 
-    if 'delete' in node_diff:
-        for key, path in dict_search_recursive(node_diff['delete'], 'geoip'):
-            set_name = f'GEOIP_CC_{path[0]}_{path[1]}_{path[3]}'
-            if (path[0] == 'route'):
-                out['deleted_name'].append(set_name)
-            elif (path[0] == 'route6'):
-                set_name = f'GEOIP_CC6_{path[0]}_{path[1]}_{path[3]}'
-                out['deleted_ipv6_name'].append(set_name)
+    if not updated:
+        changes_v6 = node_changed(conf, ['policy', 'route6'],
+                                 key_mangling=('-', '_'),
+                                 recursive=True,
+                                 expand_nodes=Diff.ADD | Diff.DELETE)
+
+        for _, path in dict_search_recursive(changes_v6, 'geoip'):
             updated = True
+            break
 
-    if updated:
-        return out
+    return updated
 
-    return False
+def geoip_sets(policy):
+    out = {'name': [], 'ipv6_name': []}
+
+    for _, path in dict_search_recursive(policy, 'geoip'):
+        if (path[0] == 'route'):
+            out['name'].append(f'GEOIP_CC_{path[0]}_{path[1]}_{path[3]}')
+        elif (path[0] == 'route6'):
+            out['ipv6_name'].append(f'GEOIP_CC6_{path[0]}_{path[1]}_{path[3]}')
+
+    return out
 
 def get_config(config=None):
     if config:
@@ -100,7 +99,9 @@ def get_config(config=None):
     if 'dynamic_group' in policy['firewall_group']:
         del policy['firewall_group']['dynamic_group']
 
-    policy['geoip_updated'] = geoip_updated(conf, policy)
+    policy['geoip_sets'] = geoip_sets(policy)
+    policy['geoip_updated'] = geoip_updated(conf)
+
     return policy
 
 def verify_rule(policy, name, rule_conf, ipv6, rule_id):
@@ -244,11 +245,12 @@ def apply(policy):
 
     apply_table_marks(policy)
 
-    if policy['geoip_updated']:
+    if policy['geoip_sets']:
         # Call helper script to Update set contents
-        if 'name' in policy['geoip_updated'] or 'ipv6_name' in policy['geoip_updated']:
-            print('Updating GeoIP. Please wait...')
-            geoip_update(policy=policy)
+        if 'name' in policy['geoip_sets'] or 'ipv6_name' in policy['geoip_sets']:
+            if policy['geoip_updated'] or not geoip_refresh():
+                print('Updating GeoIP. Please wait...')
+                geoip_update(policy=policy)
 
     return None
 
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)