#!/usr/bin/env python3
#
# Copyright (C) 2024 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/>.

# Delete 'protocols bgp address-family ipv6-unicast route-target vpn
# import/export', if 'protocols bgp address-family ipv6-unicast
# route-target vpn both' exists

from sys import argv
from sys import exit

from vyos.configtree import ConfigTree

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

file_name = argv[1]

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

config = ConfigTree(config_file)

bgp_base = ['protocols', 'bgp']
# Delete 'import/export' in default vrf if 'both' exists
if config.exists(bgp_base):
    for address_family in ['ipv4-unicast', 'ipv6-unicast']:
        rt_path = bgp_base + ['address-family', address_family, 'route-target',
                              'vpn']
        if config.exists(rt_path + ['both']):
            if config.exists(rt_path + ['import']):
                config.delete(rt_path + ['import'])
            if config.exists(rt_path + ['export']):
                config.delete(rt_path + ['export'])

# Delete import/export in vrfs if both exists
if config.exists(['vrf', 'name']):
    for vrf in config.list_nodes(['vrf', 'name']):
        vrf_base = ['vrf', 'name', vrf]
        for address_family in ['ipv4-unicast', 'ipv6-unicast']:
            rt_path = vrf_base + bgp_base + ['address-family', address_family,
                                             'route-target', 'vpn']
            if config.exists(rt_path + ['both']):
                if config.exists(rt_path + ['import']):
                    config.delete(rt_path + ['import'])
                if config.exists(rt_path + ['export']):
                    config.delete(rt_path + ['export'])

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