From 34db435e7a74ee8509777802e03927de2dd57627 Mon Sep 17 00:00:00 2001
From: sarthurdev <965089+sarthurdev@users.noreply.github.com>
Date: Mon, 13 Jun 2022 01:45:06 +0200
Subject: firewall: T4147: Use named sets for firewall groups

* Refactor nftables clean-up code
* Adds policy route test for using firewall groups
---
 data/templates/firewall/nftables-defines.j2 |  70 +++++++++++++----
 data/templates/firewall/nftables-policy.j2  |   7 +-
 data/templates/firewall/nftables.j2         |  10 ++-
 python/vyos/firewall.py                     |   8 +-
 python/vyos/template.py                     |  39 ++++------
 smoketest/scripts/cli/test_firewall.py      |  93 ++++++++++------------
 smoketest/scripts/cli/test_policy_route.py  |  71 ++++++++++-------
 src/conf_mode/firewall.py                   | 117 +++++++++++++++++-----------
 src/conf_mode/policy-route.py               |  72 ++++++++++++-----
 9 files changed, 293 insertions(+), 194 deletions(-)

diff --git a/data/templates/firewall/nftables-defines.j2 b/data/templates/firewall/nftables-defines.j2
index 12146879d..97fc123d5 100644
--- a/data/templates/firewall/nftables-defines.j2
+++ b/data/templates/firewall/nftables-defines.j2
@@ -1,38 +1,76 @@
+{% macro groups(group, is_ipv6) %}
 {% if group is vyos_defined %}
-{%     if group.address_group is vyos_defined %}
-{%         for group_name, group_conf in group.address_group | sort_nested_groups %}
+{%     set ip_type = 'ipv6_addr' if is_ipv6 else 'ipv4_addr' %}
+{%     if group.address_group is vyos_defined and not is_ipv6 %}
+{%         for group_name, group_conf in group.address_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define A_{{ group_name }} = { {{ group_conf.address | nft_nested_group(includes, 'A_') | join(",") }} }
+    set A_{{ group_name }} {
+        type {{ ip_type }}
+        flags interval
+{%             if group_conf.address is vyos_defined or includes %}
+        elements = { {{ group_conf.address | nft_nested_group(includes, group.address_group, 'address') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
-{%     if group.ipv6_address_group is vyos_defined %}
-{%         for group_name, group_conf in group.ipv6_address_group | sort_nested_groups %}
+{%     if group.ipv6_address_group is vyos_defined and is_ipv6 %}
+{%         for group_name, group_conf in group.ipv6_address_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define A6_{{ group_name }} = { {{ group_conf.address | nft_nested_group(includes, 'A6_') | join(",") }} }
+    set A6_{{ group_name }} {
+        type {{ ip_type }}
+        flags interval
+{%             if group_conf.address is vyos_defined or includes %}
+        elements = { {{ group_conf.address | nft_nested_group(includes, group.ipv6_address_group, 'address') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
 {%     if group.mac_group is vyos_defined %}
-{%         for group_name, group_conf in group.mac_group | sort_nested_groups %}
+{%         for group_name, group_conf in group.mac_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define M_{{ group_name }} = { {{ group_conf.mac_address | nft_nested_group(includes, 'M_') | join(",") }} }
+    set M_{{ group_name }} {
+        type ether_addr
+{%             if group_conf.mac_address is vyos_defined or includes %}
+        elements = { {{ group_conf.mac_address | nft_nested_group(includes, group.mac_group, 'mac_address') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
-{%     if group.network_group is vyos_defined %}
-{%         for group_name, group_conf in group.network_group | sort_nested_groups %}
+{%     if group.network_group is vyos_defined and not is_ipv6 %}
+{%         for group_name, group_conf in group.network_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define N_{{ group_name }} = { {{ group_conf.network | nft_nested_group(includes, 'N_') | join(",") }} }
+    set N_{{ group_name }} {
+        type {{ ip_type }}
+        flags interval
+{%             if group_conf.network is vyos_defined or includes %}
+        elements = { {{ group_conf.network | nft_nested_group(includes, group.network_group, 'network') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
-{%     if group.ipv6_network_group is vyos_defined %}
-{%         for group_name, group_conf in group.ipv6_network_group | sort_nested_groups %}
+{%     if group.ipv6_network_group is vyos_defined and is_ipv6 %}
+{%         for group_name, group_conf in group.ipv6_network_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define N6_{{ group_name }} = { {{ group_conf.network | nft_nested_group(includes, 'N6_') | join(",") }} }
+    set N6_{{ group_name }} {
+        type {{ ip_type }}
+        flags interval
+{%             if group_conf.network is vyos_defined or includes %}
+        elements = { {{ group_conf.network | nft_nested_group(includes, group.ipv6_network_group, 'network') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
 {%     if group.port_group is vyos_defined %}
-{%         for group_name, group_conf in group.port_group | sort_nested_groups %}
+{%         for group_name, group_conf in group.port_group.items() %}
 {%             set includes = group_conf.include if group_conf.include is vyos_defined else [] %}
-define P_{{ group_name }} = { {{ group_conf.port | nft_nested_group(includes, 'P_') | join(",") }} }
+    set P_{{ group_name }} {
+        type inet_service
+        flags interval
+{%             if group_conf.port is vyos_defined or includes %}
+        elements = { {{ group_conf.port | nft_nested_group(includes, group.port_group, 'port') | join(",") }} }
+{%             endif %}
+    }
 {%         endfor %}
 {%     endif %}
 {% endif %}
+{% endmacro %}
diff --git a/data/templates/firewall/nftables-policy.j2 b/data/templates/firewall/nftables-policy.j2
index 0154c9f7e..281525407 100644
--- a/data/templates/firewall/nftables-policy.j2
+++ b/data/templates/firewall/nftables-policy.j2
@@ -1,13 +1,13 @@
 #!/usr/sbin/nft -f
 
+{% import 'firewall/nftables-defines.j2' as group_tmpl %}
+
 {% if cleanup_commands is vyos_defined %}
 {%     for command in cleanup_commands %}
 {{ command }}
 {%     endfor %}
 {% endif %}
 
-include "/run/nftables_defines.conf"
-
 table ip mangle {
 {% if first_install is vyos_defined %}
     chain VYOS_PBR_PREROUTING {
@@ -29,6 +29,8 @@ table ip mangle {
     }
 {%     endfor %}
 {% endif %}
+
+{{ group_tmpl.groups(firewall_group, False) }}
 }
 
 table ip6 mangle {
@@ -52,4 +54,5 @@ table ip6 mangle {
     }
 {%     endfor %}
 {% endif %}
+{{ group_tmpl.groups(firewall_group, True) }}
 }
diff --git a/data/templates/firewall/nftables.j2 b/data/templates/firewall/nftables.j2
index 961b83301..ca24b7db2 100644
--- a/data/templates/firewall/nftables.j2
+++ b/data/templates/firewall/nftables.j2
@@ -1,13 +1,13 @@
 #!/usr/sbin/nft -f
 
+{% import 'firewall/nftables-defines.j2' as group_tmpl %}
+
 {% if cleanup_commands is vyos_defined %}
 {%     for command in cleanup_commands %}
 {{ command }}
 {%     endfor %}
 {% endif %}
 
-include "/run/nftables_defines.conf"
-
 table ip filter {
 {% if first_install is vyos_defined %}
     chain VYOS_FW_FORWARD {
@@ -69,6 +69,9 @@ table ip filter {
 {%         endfor %}
 {%     endif %}
 {% endif %}
+
+{{ group_tmpl.groups(group, False) }}
+
 {% if state_policy is vyos_defined %}
     chain VYOS_STATE_POLICY {
 {%     if state_policy.established is vyos_defined %}
@@ -138,6 +141,9 @@ table ip6 filter {
 {%         endfor %}
 {%     endif %}
 {% endif %}
+
+{{ group_tmpl.groups(group, True) }}
+
 {% if state_policy is vyos_defined %}
     chain VYOS_STATE_POLICY6 {
 {%     if state_policy.established is vyos_defined %}
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index a61d0a9f8..f8f913944 100644
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -192,7 +192,7 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
                     if group_name[0] == '!':
                         operator = '!='
                         group_name = group_name[1:]
-                    output.append(f'{ip_name} {prefix}addr {operator} $A{def_suffix}_{group_name}')
+                    output.append(f'{ip_name} {prefix}addr {operator} @A{def_suffix}_{group_name}')
                 # Generate firewall group domain-group
                 elif 'domain_group' in group:
                     group_name = group['domain_group']
@@ -207,14 +207,14 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
                     if group_name[0] == '!':
                         operator = '!='
                         group_name = group_name[1:]
-                    output.append(f'{ip_name} {prefix}addr {operator} $N{def_suffix}_{group_name}')
+                    output.append(f'{ip_name} {prefix}addr {operator} @N{def_suffix}_{group_name}')
                 if 'mac_group' in group:
                     group_name = group['mac_group']
                     operator = ''
                     if group_name[0] == '!':
                         operator = '!='
                         group_name = group_name[1:]
-                    output.append(f'ether {prefix}addr {operator} $M_{group_name}')
+                    output.append(f'ether {prefix}addr {operator} @M_{group_name}')
                 if 'port_group' in group:
                     proto = rule_conf['protocol']
                     group_name = group['port_group']
@@ -227,7 +227,7 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
                         operator = '!='
                         group_name = group_name[1:]
 
-                    output.append(f'{proto} {prefix}port {operator} $P_{group_name}')
+                    output.append(f'{proto} {prefix}port {operator} @P_{group_name}')
 
     if 'log' in rule_conf and rule_conf['log'] == 'enable':
         action = rule_conf['action'] if 'action' in rule_conf else 'accept'
diff --git a/python/vyos/template.py b/python/vyos/template.py
index 3feda47c8..eb7f06480 100644
--- a/python/vyos/template.py
+++ b/python/vyos/template.py
@@ -592,37 +592,24 @@ def nft_intra_zone_action(zone_conf, ipv6=False):
     return 'return'
 
 @register_filter('nft_nested_group')
-def nft_nested_group(out_list, includes, prefix):
+def nft_nested_group(out_list, includes, groups, key):
     if not vyos_defined(out_list):
         out_list = []
-    for name in includes:
-        out_list.append(f'${prefix}{name}')
-    return out_list
-
-@register_filter('sort_nested_groups')
-def sort_nested_groups(groups):
-    seen = []
-    out = {}
-
-    def include_iterate(group_name):
-        group = groups[group_name]
-        if 'include' not in group:
-            if group_name not in out:
-                out[group_name] = groups[group_name]
-            return
 
-        for inc_group_name in group['include']:
-            if inc_group_name not in seen:
-                seen.append(inc_group_name)
-                include_iterate(inc_group_name)
+    def add_includes(name):
+        if key in groups[name]:
+            for item in groups[name][key]:
+                if item in out_list:
+                    continue
+                out_list.append(item)
 
-        if group_name not in out:
-            out[group_name] = groups[group_name]
+        if 'include' in groups[name]:
+            for name_inc in groups[name]['include']:
+                add_includes(name_inc)
 
-    for group_name in groups:
-        include_iterate(group_name)
-
-    return out.items()
+    for name in includes:
+        add_includes(name)
+    return out_list
 
 @register_test('vyos_defined')
 def vyos_defined(value, test_value=None, var_type=None):
diff --git a/smoketest/scripts/cli/test_firewall.py b/smoketest/scripts/cli/test_firewall.py
index 998f1b3f3..8b8c27a9f 100755
--- a/smoketest/scripts/cli/test_firewall.py
+++ b/smoketest/scripts/cli/test_firewall.py
@@ -57,6 +57,29 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
         self.cli_delete(['firewall'])
         self.cli_commit()
 
+        # Verify chains/sets are cleaned up from nftables
+        nftables_search = [
+            ['set M_smoketest_mac'],
+            ['set N_smoketest_network'],
+            ['set P_smoketest_port'],
+            ['set smoketest_domain'],
+            ['set RECENT_smoketest_4'],
+            ['chain NAME_smoketest']
+        ]
+
+        self.verify_nftables(nftables_search, 'ip filter', inverse=True)
+
+    def verify_nftables(self, nftables_search, table, inverse=False):
+        nftables_output = cmd(f'sudo nft list table {table}')
+
+        for search in nftables_search:
+            matched = False
+            for line in nftables_output.split("\n"):
+                if all(item in line for item in search):
+                    matched = True
+                    break
+            self.assertTrue(not matched if inverse else matched, msg=search)
+
     def test_groups(self):
         hostmap_path = ['system', 'static-host-mapping', 'host-name']
         example_org = ['192.0.2.8', '192.0.2.10', '192.0.2.11']
@@ -88,23 +111,17 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
         self.cli_commit()
         nftables_search = [
             ['iifname "eth0"', 'jump NAME_smoketest'],
-            ['ip saddr { 172.16.99.0/24 }', 'ip daddr 172.16.10.10', 'th dport { 53, 123 }', 'return'],
-            ['ether saddr { 00:01:02:03:04:05 }', 'return'],
+            ['ip saddr @N_smoketest_network', 'ip daddr 172.16.10.10', 'th dport @P_smoketest_port', 'return'],
+            ['elements = { 172.16.99.0/24 }'],
+            ['elements = { 53, 123 }'],
+            ['ether saddr @M_smoketest_mac', 'return'],
+            ['elements = { 00:01:02:03:04:05 }'],
             ['set smoketest_domain'],
             ['elements = { 192.0.2.5, 192.0.2.8,'],
             ['192.0.2.10, 192.0.2.11 }'],
             ['ip saddr @smoketest_domain', 'return']
         ]
-
-        nftables_output = cmd('sudo nft list table ip filter')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched, msg=search)
+        self.verify_nftables(nftables_search, 'ip filter')
 
         self.cli_delete(['system', 'static-host-mapping'])
         self.cli_commit()
@@ -134,18 +151,12 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
 
         nftables_search = [
             ['iifname "eth0"', 'jump NAME_smoketest'],
-            ['ip saddr { 172.16.99.0/24, 172.16.101.0/24 }', 'th dport { 53, 123 }', 'return']
+            ['ip saddr @N_smoketest_network1', 'th dport @P_smoketest_port1', 'return'],
+            ['elements = { 172.16.99.0/24, 172.16.101.0/24 }'],
+            ['elements = { 53, 123 }']
         ]
 
-        nftables_output = cmd('sudo nft list table ip filter')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched, msg=search)
+        self.verify_nftables(nftables_search, 'ip filter')
 
     def test_basic_rules(self):
         self.cli_set(['firewall', 'name', 'smoketest', 'default-action', 'drop'])
@@ -169,6 +180,11 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'destination', 'port', '22'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'limit', 'rate', '5/minute'])
         self.cli_set(['firewall', 'name', 'smoketest', 'rule', '3', 'log', 'disable'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'action', 'drop'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'protocol', 'tcp'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'destination', 'port', '22'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'recent', 'count', '10'])
+        self.cli_set(['firewall', 'name', 'smoketest', 'rule', '4', 'recent', 'time', 'minute'])
 
         self.cli_set(['interfaces', 'ethernet', 'eth0', 'firewall', 'in', 'name', 'smoketest'])
 
@@ -179,18 +195,11 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
             ['saddr 172.16.20.10', 'daddr 172.16.10.10', 'log prefix "[smoketest-1-A]" level debug', 'ip ttl 15','return'],
             ['tcp flags & (syn | ack) == syn', 'tcp dport { 8888 }', 'log prefix "[smoketest-2-R]" level err', 'ip ttl > 102', 'reject'],
             ['tcp dport { 22 }', 'limit rate 5/minute', 'return'],
-            ['log prefix "[smoketest-default-D]"','smoketest default-action', 'drop']
+            ['log prefix "[smoketest-default-D]"','smoketest default-action', 'drop'],
+            ['tcp dport { 22 }', 'add @RECENT_smoketest_4 { ip saddr limit rate over 10/minute burst 10 packets }', 'drop']
         ]
 
-        nftables_output = cmd('sudo nft list table ip filter')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched, msg=search)
+        self.verify_nftables(nftables_search, 'ip filter')
 
     def test_basic_rules_ipv6(self):
         self.cli_set(['firewall', 'ipv6-name', 'v6-smoketest', 'default-action', 'drop'])
@@ -217,15 +226,7 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
             ['smoketest default-action', 'log prefix "[v6-smoketest-default-D]"', 'drop']
         ]
 
-        nftables_output = cmd('sudo nft list table ip6 filter')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched, msg=search)
+        self.verify_nftables(nftables_search, 'ip6 filter')
 
     def test_state_policy(self):
         self.cli_set(['firewall', 'state-policy', 'established', 'action', 'accept'])
@@ -273,15 +274,7 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
             ['smoketest default-action', 'drop']
         ]
 
-        nftables_output = cmd('sudo nft list table ip filter')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched, msg=search)
+        self.verify_nftables(nftables_search, 'ip filter')
 
     def test_sysfs(self):
         for name, conf in sysfs_config.items():
diff --git a/smoketest/scripts/cli/test_policy_route.py b/smoketest/scripts/cli/test_policy_route.py
index e2d70f289..534cfb082 100755
--- a/smoketest/scripts/cli/test_policy_route.py
+++ b/smoketest/scripts/cli/test_policy_route.py
@@ -47,6 +47,47 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
         self.cli_delete(['policy', 'route6'])
         self.cli_commit()
 
+        nftables_search = [
+            ['set N_smoketest_network'],
+            ['set N_smoketest_network1'],
+            ['chain VYOS_PBR_smoketest']
+        ]
+
+        self.verify_nftables(nftables_search, 'ip filter', inverse=True)
+
+    def verify_nftables(self, nftables_search, table, inverse=False):
+        nftables_output = cmd(f'sudo nft list table {table}')
+
+        for search in nftables_search:
+            matched = False
+            for line in nftables_output.split("\n"):
+                if all(item in line for item in search):
+                    matched = True
+                    break
+            self.assertTrue(not matched if inverse else matched, msg=search)
+
+    def test_pbr_group(self):
+        self.cli_set(['firewall', 'group', 'network-group', 'smoketest_network', 'network', '172.16.99.0/24'])
+        self.cli_set(['firewall', 'group', 'network-group', 'smoketest_network1', 'network', '172.16.101.0/24'])
+        self.cli_set(['firewall', 'group', 'network-group', 'smoketest_network1', 'include', 'smoketest_network'])
+
+        self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'source', 'group', 'network-group', 'smoketest_network'])
+        self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'destination', 'group', 'network-group', 'smoketest_network1'])
+        self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'set', 'mark', mark])
+
+        self.cli_set(['interfaces', 'ethernet', interface, 'policy', 'route', 'smoketest'])
+
+        self.cli_commit()
+
+        nftables_search = [
+            [f'iifname "{interface}"','jump VYOS_PBR_smoketest'],
+            ['ip daddr @N_smoketest_network1', 'ip saddr @N_smoketest_network'],
+        ]
+
+        self.verify_nftables(nftables_search, 'ip mangle')
+
+        self.cli_delete(['firewall'])
+
     def test_pbr_mark(self):
         self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'source', 'address', '172.16.20.10'])
         self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'destination', 'address', '172.16.10.10'])
@@ -63,15 +104,7 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
             ['ip daddr 172.16.10.10', 'ip saddr 172.16.20.10', 'meta mark set ' + mark_hex],
         ]
 
-        nftables_output = cmd('sudo nft list table ip mangle')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched)
+        self.verify_nftables(nftables_search, 'ip mangle')
 
     def test_pbr_table(self):
         self.cli_set(['policy', 'route', 'smoketest', 'rule', '1', 'protocol', 'tcp'])
@@ -97,15 +130,7 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
             ['tcp flags & (syn | ack) == syn', 'tcp dport { 8888 }', 'meta mark set ' + mark_hex]
         ]
 
-        nftables_output = cmd('sudo nft list table ip mangle')
-
-        for search in nftables_search:
-            matched = False
-            for line in nftables_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched)
+        self.verify_nftables(nftables_search, 'ip mangle')
 
         # IPv6
 
@@ -114,15 +139,7 @@ class TestPolicyRoute(VyOSUnitTestSHIM.TestCase):
             ['meta l4proto { tcp, udp }', 'th dport { 8888 }', 'meta mark set ' + mark_hex]
         ]
 
-        nftables6_output = cmd('sudo nft list table ip6 mangle')
-
-        for search in nftables6_search:
-            matched = False
-            for line in nftables6_output.split("\n"):
-                if all(item in line for item in search):
-                    matched = True
-                    break
-            self.assertTrue(matched)
+        self.verify_nftables(nftables6_search, 'ip6 mangle')
 
         # IP rule fwmark -> table
 
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 46b8add59..78dffe9dd 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -92,11 +92,20 @@ valid_groups = [
     'port_group'
 ]
 
-group_types = [
-    'address_group', 'network_group', 'port_group',
-    'ipv6_address_group', 'ipv6_network_group'
+nested_group_types = [
+    'address_group', 'network_group', 'mac_group',
+    'port_group', 'ipv6_address_group', 'ipv6_network_group'
 ]
 
+group_set_prefix = {
+    'A_': 'address_group',
+    'A6_': 'ipv6_address_group',
+    'M_': 'mac_group',
+    'N_': 'network_group',
+    'N6_': 'ipv6_network_group',
+    'P_': 'port_group'
+}
+
 snmp_change_type = {
     'unknown': 0,
     'add': 1,
@@ -165,18 +174,20 @@ def geoip_updated(conf, firewall):
     updated = False
 
     for key, path in dict_search_recursive(firewall, 'geoip'):
+        set_name = f'GEOIP_CC_{path[1]}_{path[3]}'
         if path[0] == 'name':
-            out['name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+            out['name'].append(set_name)
         elif path[0] == 'ipv6_name':
-            out['ipv6_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+            out['ipv6_name'].append(set_name)
         updated = True
 
     if 'delete' in node_diff:
         for key, path in dict_search_recursive(node_diff['delete'], 'geoip'):
+            set_name = f'GEOIP_CC_{path[1]}_{path[3]}'
             if path[0] == 'name':
-                out['deleted_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+                out['deleted_name'].append(set_name)
             elif path[0] == 'ipv6-name':
-                out['deleted_ipv6_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+                out['deleted_ipv6_name'].append(set_name)
             updated = True
 
     if updated:
@@ -315,7 +326,7 @@ def verify(firewall):
             raise ConfigError(f'Firewall config-trap enabled but "service snmp trap-target" is not defined')
 
     if 'group' in firewall:
-        for group_type in group_types:
+        for group_type in nested_group_types:
             if group_type in firewall['group']:
                 groups = firewall['group'][group_type]
                 for group_name, group in groups.items():
@@ -352,62 +363,75 @@ def verify(firewall):
 
     return None
 
-def cleanup_rule(table, jump_chain):
-    commands = []
-    chains = nft_iface_chains if table == 'ip filter' else nft6_iface_chains
-    for chain in chains:
-        results = cmd(f'nft -a list chain {table} {chain}').split("\n")
-        for line in results:
-            if f'jump {jump_chain}' in line:
-                handle_search = re.search('handle (\d+)', line)
-                if handle_search:
-                    commands.append(f'delete rule {table} {chain} handle {handle_search[1]}')
-    return commands
-
 def cleanup_commands(firewall):
     commands = []
-    commands_end = []
+    commands_chains = []
+    commands_sets = []
     for table in ['ip filter', 'ip6 filter']:
+        name_node = 'name' if table == 'ip filter' else 'ipv6_name'
+        chain_prefix = NAME_PREFIX if table == 'ip filter' else NAME6_PREFIX
+        state_chain = 'VYOS_STATE_POLICY' if table == 'ip filter' else 'VYOS_STATE_POLICY6'
+        iface_chains = nft_iface_chains if table == 'ip filter' else nft6_iface_chains
+
         geoip_list = []
         if firewall['geoip_updated']:
             geoip_key = 'deleted_ipv6_name' if table == 'ip6 filter' else 'deleted_name'
             geoip_list = dict_search_args(firewall, 'geoip_updated', geoip_key) or []
-
-        state_chain = 'VYOS_STATE_POLICY' if table == 'ip filter' else 'VYOS_STATE_POLICY6'
+ 
         json_str = cmd(f'nft -t -j list table {table}')
         obj = loads(json_str)
+
         if 'nftables' not in obj:
             continue
+
         for item in obj['nftables']:
             if 'chain' in item:
                 chain = item['chain']['name']
-                if chain in ['VYOS_STATE_POLICY', 'VYOS_STATE_POLICY6']:
-                    if 'state_policy' not in firewall:
-                        commands.append(f'delete chain {table} {chain}')
-                    else:
-                        commands.append(f'flush chain {table} {chain}')
-                elif chain not in preserve_chains and not chain.startswith("VZONE"):
-                    if table == 'ip filter' and dict_search_args(firewall, 'name', chain.replace(NAME_PREFIX, "", 1)) != None:
-                        commands.append(f'flush chain {table} {chain}')
-                    elif table == 'ip6 filter' and dict_search_args(firewall, 'ipv6_name', chain.replace(NAME6_PREFIX, "", 1)) != None:
-                        commands.append(f'flush chain {table} {chain}')
-                    else:
-                        commands += cleanup_rule(table, chain)
-                        commands.append(f'delete chain {table} {chain}')
-            elif 'rule' in item:
+                if chain in preserve_chains or chain.startswith("VZONE"):
+                    continue
+
+                if chain == state_chain:
+                    command = 'delete' if 'state_policy' not in firewall else 'flush'
+                    commands_chains.append(f'{command} chain {table} {chain}')
+                elif dict_search_args(firewall, name_node, chain.replace(chain_prefix, "", 1)) != None:
+                    commands.append(f'flush chain {table} {chain}')
+                else:
+                    commands_chains.append(f'delete chain {table} {chain}')
+
+            if 'rule' in item:
                 rule = item['rule']
-                if rule['chain'] in ['VYOS_FW_FORWARD', 'VYOS_FW_OUTPUT', 'VYOS_FW_LOCAL', 'VYOS_FW6_FORWARD', 'VYOS_FW6_OUTPUT', 'VYOS_FW6_LOCAL']:
-                    if 'expr' in rule and any([True for expr in rule['expr'] if dict_search_args(expr, 'jump', 'target') == state_chain]):
-                        if 'state_policy' not in firewall:
-                            chain = rule['chain']
-                            handle = rule['handle']
+                chain = rule['chain']
+                handle = rule['handle']
+
+                if chain in iface_chains:
+                    target, _ = next(dict_search_recursive(rule['expr'], 'target'))
+
+                    if target == state_chain and 'state_policy' not in firewall:
+                        commands.append(f'delete rule {table} {chain} handle {handle}')
+
+                    if target.startswith(chain_prefix):
+                        if dict_search_args(firewall, name_node, target.replace(chain_prefix, "", 1)) == None:
                             commands.append(f'delete rule {table} {chain} handle {handle}')
-            elif 'set' in item:
+
+            if 'set' in item:
                 set_name = item['set']['name']
-                if set_name.startswith('GEOIP_CC_') and set_name not in geoip_list:
+
+                if set_name.startswith('GEOIP_CC_') and set_name in geoip_list:
+                    commands_sets.append(f'delete set {table} {set_name}')
                     continue
-                commands_end.append(f'delete set {table} {set_name}')
-    return commands + commands_end
+                
+                if set_name.startswith("RECENT_"):
+                    commands_sets.append(f'delete set {table} {set_name}')
+                    continue
+
+                for prefix, group_type in group_set_prefix.items():
+                    if set_name.startswith(prefix):
+                        group_name = set_name.replace(prefix, "", 1)
+                        if dict_search_args(firewall, 'group', group_type, group_name) != None:
+                            commands_sets.append(f'flush set {table} {set_name}')
+                        else:
+                            commands_sets.append(f'delete set {table} {set_name}')
+    return commands + commands_chains + commands_sets
 
 def generate(firewall):
     if not os.path.exists(nftables_conf):
@@ -416,7 +440,6 @@ def generate(firewall):
         firewall['cleanup_commands'] = cleanup_commands(firewall)
 
     render(nftables_conf, 'firewall/nftables.j2', firewall)
-    render(nftables_defines_conf, 'firewall/nftables-defines.j2', firewall)
     return None
 
 def apply_sysfs(firewall):
diff --git a/src/conf_mode/policy-route.py b/src/conf_mode/policy-route.py
index 5de341beb..9fddbd2c6 100755
--- a/src/conf_mode/policy-route.py
+++ b/src/conf_mode/policy-route.py
@@ -25,6 +25,7 @@ from vyos.config import Config
 from vyos.template import render
 from vyos.util import cmd
 from vyos.util import dict_search_args
+from vyos.util import dict_search_recursive
 from vyos.util import run
 from vyos import ConfigError
 from vyos import airbag
@@ -33,6 +34,9 @@ airbag.enable()
 mark_offset = 0x7FFFFFFF
 nftables_conf = '/run/nftables_policy.conf'
 
+ROUTE_PREFIX = 'VYOS_PBR_'
+ROUTE6_PREFIX = 'VYOS_PBR6_'
+
 preserve_chains = [
     'VYOS_PBR_PREROUTING',
     'VYOS_PBR_POSTROUTING',
@@ -46,6 +50,16 @@ valid_groups = [
     'port_group'
 ]
 
+group_set_prefix = {
+    'A_': 'address_group',
+    'A6_': 'ipv6_address_group',
+#    'D_': 'domain_group',
+    'M_': 'mac_group',
+    'N_': 'network_group',
+    'N6_': 'ipv6_network_group',
+    'P_': 'port_group'
+}
+
 def get_policy_interfaces(conf):
     out = {}
     interfaces = conf.get_config_dict(['interfaces'], key_mangling=('-', '_'), get_first_key=True,
@@ -166,37 +180,55 @@ def verify(policy):
 
     return None
 
-def cleanup_rule(table, jump_chain):
-    commands = []
-    results = cmd(f'nft -a list table {table}').split("\n")
-    for line in results:
-        if f'jump {jump_chain}' in line:
-            handle_search = re.search('handle (\d+)', line)
-            if handle_search:
-                commands.append(f'delete rule {table} {chain} handle {handle_search[1]}')
-    return commands
-
 def cleanup_commands(policy):
     commands = []
+    commands_chains = []
+    commands_sets = []
     for table in ['ip mangle', 'ip6 mangle']:
-        json_str = cmd(f'nft -j list table {table}')
+        route_node = 'route' if table == 'ip mangle' else 'route6'
+        chain_prefix = ROUTE_PREFIX if table == 'ip mangle' else ROUTE6_PREFIX
+
+        json_str = cmd(f'nft -t -j list table {table}')
         obj = loads(json_str)
         if 'nftables' not in obj:
             continue
         for item in obj['nftables']:
             if 'chain' in item:
                 chain = item['chain']['name']
-                if not chain.startswith("VYOS_PBR"):
+                if chain in preserve_chains or not chain.startswith("VYOS_PBR"):
                     continue
+
+                if dict_search_args(policy, route_node, chain.replace(chain_prefix, "", 1)) != None:
+                    commands.append(f'flush chain {table} {chain}')
+                else:
+                    commands_chains.append(f'delete chain {table} {chain}')
+
+            if 'rule' in item:
+                rule = item['rule']
+                chain = rule['chain']
+                handle = rule['handle']
+
                 if chain not in preserve_chains:
-                    if table == 'ip mangle' and dict_search_args(policy, 'route', chain.replace("VYOS_PBR_", "", 1)):
-                        commands.append(f'flush chain {table} {chain}')
-                    elif table == 'ip6 mangle' and dict_search_args(policy, 'route6', chain.replace("VYOS_PBR6_", "", 1)):
-                        commands.append(f'flush chain {table} {chain}')
-                    else:
-                        commands += cleanup_rule(table, chain)
-                        commands.append(f'delete chain {table} {chain}')
-    return commands
+                    continue
+
+                target, _ = next(dict_search_recursive(rule['expr'], 'target'))
+
+                if target.startswith(chain_prefix):
+                    if dict_search_args(policy, route_node, target.replace(chain_prefix, "", 1)) == None:
+                        commands.append(f'delete rule {table} {chain} handle {handle}')
+
+            if 'set' in item:
+                set_name = item['set']['name']
+
+                for prefix, group_type in group_set_prefix.items():
+                    if set_name.startswith(prefix):
+                        group_name = set_name.replace(prefix, "", 1)
+                        if dict_search_args(policy, 'firewall_group', group_type, group_name) != None:
+                            commands_sets.append(f'flush set {table} {set_name}')
+                        else:
+                            commands_sets.append(f'delete set {table} {set_name}')
+
+    return commands + commands_chains + commands_sets
 
 def generate(policy):
     if not os.path.exists(nftables_conf):
-- 
cgit v1.2.3


From 7e59b2a3f31edd4793264876d87af725771a222d Mon Sep 17 00:00:00 2001
From: sarthurdev <965089+sarthurdev@users.noreply.github.com>
Date: Tue, 14 Jun 2022 16:19:55 +0200
Subject: firewall: T970: Use set prefix to domain groups

---
 data/templates/firewall/nftables.j2      | 2 +-
 python/vyos/firewall.py                  | 2 +-
 smoketest/scripts/cli/test_firewall.py   | 6 +++---
 src/conf_mode/firewall.py                | 5 +++--
 src/helpers/vyos-domain-group-resolve.py | 2 +-
 5 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/data/templates/firewall/nftables.j2 b/data/templates/firewall/nftables.j2
index ca24b7db2..b91fed615 100644
--- a/data/templates/firewall/nftables.j2
+++ b/data/templates/firewall/nftables.j2
@@ -47,7 +47,7 @@ table ip filter {
 {%     endfor %}
 {%     if group is vyos_defined and group.domain_group is vyos_defined %}
 {%         for name, name_config in group.domain_group.items() %}
-    set {{ name }} {
+    set D_{{ name }} {
         type ipv4_addr
         flags interval
     }
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index f8f913944..7d1278d0e 100644
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -200,7 +200,7 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
                     if group_name[0] == '!':
                         operator = '!='
                         group_name = group_name[1:]
-                    output.append(f'{ip_name} {prefix}addr {operator} @{group_name}')
+                    output.append(f'{ip_name} {prefix}addr {operator} @D_{group_name}')
                 elif 'network_group' in group:
                     group_name = group['network_group']
                     operator = ''
diff --git a/smoketest/scripts/cli/test_firewall.py b/smoketest/scripts/cli/test_firewall.py
index 8b8c27a9f..ce06b9074 100755
--- a/smoketest/scripts/cli/test_firewall.py
+++ b/smoketest/scripts/cli/test_firewall.py
@@ -62,7 +62,7 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
             ['set M_smoketest_mac'],
             ['set N_smoketest_network'],
             ['set P_smoketest_port'],
-            ['set smoketest_domain'],
+            ['set D_smoketest_domain'],
             ['set RECENT_smoketest_4'],
             ['chain NAME_smoketest']
         ]
@@ -116,10 +116,10 @@ class TestFirewall(VyOSUnitTestSHIM.TestCase):
             ['elements = { 53, 123 }'],
             ['ether saddr @M_smoketest_mac', 'return'],
             ['elements = { 00:01:02:03:04:05 }'],
-            ['set smoketest_domain'],
+            ['set D_smoketest_domain'],
             ['elements = { 192.0.2.5, 192.0.2.8,'],
             ['192.0.2.10, 192.0.2.11 }'],
-            ['ip saddr @smoketest_domain', 'return']
+            ['ip saddr @D_smoketest_domain', 'return']
         ]
         self.verify_nftables(nftables_search, 'ip filter')
 
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 78dffe9dd..07eca722f 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -100,6 +100,7 @@ nested_group_types = [
 group_set_prefix = {
     'A_': 'address_group',
     'A6_': 'ipv6_address_group',
+    'D_': 'domain_group',
     'M_': 'mac_group',
     'N_': 'network_group',
     'N6_': 'ipv6_network_group',
@@ -535,8 +536,8 @@ def apply(firewall):
                 # and add elements to nft set
                 ip_dict = get_ips_domains_dict(domains)
                 elements = sum(ip_dict.values(), [])
-                nft_init_set(group)
-                nft_add_set_elements(group, elements)
+                nft_init_set(f'D_{group}')
+                nft_add_set_elements(f'D_{group}', elements)
         else:
             call('systemctl stop vyos-domain-group-resolve.service')
 
diff --git a/src/helpers/vyos-domain-group-resolve.py b/src/helpers/vyos-domain-group-resolve.py
index e8501cfc6..6b677670b 100755
--- a/src/helpers/vyos-domain-group-resolve.py
+++ b/src/helpers/vyos-domain-group-resolve.py
@@ -56,5 +56,5 @@ if __name__ == '__main__':
 
                 # Resolve successful
                 if elements:
-                    nft_update_set_elements(set_name, elements)
+                    nft_update_set_elements(f'D_{set_name}', elements)
         time.sleep(timeout)
-- 
cgit v1.2.3