T6313: Add "NAT" to "generate" command for rule resequence

1 files changed, 145 insertions, 0 deletions

diff --git a/src/op_mode/generate_service_rule-resequence.py b/src/op_mode/generate_service_rule-resequence.py
new file mode 100755
index 000000000..9333d6353
--- /dev/null
+++ b/src/op_mode/generate_service_rule-resequence.py
@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2023 VyOS maintainers and contributors
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 or later as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+import argparse
+from vyos.configquery import ConfigTreeQuery
+
+
+def convert_to_set_commands(config_dict, parent_key=''):
+    """
+    Converts a configuration dictionary into a list of set commands.
+
+    Args:
+        config_dict (dict): The configuration dictionary.
+        parent_key (str): The parent key for nested dictionaries.
+
+    Returns:
+        list: A list of set commands.
+    """
+    commands = []
+    for key, value in config_dict.items():
+        current_key = parent_key + key if parent_key else key
+
+        if isinstance(value, dict):
+            if not value:
+                commands.append(f"set {current_key}")
+            else:
+                commands.extend(
+                    convert_to_set_commands(value, f"{current_key} "))
+
+        elif isinstance(value, list):
+            for item in value:
+                commands.append(f"set {current_key} '{item}'")
+
+        elif isinstance(value, str):
+            commands.append(f"set {current_key} '{value}'")
+
+    return commands
+
+
+def change_rule_numbers(config_dict, start, step):
+    """
+    Changes rule numbers in the configuration dictionary.
+
+    Args:
+        config_dict (dict): The configuration dictionary.
+        start (int): The starting rule number.
+        step (int): The step to increment the rule numbers.
+
+    Returns:
+        None
+    """
+    if 'rule' in config_dict:
+        rule_dict = config_dict['rule']
+        updated_rule_dict = {}
+        rule_num = start
+        for rule_key in sorted(rule_dict.keys()):
+            updated_rule_dict[str(rule_num)] = rule_dict[rule_key]
+            rule_num += step
+        config_dict['rule'] = updated_rule_dict
+
+    for key in config_dict:
+        if isinstance(config_dict[key], dict):
+            change_rule_numbers(config_dict[key], start, step)
+
+
+def convert_rule_keys_to_int(config_dict, prev_key=None):
+    """
+    Converts rule keys in the configuration dictionary to integers.
+
+    Args:
+        config_dict (dict or list): The configuration dictionary or list.
+
+    Returns:
+        dict or list: The modified dictionary or list.
+    """
+    if isinstance(config_dict, dict):
+        new_dict = {}
+        for key, value in config_dict.items():
+            # Convert key to integer if possible
+            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, key)
+            else:
+                new_value = value
+
+            new_dict[new_key] = new_value
+
+        return new_dict
+    elif isinstance(config_dict, list):
+        return [convert_rule_keys_to_int(item) for item in config_dict]
+    else:
+        return 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('--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(args.service):
+        print(f'{args.service} is not configured')
+        exit(1)
+
+    config_dict = config.get_config_dict(args.service)
+
+    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]
+
+    # 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)
+
+    # Apply rule number modifications
+    change_rule_numbers(config_dict, start=args.start, step=args.step)
+
+    # Convert to 'set' commands
+    set_commands = convert_to_set_commands(config_dict)
+
+    print()
+    for command in set_commands:
+        print(command)
+    print()