summaryrefslogtreecommitdiff
path: root/src/op_mode/generate_firewall_rule-resequence.py
blob: 4362b484a89fd8648ff33f88eda79cbde5db364d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#!/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, 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):
    """
    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() else key

            # Recur for nested dictionaries
            if isinstance(value, dict):
                new_value = convert_rule_keys_to_int(value)
            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('--start', type=int, default=100, help='Start rule number')
    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')
        exit(1)

    config_dict = config.get_config_dict('firewall')

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

    # 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()