#!/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/>.

# T5160: Firewall re-writing

#  cli changes from:       
#  set firewall name <name> ...
#  set firewall ipv6-name <name> ...
#  To
#  set firewall ip name <name> 
#  set firewall ipv6 ipv6-name <name> 

## Also from 'firewall interface' removed.
## in and out:
    # set firewall interface <iface> [in|out] [name | ipv6-name] <name>
    # To
    # set firewall [ip | ipv6] forward filter rule <5,10,15,...> [inbound-interface | outboubd-interface] interface-name <iface>
    # set firewall [ip | ipv6] forward filter rule <5,10,15,...> action jump
    # set firewall [ip | ipv6] forward filter rule <5,10,15,...> jump-target <name>
## local:
    # set firewall interface <iface> local [name | ipv6-name] <name>
    # To
    # set firewall [ip | ipv6] input filter rule <5,10,15,...> inbound-interface interface-name <iface>
    # set firewall [ip | ipv6] input filter rule <5,10,15,...> action jump
    # set firewall [ip | ipv6] input filter rule <5,10,15,...> jump-target <name>

import re

from sys import argv
from sys import exit

from vyos.configtree import ConfigTree
from vyos.ifconfig import Section

if (len(argv) < 1):
    print("Must specify file name!")
    exit(1)

file_name = argv[1]

with open(file_name, 'r') as f:
    config_file = f.read()

base = ['firewall']
config = ConfigTree(config_file)

if not config.exists(base):
    # Nothing to do
    exit(0)

### Migration of state policies
if config.exists(base + ['state-policy']):
    for family in ['ip', 'ipv6']:
        for hook in ['forward', 'input', 'output']:
            for priority in ['filter']:
                # Add default-action== accept for compatibility reasons:
                config.set(base + [family, hook, priority, 'default-action'], value='accept')
                position = 1
                for state in config.list_nodes(base + ['state-policy']):
                    action = config.return_value(base + ['state-policy', state, 'action'])
                    config.set(base + [family, hook, priority, 'rule'])
                    config.set_tag(base + [family, hook, priority, 'rule'])
                    config.set(base + [family, hook, priority, 'rule', position, 'state', state], value='enable')
                    config.set(base + [family, hook, priority, 'rule', position, 'action'], value=action)
                    position = position + 1
    config.delete(base + ['state-policy'])
############

## migration of global options:
for option in ['all-ping', 'broadcast-ping', 'config-trap', 'ip-src-route', 'ipv6-receive-redirects', 'ipv6-src-route', 'log-martians',
                'receive-redirects', 'resolver-cache', 'resolver-internal', 'send-redirects', 'source-validation', 'syn-cookies', 'twa-hazards-protection']:
    if config.exists(base + [option]):
        val = config.return_value(base + [option])
        config.set(base + ['global-options', option], value=val)
        config.delete(base + [option])

### Migration of firewall name and ipv6-name
if config.exists(base + ['name']):
    config.set(['firewall', 'ip', 'name'])
    config.set_tag(['firewall', 'ip', 'name'])

    for ipv4name in config.list_nodes(base + ['name']):
        config.copy(base + ['name', ipv4name], base + ['ip', 'name', ipv4name])
    config.delete(base + ['name'])

if config.exists(base + ['ipv6-name']):
    config.set(['firewall', 'ipv6', 'ipv6-name'])
    config.set_tag(['firewall', 'ipv6', 'ipv6-name'])

    for ipv6name in config.list_nodes(base + ['ipv6-name']):
        config.copy(base + ['ipv6-name', ipv6name], base + ['ipv6', 'ipv6-name', ipv6name])
    config.delete(base + ['ipv6-name'])

### Migration of firewall interface
if config.exists(base + ['interface']):
    fwd_ipv4_rule = 5
    inp_ipv4_rule = 5
    fwd_ipv6_rule = 5
    inp_ipv6_rule = 5
    for iface in config.list_nodes(base + ['interface']):
        for direction in ['in', 'out', 'local']:
            if config.exists(base + ['interface', iface, direction]):
                if config.exists(base + ['interface', iface, direction, 'name']):
                    target = config.return_value(base + ['interface', iface, direction, 'name'])
                    if direction == 'in':
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ip', 'forward', 'filter', 'default-action'], value='accept')
                        new_base = base + ['ip', 'forward', 'filter', 'rule']
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [fwd_ipv4_rule, 'inbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [fwd_ipv4_rule, 'action'], value='jump')
                        config.set(new_base + [fwd_ipv4_rule, 'jump-target'], value=target)
                        fwd_ipv4_rule = fwd_ipv4_rule + 5
                    elif direction == 'out':
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ip', 'forward', 'filter', 'default-action'], value='accept')
                        new_base = base + ['ip', 'forward', 'filter', 'rule']
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [fwd_ipv4_rule, 'outbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [fwd_ipv4_rule, 'action'], value='jump')
                        config.set(new_base + [fwd_ipv4_rule, 'jump-target'], value=target)
                        fwd_ipv4_rule = fwd_ipv4_rule + 5
                    else:
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ip', 'input', 'filter', 'default-action'], value='accept')
                        new_base = base + ['ip', 'input', 'filter', 'rule']
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [inp_ipv4_rule, 'inbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [inp_ipv4_rule, 'action'], value='jump')
                        config.set(new_base + [inp_ipv4_rule, 'jump-target'], value=target)
                        inp_ipv4_rule = inp_ipv4_rule + 5

                if config.exists(base + ['interface', iface, direction, 'ipv6-name']):
                    target = config.return_value(base + ['interface', iface, direction, 'ipv6-name'])
                    if direction == 'in':
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ipv6', 'forward', 'filter', 'default-action'], value='accept')
                        new_base = base + ['ipv6', 'forward', 'filter', 'rule']
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [fwd_ipv6_rule, 'inbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [fwd_ipv6_rule, 'action'], value='jump')
                        config.set(new_base + [fwd_ipv6_rule, 'jump-target'], value=target)
                        fwd_ipv6_rule = fwd_ipv6_rule + 5
                    elif direction == 'out':
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ipv6', 'forward', 'filter', 'default-action'], value='accept')
                        new_base = base + ['ipv6', 'forward', 'filter', 'rule']
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [fwd_ipv6_rule, 'outbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [fwd_ipv6_rule, 'action'], value='jump')
                        config.set(new_base + [fwd_ipv6_rule, 'jump-target'], value=target)
                        fwd_ipv6_rule = fwd_ipv6_rule + 5
                    else:
                        new_base = base + ['ipv6', 'input', 'filter', 'rule']
                        # Add default-action== accept for compatibility reasons:
                        config.set(base + ['ipv6', 'input', 'filter', 'default-action'], value='accept')
                        config.set(new_base)
                        config.set_tag(new_base)
                        config.set(new_base + [inp_ipv6_rule, 'inbound-interface', 'interface-name'], value=iface)
                        config.set(new_base + [inp_ipv6_rule, 'action'], value='jump')
                        config.set(new_base + [inp_ipv6_rule, 'jump-target'], value=target)
                        inp_ipv6_rule = inp_ipv6_rule + 5

    config.delete(base + ['interface'])


### Migration of zones config v2:
### User interface groups 
if config.exists(base + ['zone']):
    inp_ipv4_rule = 101
    inp_ipv6_rule = 101
    fwd_ipv4_rule = 101
    fwd_ipv6_rule = 101
    out_ipv4_rule = 101
    out_ipv6_rule = 101
    local_zone = 'False'

    for zone in config.list_nodes(base + ['zone']):
        if config.exists(base + ['zone', zone, 'local-zone']):
            local_zone = 'True'
            # Add default-action== accept for compatibility reasons:
            config.set(base + ['ip', 'input', 'filter', 'default-action'], value='accept')
            config.set(base + ['ipv6', 'input', 'filter', 'default-action'], value='accept')
            config.set(base + ['ip', 'output', 'filter', 'default-action'], value='accept')
            config.set(base + ['ipv6', 'output', 'filter', 'default-action'], value='accept')
            for from_zone in config.list_nodes(base + ['zone', zone, 'from']):
                group_name = 'IG_' + from_zone
                if config.exists(base + ['zone', zone, 'from', from_zone, 'firewall', 'name']):
                    # ipv4 input ruleset
                    target_ipv4_chain = config.return_value(base + ['zone', zone, 'from', from_zone, 'firewall', 'name'])
                    config.set(base + ['ip', 'input', 'filter', 'rule'])
                    config.set_tag(base + ['ip', 'input', 'filter', 'rule'])
                    config.set(base + ['ip', 'input', 'filter', 'rule', inp_ipv4_rule, 'inbound-interface', 'interface-group'], value=group_name)
                    config.set(base + ['ip', 'input', 'filter', 'rule', inp_ipv4_rule, 'action'], value='jump')
                    config.set(base + ['ip', 'input', 'filter', 'rule', inp_ipv4_rule, 'jump-target'], value=target_ipv4_chain)
                    inp_ipv4_rule = inp_ipv4_rule + 5
                if config.exists(base + ['zone', zone, 'from', from_zone, 'firewall', 'ipv6-name']):
                    # ipv6 input ruleset
                    target_ipv6_chain = config.return_value(base + ['zone', zone, 'from', from_zone, 'firewall', 'ipv6-name'])
                    config.set(base + ['ipv6', 'input', 'filter', 'rule'])
                    config.set_tag(base + ['ipv6', 'input', 'filter', 'rule'])
                    config.set(base + ['ipv6', 'input', 'filter', 'rule', inp_ipv6_rule, 'inbound-interface', 'interface-group'], value=group_name)
                    config.set(base + ['ipv6', 'input', 'filter', 'rule', inp_ipv6_rule, 'action'], value='jump')
                    config.set(base + ['ipv6', 'input', 'filter', 'rule', inp_ipv6_rule, 'jump-target'], value=target_ipv6_chain)
                    inp_ipv6_rule = inp_ipv6_rule + 5

            # Migrate: set firewall zone <zone> default-action <action>
            # Options: drop or reject. If not specified, is drop
            if config.exists(base + ['zone', zone, 'default-action']):
                local_def_action = config.return_value(base + ['zone', zone, 'default-action'])
            else:
                local_def_action = 'drop'
            config.set(base + ['ip', 'input', 'filter', 'rule'])
            config.set_tag(base + ['ip', 'input', 'filter', 'rule'])
            config.set(base + ['ip', 'input', 'filter', 'rule', inp_ipv4_rule, 'action'], value=local_def_action)
            config.set(base + ['ipv6', 'input', 'filter', 'rule'])
            config.set_tag(base + ['ipv6', 'input', 'filter', 'rule'])
            config.set(base + ['ipv6', 'input', 'filter', 'rule', inp_ipv6_rule, 'action'], value=local_def_action)
            if config.exists(base + ['zone', zone, 'enable-default-log']):
                config.set(base + ['ip', 'input', 'filter', 'rule', inp_ipv4_rule, 'log'], value='enable')
                config.set(base + ['ipv6', 'input', 'filter', 'rule', inp_ipv6_rule, 'log'], value='enable')

        else:
        # It's not a local zone
            group_name = 'IG_' + zone
            # Add default-action== accept for compatibility reasons:
            config.set(base + ['ip', 'forward', 'filter', 'default-action'], value='accept')
            config.set(base + ['ipv6', 'forward', 'filter', 'default-action'], value='accept')
            # intra-filtering migration. By default accept
            intra_zone_ipv4_action = 'accept'
            intra_zone_ipv6_action = 'accept'
            
            if config.exists(base + ['zone', zone, 'intra-zone-filtering', 'action']):
                intra_zone_ipv4_action = config.return_value(base + ['zone', zone, 'intra-zone-filtering', 'action'])
                intra_zone_ipv6_action = intra_zone_ipv4_action
            else:
                if config.exists(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'name']):
                    intra_zone_ipv4_target = config.return_value(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'name'])
                    intra_zone_ipv4_action = 'jump'
                if config.exists(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'ipv6-name']):
                    intra_zone_ipv6_target = config.return_value(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'ipv6-name'])
                    intra_zone_ipv6_action = 'jump'
            config.set(base + ['ip', 'forward', 'filter', 'rule'])
            config.set_tag(base + ['ip', 'forward', 'filter', 'rule'])
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'outbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'inbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'action'], value=intra_zone_ipv4_action)
            config.set_tag(base + ['ipv6', 'forward', 'filter', 'rule'])
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'outbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'inbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'action'], value=intra_zone_ipv6_action)
            if intra_zone_ipv4_action == 'jump':
                if config.exists(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'name']):
                    intra_zone_ipv4_target = config.return_value(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'name'])
                    config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'jump-target'], value=intra_zone_ipv4_target)
            if intra_zone_ipv6_action == 'jump':
                if config.exists(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'ipv6-name']):
                    intra_zone_ipv6_target = config.return_value(base + ['zone', zone, 'intra-zone-filtering', 'firewall', 'ipv6-name'])
                    config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'jump-target'], value=intra_zone_ipv6_target)
            fwd_ipv4_rule = fwd_ipv4_rule + 5
            fwd_ipv6_rule = fwd_ipv6_rule + 5

            if config.exists(base + ['zone', zone, 'interface']):
                # Create interface group IG_<zone>
                group_name = 'IG_' + zone
                config.set(base + ['group', 'interface-group'], value=group_name)
                config.set_tag(base + ['group', 'interface-group'])
                for iface in config.return_values(base + ['zone', zone, 'interface']):
                    config.set(base + ['group', 'interface-group', group_name, 'interface'], value=iface, replace=False)

            if config.exists(base + ['zone', zone, 'from']):
                for from_zone in config.list_nodes(base + ['zone', zone, 'from']):
                    from_group = 'IG_' + from_zone
                    if config.exists(base + ['zone', zone, 'from', from_zone, 'firewall', 'name']):
                        target_ipv4_chain = config.return_value(base + ['zone', zone, 'from', from_zone, 'firewall', 'name'])
                        if config.exists(base + ['zone', from_zone, 'local-zone']):
                            # It's from LOCAL zone -> Output filtering 
                            config.set(base + ['ip', 'output', 'filter', 'rule'])
                            config.set_tag(base + ['ip', 'output', 'filter', 'rule'])
                            config.set(base + ['ip', 'output', 'filter', 'rule', out_ipv4_rule, 'outbound-interface', 'interface-group'], value=group_name)
                            config.set(base + ['ip', 'output', 'filter', 'rule', out_ipv4_rule, 'action'], value='jump')
                            config.set(base + ['ip', 'output', 'filter', 'rule', out_ipv4_rule, 'jump-target'], value=target_ipv4_chain)
                            out_ipv4_rule = out_ipv4_rule + 5
                        else:
                            # It's not LOCAL zone -> forward filtering
                            config.set(base + ['ip', 'forward', 'filter', 'rule'])
                            config.set_tag(base + ['ip', 'forward', 'filter', 'rule'])
                            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'outbound-interface', 'interface-group'], value=group_name)
                            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'inbound-interface', 'interface-group'], value=from_group)
                            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'action'], value='jump')
                            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'jump-target'], value=target_ipv4_chain)
                            fwd_ipv4_rule = fwd_ipv4_rule + 5
                    if config.exists(base + ['zone', zone, 'from', from_zone, 'firewall', 'ipv6-name']):
                        target_ipv6_chain = config.return_value(base + ['zone', zone, 'from', from_zone, 'firewall', 'ipv6-name'])
                        if config.exists(base + ['zone', from_zone, 'local-zone']):
                            # It's from LOCAL zone -> Output filtering
                            config.set(base + ['ipv6', 'output', 'filter', 'rule'])
                            config.set_tag(base + ['ipv6', 'output', 'filter', 'rule'])
                            config.set(base + ['ipv6', 'output', 'filter', 'rule', out_ipv6_rule, 'outbound-interface', 'interface-group'], value=group_name)
                            config.set(base + ['ipv6', 'output', 'filter', 'rule', out_ipv6_rule, 'action'], value='jump')
                            config.set(base + ['ipv6', 'output', 'filter', 'rule', out_ipv6_rule, 'jump-target'], value=target_ipv6_chain)
                            out_ipv6_rule = out_ipv6_rule + 5
                        else:
                            # It's not LOCAL zone -> forward filtering
                            config.set(base + ['ipv6', 'forward', 'filter', 'rule'])
                            config.set_tag(base + ['ipv6', 'forward', 'filter', 'rule'])
                            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'outbound-interface', 'interface-group'], value=group_name)
                            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'inbound-interface', 'interface-group'], value=from_group)
                            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'action'], value='jump')
                            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'jump-target'], value=target_ipv6_chain)
                            fwd_ipv6_rule = fwd_ipv6_rule + 5

            ## Now need to migrate: set firewall zone <zone> default-action <action>    # action=drop if not specified.
            if config.exists(base + ['zone', zone, 'default-action']):
                def_action = config.return_value(base + ['zone', zone, 'default-action'])
            else:
                def_action = 'drop'
            config.set(base + ['ip', 'forward', 'filter', 'rule'])
            config.set_tag(base + ['ip', 'forward', 'filter', 'rule'])
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'outbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'action'], value=def_action)
            description = 'zone_' + zone + ' default-action'
            config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'description'], value=description)
            config.set(base + ['ipv6', 'forward', 'filter', 'rule'])
            config.set_tag(base + ['ipv6', 'forward', 'filter', 'rule'])
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'outbound-interface', 'interface-group'], value=group_name)
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'action'], value=def_action)
            config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'description'], value=description)

            if config.exists(base + ['zone', zone, 'enable-default-log']):
                config.set(base + ['ip', 'forward', 'filter', 'rule', fwd_ipv4_rule, 'log'], value='enable')
                config.set(base + ['ipv6', 'forward', 'filter', 'rule', fwd_ipv6_rule, 'log'], value='enable')
            fwd_ipv4_rule = fwd_ipv4_rule + 5
            fwd_ipv6_rule = fwd_ipv6_rule + 5

    # Migrate default-action (force to be drop in output chain) if local zone is defined
    if local_zone == 'True':
        # General drop in output change if needed
        config.set(base + ['ip', 'output', 'filter', 'rule'])
        config.set_tag(base + ['ip', 'output', 'filter', 'rule'])
        config.set(base + ['ip', 'output', 'filter', 'rule', out_ipv4_rule, 'action'], value=local_def_action)
        config.set(base + ['ipv6', 'output', 'filter', 'rule'])
        config.set_tag(base + ['ipv6', 'output', 'filter', 'rule'])
        config.set(base + ['ipv6', 'output', 'filter', 'rule', out_ipv6_rule, 'action'], value=local_def_action)

    config.delete(base + ['zone'])

###### END migration zones v2

try:
    with open(file_name, 'w') as f:
        f.write(config.to_string())
except OSError as e:
    print("Failed to save the modified config: {}".format(e))
    exit(1)