migration: T6007: convert all migration scripts to load as module

(cherry picked from commit 26740a8d583f64dc0a27b59dd4ae303056972c0b)

1 files changed, 58 insertions, 79 deletions

diff --git a/src/migration-scripts/ipsec/7-to-8 b/src/migration-scripts/ipsec/7-to-8
index 9acc737d5..481f00d29 100755..100644
--- a/src/migration-scripts/ipsec/7-to-8
+++ b/src/migration-scripts/ipsec/7-to-8
@@ -1,18 +1,17 @@
-#!/usr/bin/env python3
+# Copyright 2021-2024 VyOS maintainers and contributors <maintainers@vyos.io>
 #
-# Copyright (C) 2021-2024 VyOS maintainers and contributors
+# 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 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,
+# 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 General Public License for more details.
+# 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 General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# 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
 
@@ -22,29 +21,15 @@ import struct
 
 from cryptography.hazmat.primitives.asymmetric import rsa
 
-from sys import argv
-from sys import exit
-
 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
 
-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()
-
 pki_base = ['pki']
 ipsec_site_base = ['vpn', 'ipsec', 'site-to-site', 'peer']
 rsa_keys_base = ['vpn', 'rsa-keys']
 
-config = ConfigTree(config_file)
-
 LOCAL_KEY_PATHS = ['/config/auth/', '/config/ipsec.d/rsa-keys/']
 
 def migrate_from_vyatta_key(data):
@@ -60,65 +45,59 @@ def wrapped_pem_to_config_value(pem):
 
 local_key_name = 'localhost'
 
-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
+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')
 
-            vyatta_key = config.return_value(rsa_keys_base + ['rsa-key-name', rsa_name, 'rsa-key'])
-            public_key = migrate_from_vyatta_key(vyatta_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
 
-            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}"')
+                vyatta_key = config.return_value(rsa_keys_base + ['rsa-key-name', rsa_name, 'rsa-key'])
+                public_key = migrate_from_vyatta_key(vyatta_key)
 
-    config.delete(rsa_keys_base)
+                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}"')
 
-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'])
+        config.delete(rsa_keys_base)
 
-        if mode != 'rsa':
-            continue
+    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'])
 
-        config.set(ipsec_site_base + [peer, 'authentication', 'rsa', 'local-key'], value=local_key_name)
+            if mode != 'rsa':
+                continue
 
-        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'])
+            config.set(ipsec_site_base + [peer, 'authentication', 'rsa', 'local-key'], value=local_key_name)
 
-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))
-    sys.exit(1)
+            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'])