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

from sys import argv
from sys import exit

from vyos.configtree import ConfigTree
from vyos.template import is_ipv4
from vyos.template import is_ipv6


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 = ['vpn', 'ipsec']
config = ConfigTree(config_file)

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

# IKE changes, T4118:
if config.exists(base + ['ike-group']):
    for ike_group in config.list_nodes(base + ['ike-group']):
        # replace 'ipsec ike-group <tag> mobike disable'
        #      => 'ipsec ike-group <tag> disable-mobike'
        mobike = base + ['ike-group', ike_group, 'mobike']
        if config.exists(mobike):
            if config.return_value(mobike) == 'disable':
                config.set(base + ['ike-group', ike_group, 'disable-mobike'])
            config.delete(mobike)

        # replace 'ipsec ike-group <tag> ikev2-reauth yes'
        #      => 'ipsec ike-group <tag> ikev2-reauth'
        reauth = base + ['ike-group', ike_group, 'ikev2-reauth']
        if config.exists(reauth):
            if config.return_value(reauth) == 'yes':
                config.delete(reauth)
                config.set(reauth)
            else:
                config.delete(reauth)

# ESP changes
# replace 'ipsec esp-group <tag> compression enable'
#      => 'ipsec esp-group <tag> compression'
if config.exists(base + ['esp-group']):
    for esp_group in config.list_nodes(base + ['esp-group']):
        compression = base + ['esp-group', esp_group, 'compression']
        if config.exists(compression):
            if config.return_value(compression) == 'enable':
                config.delete(compression)
                config.set(compression)
            else:
                config.delete(compression)

# PEER changes
if config.exists(base + ['site-to-site', 'peer']):
    for peer in config.list_nodes(base + ['site-to-site', 'peer']):
        # replace: 'peer <tag> id x'
        #       => 'peer <tag> local-id x'
        if config.exists(base + ['site-to-site', 'peer', peer, 'authentication', 'id']):
            config.rename(base + ['site-to-site', 'peer', peer, 'authentication', 'id'], 'local-id')

        # For the peer '@foo' set remote-id 'foo' if remote-id is not defined
        if peer.startswith('@'):
            if not config.exists(base + ['site-to-site', 'peer', peer, 'authentication', 'remote-id']):
                tmp = peer.replace('@', '')
                config.set(base + ['site-to-site', 'peer', peer, 'authentication', 'remote-id'], value=tmp)

        # replace: 'peer <tag> force-encapsulation enable'
        #       => 'peer <tag> force-udp-encapsulation'
        force_enc = base + ['site-to-site', 'peer', peer, 'force-encapsulation']
        if config.exists(force_enc):
            if config.return_value(force_enc) == 'enable':
                config.delete(force_enc)
                config.set(base + ['site-to-site', 'peer', peer, 'force-udp-encapsulation'])
            else:
                config.delete(force_enc)

        # add option: 'peer <tag> remote-address x.x.x.x'
        remote_address = peer
        if peer.startswith('@'):
            remote_address = 'any'
        config.set(base + ['site-to-site', 'peer', peer, 'remote-address'], value=remote_address)
        # Peer name it is swanctl connection name and shouldn't contain dots or colons
        # rename peer:
        #   peer 192.0.2.1   => peer peer_192-0-2-1
        #   peer 2001:db8::2 => peer peer_2001-db8--2
        #   peer @foo        => peer peer_foo
        re_peer_name = re.sub(':|\.', '-', peer)
        if re_peer_name.startswith('@'):
            re_peer_name = re.sub('@', '', re_peer_name)
        new_peer_name = f'peer_{re_peer_name}'

        config.rename(base + ['site-to-site', 'peer', peer], new_peer_name)

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)