geoip: T5746: Add GeoIP ASN support

6 files changed, 134 insertions, 35 deletions

diff --git a/interface-definitions/include/firewall/geoip.xml.i b/interface-definitions/include/firewall/geoip.xml.i
index 09b2da663..55765e6c5 100644
--- a/interface-definitions/include/firewall/geoip.xml.i
+++ b/interface-definitions/include/firewall/geoip.xml.i
@@ -4,6 +4,19 @@
     <help>GeoIP options - Data provided by DB-IP.com</help>
   </properties>
   <children>
+    <leafNode name="asn">
+      <properties>
+        <help>Autonomous system number</help>
+        <valueHelp>
+          <format>u32:1-4294967294</format>
+          <description>Autonomous system number</description>
+        </valueHelp>
+        <constraint>
+          <validator name="numeric" argument="--range 1-4294967294"/>
+        </constraint>
+        <multi />
+      </properties>
+    </leafNode>
     <leafNode name="country-code">
       <properties>
         <help>GeoIP country code</help>
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,
diff --git a/smoketest/scripts/cli/test_firewall.py b/smoketest/scripts/cli/test_firewall.py
index 224c29a5e..baafc9d6e 100755
--- a/smoketest/scripts/cli/test_firewall.py
+++ b/smoketest/scripts/cli/test_firewall.py
@@ -93,12 +93,19 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
         self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '2', 'source', 'geoip', 'country-code', 'de'])
         self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '2', 'source', 'geoip', 'country-code', 'fr'])
         self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '2', 'source', 'geoip', 'inverse-match'])
+        self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '3', 'action', 'drop'])
+        self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '3', 'source', 'geoip', 'asn', '13335'])
+        self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '4', 'action', 'accept'])
+        self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '4', 'source', 'geoip', 'asn', '15169'])
+        self.cli_set(['firewall', 'ipv4', 'name', 'smoketest', 'rule', '4', 'source', 'geoip', 'inverse-match'])
 
         self.cli_commit()
 
         nftables_search = [
             ['ip saddr @GEOIP_CC_name_smoketest_1', 'drop'],
-            ['ip saddr != @GEOIP_CC_name_smoketest_2', 'accept']
+            ['ip saddr != @GEOIP_CC_name_smoketest_2', 'accept'],
+            ['ip saddr @GEOIP_ASN_name_smoketest_3', 'drop'],
+            ['ip saddr != @GEOIP_ASN_name_smoketest_4', 'accept']
         ]
 
         # -t prevents 1000+ GeoIP elements being returned
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 4a2706a03..4b6a4d272 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -83,12 +83,18 @@ snmp_trap_name = 'mgmtEventTrap'
 def geoip_sets(firewall):
     out = {'name': [], 'ipv6_name': []}
 
-    for _, path in dict_search_recursive(firewall, 'geoip'):
+    for _, path in dict_search_recursive(firewall, 'country_code'):
         if (path[0] == 'ipv4'):
             out['name'].append(f'GEOIP_CC_{path[1]}_{path[2]}_{path[4]}')
         elif (path[0] == 'ipv6'):
             out['ipv6_name'].append(f'GEOIP_CC6_{path[1]}_{path[2]}_{path[4]}')
 
+    for _, path in dict_search_recursive(firewall, 'asn'):
+        if (path[0] == 'ipv4'):
+            out['name'].append(f'GEOIP_ASN_{path[1]}_{path[2]}_{path[4]}')
+        elif (path[0] == 'ipv6'):
+            out['ipv6_name'].append(f'GEOIP_ASN6_{path[1]}_{path[2]}_{path[4]}')
+
     return out
 
 def geoip_updated(conf):
@@ -366,6 +372,10 @@ def verify_rule(firewall, family, hook, priority, rule_id, rule_conf):
             if len({'address', 'fqdn', 'geoip'} & set(side_conf)) > 1:
                 raise ConfigError('Only one of address, fqdn or geoip can be specified')
 
+            if 'geoip' in side_conf:
+                if len({'asn', 'country_code'} & set(side_conf['geoip'])) > 1:
+                    raise ConfigError('Only one of asn or country-code can be specified')
+
             if 'group' in side_conf:
                 if len({'address_group', 'network_group', 'domain_group', 'remote_group'} & set(side_conf['group'])) > 1:
                     raise ConfigError('Only one address-group, network-group, remote-group or domain-group can be specified')
diff --git a/src/conf_mode/policy_route.py b/src/conf_mode/policy_route.py
index 3cfdad913..64915632c 100755
--- a/src/conf_mode/policy_route.py
+++ b/src/conf_mode/policy_route.py
@@ -61,12 +61,18 @@ def geoip_updated(conf):
 def geoip_sets(policy):
     out = {'name': [], 'ipv6_name': []}
 
-    for _, path in dict_search_recursive(policy, 'geoip'):
+    for _, path in dict_search_recursive(policy, 'country_code'):
         if (path[0] == 'route'):
             out['name'].append(f'GEOIP_CC_{path[0]}_{path[1]}_{path[3]}')
         elif (path[0] == 'route6'):
             out['ipv6_name'].append(f'GEOIP_CC6_{path[0]}_{path[1]}_{path[3]}')
 
+    for _, path in dict_search_recursive(policy, 'asn'):
+        if (path[0] == 'route'):
+            out['name'].append(f'GEOIP_ASN_{path[0]}_{path[1]}_{path[3]}')
+        elif (path[0] == 'route6'):
+            out['ipv6_name'].append(f'GEOIP_ASN6_{path[0]}_{path[1]}_{path[3]}')
+
     return out
 
 def get_config(config=None):
@@ -141,6 +147,10 @@ def verify_rule(policy, name, rule_conf, ipv6, rule_id):
         if side in rule_conf:
             side_conf = rule_conf[side]
 
+            if 'geoip' in side_conf:
+                if len({'asn', 'country_code'} & set(side_conf['geoip'])) > 1:
+                    raise ConfigError('Only one of asn or country-code can be specified')
+
             if 'group' in side_conf:
                 if len({'address_group', 'domain_group', 'network_group'} & set(side_conf['group'])) > 1:
                     raise ConfigError('Only one address-group, domain-group or network-group can be specified')