diff options
Diffstat (limited to 'python')
| -rwxr-xr-x | python/vyos/firewall.py | 152 | ||||
| -rw-r--r-- | python/vyos/geoip.py | 258 |
2 files changed, 258 insertions, 152 deletions
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py index b136b6fca..0025ed98a 100755 --- a/python/vyos/firewall.py +++ b/python/vyos/firewall.py @@ -12,24 +12,16 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see <http://www.gnu.org/licenses/>. -import csv -import gzip -import os import re -from pathlib import Path from socket import AF_INET from socket import AF_INET6 from socket import getaddrinfo -from time import strftime -from vyos.remote import download from vyos.template import is_ipv4 -from vyos.template import render from vyos.utils.dict import dict_search_args from vyos.utils.dict import dict_search_recursive from vyos.utils.process import cmd -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 @@ -691,147 +683,3 @@ def parse_time(time): out_days = [f'"{day}"' for day in days if day[0] != '!'] out.append(f'day {{{",".join(out_days)}}}') return " ".join(out) - -# GeoIP - -nftables_geoip_conf = '/run/nftables-geoip.conf' -geoip_database = '/usr/share/vyos-geoip/dbip-country-lite.csv.gz' -geoip_lock_file = '/run/vyos-geoip.lock' - -def geoip_load_data(codes=[]): - data = None - - if not os.path.exists(geoip_database): - return [] - - try: - with gzip.open(geoip_database, mode='rt') as csv_fh: - reader = csv.reader(csv_fh) - out = [] - for start, end, code in reader: - if code.lower() in codes: - out.append([start, end, code.lower()]) - return out - except: - print('Error: Failed to open GeoIP database') - return [] - -def geoip_download_data(): - url = 'https://download.db-ip.com/free/dbip-country-lite-{}.csv.gz'.format(strftime("%Y-%m")) - try: - dirname = os.path.dirname(geoip_database) - if not os.path.exists(dirname): - os.mkdir(dirname) - - download(geoip_database, url) - print("Downloaded GeoIP database") - return True - except: - print("Error: Failed to download GeoIP database") - return False - -class GeoIPLock(object): - def __init__(self, file): - self.file = file - - def __enter__(self): - if os.path.exists(self.file): - return False - - Path(self.file).touch() - return True - - def __exit__(self, exc_type, exc_value, tb): - os.unlink(self.file) - -def geoip_update(firewall=None, policy=None, force=False): - with GeoIPLock(geoip_lock_file) as lock: - if not lock: - print("Script is already running") - return False - - if not firewall and not policy: - print("Firewall and policy are not configured") - return True - - if not os.path.exists(geoip_database): - if not geoip_download_data(): - return False - elif force: - geoip_download_data() - - ipv4_codes = {} - ipv6_codes = {} - - ipv4_sets = {} - ipv6_sets = {} - - ipv4_codes_policy = {} - ipv6_codes_policy = {} - - ipv4_sets_policy = {} - ipv6_sets_policy = {} - - # Map country codes to set names - if firewall: - for codes, path in dict_search_recursive(firewall, 'country_code'): - set_name = f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}' - if ( path[0] == 'ipv4'): - for code in codes: - ipv4_codes.setdefault(code, []).append(set_name) - elif ( path[0] == 'ipv6' ): - set_name = f'GEOIP_CC6_{path[1]}_{path[2]}_{path[4]}' - for code in codes: - ipv6_codes.setdefault(code, []).append(set_name) - - if policy: - for codes, path in dict_search_recursive(policy, 'country_code'): - set_name = f'GEOIP_CC_{path[0]}_{path[1]}_{path[3]}' - if ( path[0] == 'route'): - for code in codes: - ipv4_codes_policy.setdefault(code, []).append(set_name) - elif ( path[0] == 'route6' ): - set_name = f'GEOIP_CC6_{path[0]}_{path[1]}_{path[3]}' - for code in codes: - ipv6_codes_policy.setdefault(code, []).append(set_name) - - if not ipv4_codes and not ipv6_codes and not ipv4_codes_policy and not ipv6_codes_policy: - if force: - print("GeoIP not in use by firewall and policy") - return True - - geoip_data = geoip_load_data([*ipv4_codes, *ipv6_codes, *ipv4_codes_policy, *ipv6_codes_policy]) - - # Iterate IP blocks to assign to sets - for start, end, code in geoip_data: - ipv4 = is_ipv4(start) - if code in ipv4_codes and ipv4: - ip_range = f'{start}-{end}' if start != end else start - for setname in ipv4_codes[code]: - ipv4_sets.setdefault(setname, []).append(ip_range) - if code in ipv4_codes_policy and ipv4: - ip_range = f'{start}-{end}' if start != end else start - for setname in ipv4_codes_policy[code]: - ipv4_sets_policy.setdefault(setname, []).append(ip_range) - if code in ipv6_codes and not ipv4: - ip_range = f'{start}-{end}' if start != end else start - for setname in ipv6_codes[code]: - ipv6_sets.setdefault(setname, []).append(ip_range) - if code in ipv6_codes_policy and not ipv4: - ip_range = f'{start}-{end}' if start != end else start - for setname in ipv6_codes_policy[code]: - ipv6_sets_policy.setdefault(setname, []).append(ip_range) - - render(nftables_geoip_conf, 'firewall/nftables-geoip-update.j2', { - 'ipv4_sets': ipv4_sets, - 'ipv6_sets': ipv6_sets, - 'ipv4_sets_policy': ipv4_sets_policy, - 'ipv6_sets_policy': ipv6_sets_policy, - }) - - result = run(f'nft --file {nftables_geoip_conf}') - if result != 0: - print('Error: GeoIP failed to update firewall/policy') - return False - - return True diff --git a/python/vyos/geoip.py b/python/vyos/geoip.py new file mode 100644 index 000000000..9a2b40899 --- /dev/null +++ b/python/vyos/geoip.py @@ -0,0 +1,258 @@ + +import csv +import gzip +import os +import sqlite3 +import zipfile + +from io import TextIOWrapper +from pathlib import Path +from time import strftime + +from vyos.remote import download +from vyos.template import is_ipv4, render +from vyos.utils.dict import dict_search_recursive +from vyos.utils.process import run + +nftables_geoip_conf = '/run/nftables-geoip.conf' +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' + +# Raw data + +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(dbip_database_raw) + if not os.path.exists(dirname): + os.mkdir(dirname) + + 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 + +# VyOS database + +def db_is_initialised(): + if not os.path.exists(geoip_database_path): + return False + + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() + cur.execute("PRAGMA table_info(geoip_ranges);") + rows = cur.fetchall() + return len(rows) > 0 + +def db_initialise(): + dirname = os.path.dirname(geoip_database_path) + if not os.path.exists(dirname): + os.mkdir(dirname) + + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() + cur.execute(""" + CREATE TABLE IF NOT EXISTS geoip_ranges ( + country_code TEXT NOT NULL, + range TEXT NOT NULL, + version INT NOT NULL + ) + """) + cur.execute('CREATE INDEX IF NOT EXISTS idx_cc_version ON geoip_ranges(country_code, version)') + conn.commit() + +def db_import_dbip_ranges(replace=True, delete_file=False): + if not os.path.exists(dbip_database_raw): + return False + + if not os.path.exists(geoip_database_path): + return False + + try: + with gzip.open(dbip_database_raw, mode='rt') as csv_fh: + reader = csv.reader(csv_fh) + + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() + + if replace: + cur.execute('DELETE FROM geoip_ranges') + + for start, end, code in reader: + version = 4 if is_ipv4(start) else 6 + cur.execute('INSERT INTO geoip_ranges (country_code, range, version) VALUES (?, ?, ?)', (code.lower(), f'{start}-{end}', version)) + conn.commit() + + if delete_file: + 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: + return False + +def db_return_ranges(codes, version): + out = [] + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() + ph = ','.join(['?'] * len(codes)) + for row in cur.execute(f'SELECT range FROM geoip_ranges WHERE version = ? AND country_code IN ({ph})', [version, *codes]): + out.append(row[0]) + return out + +# Update + +def geoip_refresh(): + with GeoIPLock(geoip_lock_file) as lock: + if not lock: + return True + + if not os.path.exists(nftables_geoip_conf): + return False + + result = run(f'nft --file {nftables_geoip_conf}') + if result != 0: + return False + + return True + +def geoip_update(firewall=None, policy=None): + with GeoIPLock(geoip_lock_file) as lock: + if not lock: + print("Script is already running") + return False + + if not firewall and not policy: + print("Firewall and policy are not configured") + return True + + if not os.path.exists(geoip_database_path): + print("Running one-time database initialisation") + db_initialise() + db_import_dbip_ranges() + + firewall_sets = {'v4': {}, 'v6': {}} + policy_sets = {'v4': {}, 'v6': {}} + + if firewall: + for codes, path in dict_search_recursive(firewall, 'country_code'): + version = 6 if path[0] == 'ipv6' else 4 + vprefix = '6' if version == 6 else '' + set_name = f'GEOIP_CC{vprefix}_{path[1]}_{path[2]}_{path[4]}' + firewall_sets[f'v{version}'][set_name] = db_return_ranges(codes, version) + + if policy: + for codes, path in dict_search_recursive(policy, 'country_code'): + version = 6 if path[0] == 'route6' else 4 + vprefix = '6' if version == 6 else '' + set_name = f'GEOIP_CC{vprefix}_{path[0]}_{path[1]}_{path[3]}' + policy_sets[f'v{version}'][set_name] = db_return_ranges(codes, version) + + render(nftables_geoip_conf, 'firewall/nftables-geoip-update.j2', { + 'firewall_sets': firewall_sets, + 'policy_sets': policy_sets + }) + + result = run(f'nft --file {nftables_geoip_conf}') + if result != 0: + print('Error: GeoIP failed to update firewall and/or policy') + return False + + return True + +# Utility + +class GeoIPLock(object): + def __init__(self, file): + self.file = file + + def __enter__(self): + if os.path.exists(self.file): + return False + + Path(self.file).touch() + return True + + def __exit__(self, exc_type, exc_value, tb): + os.unlink(self.file) |