Merge pull request #3906 from vyos/mergify/bp/circinus/pr-3715

T6313: Add "NAT" to "generate" command for rule resequence (backport #3715)

6 files changed, 89 insertions, 44 deletions

diff --git a/op-mode-definitions/generate_firewall_rule-resequence.xml.in b/op-mode-definitions/generate_firewall_rule-resequence.xml.in
index 66078deb9..ef81579fa 100644
--- a/op-mode-definitions/generate_firewall_rule-resequence.xml.in
+++ b/op-mode-definitions/generate_firewall_rule-resequence.xml.in
@@ -7,34 +7,7 @@
           <help>Firewall</help>
         </properties>
         <children>
-          <node name="rule-resequence">
-            <properties>
-              <help>Resequence the firewall rules</help>
-            </properties>
-            <command>${vyos_op_scripts_dir}/generate_firewall_rule-resequence.py</command>
-            <children>
-              <tagNode name="start">
-                <properties>
-                  <help>Set the first sequence number</help>
-                  <completionHelp>
-                    <list>1-1000</list>
-                  </completionHelp>
-                </properties>
-                <command>${vyos_op_scripts_dir}/generate_firewall_rule-resequence.py --start $5</command>
-                <children>
-                  <tagNode name="step">
-                    <properties>
-                      <help>Step between rules</help>
-                      <completionHelp>
-                        <list>1-1000</list>
-                      </completionHelp>
-                    </properties>
-                    <command>${vyos_op_scripts_dir}/generate_firewall_rule-resequence.py --start $5 --step $7</command>
-                  </tagNode>
-                </children>
-              </tagNode>
-            </children>
-          </node>
+          #include <include/rule-resequence.xml.i>
         </children>
       </node>
     </children>
diff --git a/op-mode-definitions/generate_nat64_rule-resequence.xml.in b/op-mode-definitions/generate_nat64_rule-resequence.xml.in
new file mode 100644
index 000000000..399253b37
--- /dev/null
+++ b/op-mode-definitions/generate_nat64_rule-resequence.xml.in
@@ -0,0 +1,15 @@
+<?xml version="1.0"?>
+<interfaceDefinition>
+  <node name="generate">
+    <children>
+      <node name="nat64">
+        <properties>
+          <help>Network Address Translation (NAT64)</help>
+        </properties>
+        <children>
+          #include <include/rule-resequence.xml.i>
+        </children>
+      </node>
+    </children>
+  </node>
+</interfaceDefinition>
diff --git a/op-mode-definitions/generate_nat66_rule-resequence.xml.in b/op-mode-definitions/generate_nat66_rule-resequence.xml.in
new file mode 100644
index 000000000..d7159cf60
--- /dev/null
+++ b/op-mode-definitions/generate_nat66_rule-resequence.xml.in
@@ -0,0 +1,15 @@
+<?xml version="1.0"?>
+<interfaceDefinition>
+  <node name="generate">
+    <children>
+      <node name="nat66">
+        <properties>
+          <help>Network Prefix Translation (NAT66/NPTv6)</help>
+        </properties>
+        <children>
+          #include <include/rule-resequence.xml.i>
+        </children>
+      </node>
+    </children>
+  </node>
+</interfaceDefinition>
diff --git a/op-mode-definitions/generate_nat_rule-resequence.xml.in b/op-mode-definitions/generate_nat_rule-resequence.xml.in
new file mode 100644
index 000000000..e32a89e08
--- /dev/null
+++ b/op-mode-definitions/generate_nat_rule-resequence.xml.in
@@ -0,0 +1,15 @@
+<?xml version="1.0"?>
+<interfaceDefinition>
+  <node name="generate">
+    <children>
+      <node name="nat">
+        <properties>
+          <help>Network Address Translation (NAT)</help>
+        </properties>
+        <children>
+          #include <include/rule-resequence.xml.i>
+        </children>
+      </node>
+    </children>
+  </node>
+</interfaceDefinition>
diff --git a/op-mode-definitions/include/rule-resequence.xml.i b/op-mode-definitions/include/rule-resequence.xml.i
new file mode 100644
index 000000000..987bf634e
--- /dev/null
+++ b/op-mode-definitions/include/rule-resequence.xml.i
@@ -0,0 +1,30 @@
+<!-- included start from show-nht.xml.i -->
+<node name="rule-resequence">
+  <properties>
+    <help>Resequence rules</help>
+  </properties>
+  <command>${vyos_op_scripts_dir}/generate_service_rule-resequence.py --service $2</command>
+  <children>
+    <tagNode name="start">
+      <properties>
+        <help>Set the first sequence number</help>
+        <completionHelp>
+          <list>1-1000</list>
+        </completionHelp>
+      </properties>
+      <command>${vyos_op_scripts_dir}/generate_service_rule-resequence.py --service $2 --start $5</command>
+      <children>
+        <tagNode name="step">
+          <properties>
+            <help>Step between rules</help>
+            <completionHelp>
+              <list>1-1000</list>
+            </completionHelp>
+          </properties>
+          <command>${vyos_op_scripts_dir}/generate_service_rule-resequence.py --service $2 --start $5 --step $7</command>
+        </tagNode>
+      </children>
+    </tagNode>
+  </children>
+</node>
+<!-- included end -->
diff --git a/src/op_mode/generate_firewall_rule-resequence.py b/src/op_mode/generate_service_rule-resequence.py
index 21441f689..9333d6353 100755
--- a/src/op_mode/generate_firewall_rule-resequence.py
+++ b/src/op_mode/generate_service_rule-resequence.py
@@ -77,7 +77,7 @@ def change_rule_numbers(config_dict, start, step):
             change_rule_numbers(config_dict[key], start, step)
 
 
-def convert_rule_keys_to_int(config_dict):
+def convert_rule_keys_to_int(config_dict, prev_key=None):
     """
     Converts rule keys in the configuration dictionary to integers.
 
@@ -91,11 +91,11 @@ def convert_rule_keys_to_int(config_dict):
         new_dict = {}
         for key, value in config_dict.items():
             # Convert key to integer if possible
-            new_key = int(key) if key.isdigit() else key
+            new_key = int(key) if key.isdigit() and prev_key == 'rule' else key
 
             # Recur for nested dictionaries
             if isinstance(value, dict):
-                new_value = convert_rule_keys_to_int(value)
+                new_value = convert_rule_keys_to_int(value, key)
             else:
                 new_value = value
 
@@ -111,27 +111,24 @@ def convert_rule_keys_to_int(config_dict):
 if __name__ == "__main__":
     # Parse command-line arguments
     parser = argparse.ArgumentParser(description='Convert dictionary to set commands with rule number modifications.')
-    parser.add_argument('--start', type=int, default=100, help='Start rule number')
+    parser.add_argument('--service', type=str, help='Name of service')
+    parser.add_argument('--start', type=int, default=100, help='Start rule number (default: 100)')
     parser.add_argument('--step', type=int, default=10, help='Step for rule numbers (default: 10)')
     args = parser.parse_args()
 
     config = ConfigTreeQuery()
-    if not config.exists('firewall'):
-        print('Firewall is not configured')
+    if not config.exists(args.service):
+        print(f'{args.service} is not configured')
         exit(1)
 
-    config_dict = config.get_config_dict('firewall')
+    config_dict = config.get_config_dict(args.service)
 
-    # Remove global-options, group and flowtable as they don't need sequencing
-    if 'global-options' in config_dict['firewall']:
-        del config_dict['firewall']['global-options']
+    if 'firewall' in config_dict:
+        # Remove global-options, group and flowtable as they don't need sequencing
+        for item in ['global-options', 'group', 'flowtable']:
+            if item in config_dict['firewall']:
+                del config_dict['firewall'][item]
 
-    if 'group' in config_dict['firewall']:
-        del config_dict['firewall']['group']
-
-    if 'flowtable' in config_dict['firewall']:
-        del config_dict['firewall']['flowtable']
-    
     # Convert rule keys to integers, rule "10" -> rule 10
     # This is necessary for sorting the rules
     config_dict = convert_rule_keys_to_int(config_dict)