diff options
Diffstat (limited to 'python')
| -rw-r--r-- | python/vyos/geoip.py | 95 |
1 files changed, 89 insertions, 6 deletions
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: |