geoip: T8987: Support updates via source-address/vrf

4 files changed, 34 insertions, 33 deletions

diff --git a/interface-definitions/include/firewall/global-options.xml.i b/interface-definitions/include/firewall/global-options.xml.i
index 5ff5acfd8..4827f6e82 100644
--- a/interface-definitions/include/firewall/global-options.xml.i
+++ b/interface-definitions/include/firewall/global-options.xml.i
@@ -171,6 +171,8 @@
             <valueless/>
           </properties>
         </leafNode>
+        #include <include/source-address-ipv4-ipv6.xml.i>
+        #include <include/interface/vrf.xml.i>
       </children>
     </node>
     <leafNode name="ip-src-route">
diff --git a/op-mode-definitions/geoip.xml.in b/op-mode-definitions/geoip.xml.in
index 8c51817f2..6d667ff6a 100644
--- a/op-mode-definitions/geoip.xml.in
+++ b/op-mode-definitions/geoip.xml.in
@@ -6,17 +6,6 @@
         <properties>
           <help>Update GeoIP database and firewall sets</help>
         </properties>
-        <children>
-          <tagNode name="vrf">
-            <properties>
-              <help>VRF to use for GeoIP update</help>
-              <completionHelp>
-                <path>vrf name</path>
-              </completionHelp>
-            </properties>
-            <command>sudo ip vrf exec $4 ${vyos_libexec_dir}/geoip-update.py</command>
-          </tagNode>
-        </children>
         <command>sudo ${vyos_libexec_dir}/geoip-update.py</command>
       </node>
     </children>
diff --git a/python/vyos/geoip.py b/python/vyos/geoip.py
index 0f3e7f01c..70c74b04f 100644
--- a/python/vyos/geoip.py
+++ b/python/vyos/geoip.py
@@ -24,7 +24,7 @@ geoip_lock_file = '/var/lock/vyos-geoip.lock'
 
 # Raw data
 
-def geoip_download_dbip():
+def geoip_download_dbip(source_address=None):
     url = 'https://download.db-ip.com/free/dbip-country-lite-{}.csv.gz'.format(strftime("%Y-%m"))
     asn_url = 'https://download.db-ip.com/free/dbip-asn-lite-{}.csv.gz'.format(strftime("%Y-%m"))
     try:
@@ -32,13 +32,13 @@ def geoip_download_dbip():
         if not os.path.exists(dirname):
             os.mkdir(dirname)
 
-        download(dbip_database_raw, url)
-        download(dbip_asn_database_raw, asn_url)
+        download(dbip_database_raw, url, source_host=source_address)
+        download(dbip_asn_database_raw, asn_url, source_host=source_address)
         return True
     except:
         return False
 
-def geoip_download_maxmind(account_id : str, license_key: str, lite : bool) -> bool:
+def geoip_download_maxmind(account_id : str, license_key: str, lite : bool, source_address: str = None) -> 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'
     asn_url = f'https://{account_id}:{license_key}@download.maxmind.com/geoip/databases/{db_str}-ASN-CSV/download?suffix=zip'
@@ -47,8 +47,8 @@ def geoip_download_maxmind(account_id : str, license_key: str, lite : bool) -> b
         if not os.path.exists(dirname):
             os.mkdir(dirname)
 
-        download(mm_database_raw, url)
-        download(mm_asn_database_raw, asn_url)
+        download(mm_database_raw, url, source_host=source_address)
+        download(mm_asn_database_raw, asn_url, source_host=source_address)
         return True
     except:
         return False
diff --git a/src/helpers/geoip-update.py b/src/helpers/geoip-update.py
index 1879ab7ef..ccdcfb454 100755
--- a/src/helpers/geoip-update.py
+++ b/src/helpers/geoip-update.py
@@ -27,6 +27,7 @@ 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
+from vyos.utils.process import run
 
 def get_config(conf):
     return (
@@ -41,6 +42,7 @@ def get_config(conf):
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("--init", help="Initialise", action="store_true")
+    parser.add_argument('--download', help="Download", action="store_true")
     args = parser.parse_args()
 
     if args.init:
@@ -56,31 +58,39 @@ if __name__ == '__main__':
 
     options, firewall, policy = get_config(conf)
 
+    src_addr = options.get('source_address', '')
+    vrf = options.get('vrf', None)
+
+    if args.download or not vrf:
+        if options['provider'] == 'db-ip':
+            print('Downloading latest DB-IP database...')
+            if not geoip_download_dbip(source_address=src_addr):
+                print('Failed to download, 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('Downloading latest MaxMind database...')
+            if not geoip_download_maxmind(account_id, license_key, lite, source_address=src_addr):
+                print('Failed to download, aborting.')
+                sys.exit(1)
+        if args.download:
+            sys.exit(0)
+    elif vrf:
+        run(['python3', __file__, '--download'], stdout=None, stderr=None, vrf=vrf)
+
     if not db_is_initialised():
         db_initialise()
 
+    print('Extracting database...')
     if options['provider'] == 'db-ip':
-        print('Downloading 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('Downloading 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)