# Copyright 2021-2024 VyOS maintainers and contributors <maintainers@vyos.io>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this library.  If not, see <http://www.gnu.org/licenses/>.

# Migrate rsa keys into PKI configuration

import base64
import os
import struct

from cryptography.hazmat.primitives.asymmetric import rsa

from vyos.configtree import ConfigTree
from vyos.pki import load_private_key
from vyos.pki import encode_public_key
from vyos.pki import encode_private_key

pki_base = ['pki']
ipsec_site_base = ['vpn', 'ipsec', 'site-to-site', 'peer']
rsa_keys_base = ['vpn', 'rsa-keys']

LOCAL_KEY_PATHS = ['/config/auth/', '/config/ipsec.d/rsa-keys/']

def migrate_from_vyatta_key(data):
    data = base64.b64decode(data[2:])
    length = struct.unpack('B', data[:1])[0]
    e = int.from_bytes(data[1:1+length], 'big')
    n = int.from_bytes(data[1+length:], 'big')
    public_numbers = rsa.RSAPublicNumbers(e, n)
    return public_numbers.public_key()

def wrapped_pem_to_config_value(pem):
    return "".join(pem.strip().split("\n")[1:-1])

local_key_name = 'localhost'

def migrate(config: ConfigTree) -> None:
    if config.exists(rsa_keys_base):
        if not config.exists(pki_base + ['key-pair']):
            config.set(pki_base + ['key-pair'])
            config.set_tag(pki_base + ['key-pair'])

        if config.exists(rsa_keys_base + ['local-key', 'file']):
            local_file = config.return_value(rsa_keys_base + ['local-key', 'file'])
            local_path = None
            local_key = None

            for path in LOCAL_KEY_PATHS:
                full_path = os.path.join(path, local_file)
                if os.path.exists(full_path):
                    local_path = full_path
                    break

            if local_path:
                with open(local_path, 'r') as f:
                    local_key_data = f.read()
                    local_key = load_private_key(local_key_data, wrap_tags=False)

            if local_key:
                local_key_pem = encode_private_key(local_key)
                config.set(pki_base + ['key-pair', local_key_name, 'private', 'key'], value=wrapped_pem_to_config_value(local_key_pem))
            else:
                print('Failed to migrate local RSA key')

        if config.exists(rsa_keys_base + ['rsa-key-name']):
            for rsa_name in config.list_nodes(rsa_keys_base + ['rsa-key-name']):
                if not config.exists(rsa_keys_base + ['rsa-key-name', rsa_name, 'rsa-key']):
                    continue

                vyatta_key = config.return_value(rsa_keys_base + ['rsa-key-name', rsa_name, 'rsa-key'])
                public_key = migrate_from_vyatta_key(vyatta_key)

                if public_key:
                    public_key_pem = encode_public_key(public_key)
                    config.set(pki_base + ['key-pair', rsa_name, 'public', 'key'], value=wrapped_pem_to_config_value(public_key_pem))
                else:
                    print(f'Failed to migrate rsa-key "{rsa_name}"')

        config.delete(rsa_keys_base)

    if config.exists(ipsec_site_base):
        for peer in config.list_nodes(ipsec_site_base):
            mode = config.return_value(ipsec_site_base + [peer, 'authentication', 'mode'])

            if mode != 'rsa':
                continue

            config.set(ipsec_site_base + [peer, 'authentication', 'rsa', 'local-key'], value=local_key_name)

            remote_key_name = config.return_value(ipsec_site_base + [peer, 'authentication', 'rsa-key-name'])
            config.set(ipsec_site_base + [peer, 'authentication', 'rsa', 'remote-key'], value=remote_key_name)
            config.delete(ipsec_site_base + [peer, 'authentication', 'rsa-key-name'])