diff options
| -rw-r--r-- | interface-definitions/include/firewall/global-options.xml.i | 2 | ||||
| -rw-r--r-- | op-mode-definitions/geoip.xml.in | 11 | ||||
| -rw-r--r-- | python/vyos/geoip.py | 12 | ||||
| -rwxr-xr-x | src/helpers/geoip-update.py | 42 |
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) |