diff options
Diffstat (limited to 'python')
| -rwxr-xr-x | python/vyos/firewall.py | 7 | ||||
| -rw-r--r-- | python/vyos/geoip.py | 116 |
2 files changed, 91 insertions, 32 deletions
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py index c4f6d0936..e33ee5b12 100755 --- a/python/vyos/firewall.py +++ b/python/vyos/firewall.py @@ -213,7 +213,10 @@ def parse_rule(rule_conf, hook, fw_name, rule_id, ip_name): hook_name = f'name{def_suffix}' output.append(f'{ip_name} {prefix}addr {operator} @FQDN_{hook_name}_{fw_name}_{rule_id}_{prefix}') - if dict_search_args(side_conf, 'geoip', 'country_code'): + country_code = dict_search_args(side_conf, 'geoip', 'country_code') + asn = dict_search_args(side_conf, 'geoip', 'asn') + if country_code or asn: + geoip_prefix = 'CC' if country_code else 'ASN' operator = '' hook_name = '' if dict_search_args(side_conf, 'geoip', 'inverse_match') != None: @@ -231,7 +234,7 @@ def parse_rule(rule_conf, hook, fw_name, rule_id, ip_name): # for policy if hook == 'route' or hook == 'route6': hook_name = hook - output.append(f'{ip_name} {prefix}addr {operator} @GEOIP_CC{def_suffix}_{hook_name}_{fw_name}_{rule_id}') + output.append(f'{ip_name} {prefix}addr {operator} @GEOIP_{geoip_prefix}{def_suffix}_{hook_name}_{fw_name}_{rule_id}') if 'mac_address' in side_conf: suffix = side_conf["mac_address"] diff --git a/python/vyos/geoip.py b/python/vyos/geoip.py index db2313ab9..0f3e7f01c 100644 --- a/python/vyos/geoip.py +++ b/python/vyos/geoip.py @@ -16,7 +16,9 @@ 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' +dbip_asn_database_raw = '/usr/share/vyos-geoip/dbip-asn-lite.csv.gz' mm_database_raw = '/usr/share/vyos-geoip/maxmind-country.zip' +mm_asn_database_raw = '/usr/share/vyos-geoip/maxmind-asn.zip' geoip_database_path = '/var/cache/vyos/geoip-lookup.db' geoip_lock_file = '/var/lock/vyos-geoip.lock' @@ -24,12 +26,14 @@ 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")) + asn_url = 'https://download.db-ip.com/free/dbip-asn-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) + download(dbip_asn_database_raw, asn_url) return True except: return False @@ -37,12 +41,14 @@ def geoip_download_dbip(): 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' + asn_url = f'https://{account_id}:{license_key}@download.maxmind.com/geoip/databases/{db_str}-ASN-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) + download(mm_asn_database_raw, asn_url) return True except: return False @@ -68,12 +74,14 @@ def db_initialise(): cur = conn.cursor() cur.execute(""" CREATE TABLE IF NOT EXISTS geoip_ranges ( - country_code TEXT NOT NULL, + country_code TEXT, + asn INT, range TEXT NOT NULL, version INT NOT NULL ) """) cur.execute('CREATE INDEX IF NOT EXISTS idx_cc_version ON geoip_ranges(country_code, version)') + cur.execute('CREATE INDEX IF NOT EXISTS idx_asn_version ON geoip_ranges(asn, version)') conn.commit() def db_import_dbip_ranges(replace=True, delete_file=False): @@ -84,22 +92,29 @@ def db_import_dbip_ranges(replace=True, delete_file=False): 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() - with sqlite3.connect(geoip_database_path) as conn: - cur = conn.cursor() - - if replace: - cur.execute('DELETE FROM geoip_ranges') + if replace: + cur.execute('DELETE FROM geoip_ranges') + with gzip.open(dbip_database_raw, mode='rt') as csv_fh: + reader = csv.reader(csv_fh) 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() + + with gzip.open(dbip_asn_database_raw, mode='rt') as csv_fh: + reader = csv.reader(csv_fh) + for start, end, asn, _ in reader: + version = 4 if is_ipv4(start) else 6 + cur.execute('INSERT INTO geoip_ranges (asn, range, version) VALUES (?, ?, ?)', (asn, f'{start}-{end}', version)) + + conn.commit() if delete_file: os.unlink(dbip_database_raw) + os.unlink(dbip_asn_database_raw) return True except: @@ -116,28 +131,28 @@ def db_import_maxmind_ranges(replace=True, delete_file=False): 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' + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() - 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 = {} + if replace: + cur.execute('DELETE FROM geoip_ranges') - 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) + 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' - for row in reader: - id = row['geoname_id'] - locations_map[id] = row['country_iso_code'] + 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 sqlite3.connect(geoip_database_path) as conn: - cur = conn.cursor() + 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) - if replace: - cur.execute('DELETE FROM geoip_ranges') + for row in reader: + id = row['geoname_id'] + locations_map[id] = row['country_iso_code'] with zip_fh.open(ipv4_file) as raw_csv_fh: with TextIOWrapper(raw_csv_fh, encoding='utf-8') as csv_fh: @@ -163,16 +178,36 @@ def db_import_maxmind_ranges(replace=True, delete_file=False): code = locations_map[id] cur.execute('INSERT INTO geoip_ranges (country_code, range, version) VALUES (?, ?, 6)', (code.lower(), row['network'])) - conn.commit() + with zipfile.ZipFile(mm_asn_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}-ASN-Blocks-IPv4.csv' + ipv6_file = f'{directory}/{prefix}-ASN-Blocks-IPv6.csv' + + 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: + cur.execute('INSERT INTO geoip_ranges (asn, range, version) VALUES (?, ?, 4)', (row['autonomous_system_number'], 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: + cur.execute('INSERT INTO geoip_ranges (asn, range, version) VALUES (?, ?, 6)', (row['autonomous_system_number'], row['network'])) + + conn.commit() if delete_file: os.unlink(mm_database_raw) + os.unlink(mm_asn_database_raw) return True except: return False -def db_return_ranges(codes, version): +def db_return_cc_ranges(codes, version): out = [] with sqlite3.connect(geoip_database_path) as conn: cur = conn.cursor() @@ -181,6 +216,15 @@ def db_return_ranges(codes, version): out.append(row[0]) return out +def db_return_asn_ranges(asn, version): + out = [] + with sqlite3.connect(geoip_database_path) as conn: + cur = conn.cursor() + ph = ','.join(['?'] * len(asn)) + for row in cur.execute(f'SELECT range FROM geoip_ranges WHERE version = ? AND asn IN ({ph})', [version, *asn]): + out.append(row[0]) + return out + # Update def geoip_refresh(): @@ -223,7 +267,13 @@ def geoip_update(firewall=None, policy=None): 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) + firewall_sets[f'v{version}'][set_name] = db_return_cc_ranges(codes, version) + + for asns, path in dict_search_recursive(firewall, 'asn'): + version = 6 if path[0] == 'ipv6' else 4 + vprefix = '6' if version == 6 else '' + set_name = f'GEOIP_ASN{vprefix}_{path[1]}_{path[2]}_{path[4]}' + firewall_sets[f'v{version}'][set_name] = db_return_asn_ranges(asns, version) if policy: for codes, path in dict_search_recursive(policy, 'country_code'): @@ -233,7 +283,13 @@ def geoip_update(firewall=None, policy=None): 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) + policy_sets[f'v{version}'][set_name] = db_return_cc_ranges(codes, version) + + for asns, path in dict_search_recursive(policy, 'asn'): + version = 6 if path[0] == 'route6' else 4 + vprefix = '6' if version == 6 else '' + set_name = f'GEOIP_ASN{vprefix}_{path[0]}_{path[1]}_{path[3]}' + policy_sets[f'v{version}'][set_name] = db_return_asn_ranges(asns, version) render( nftables_geoip_conf, |