geoip: T5746: Add GeoIP ASN support

2 files changed, 91 insertions, 32 deletions

diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index c655e2de8..c2e6f0d68 100755
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -210,7 +210,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:
@@ -228,7 +231,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,