geoip: T5746: Add GeoIP ASN support

2 files changed, 22 insertions, 2 deletions

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')