image-tools: T5758: restore saving previous data on install

Restore scanning previous installations for config data and ssh host
keys on install.

2 files changed, 96 insertions, 2 deletions

diff --git a/python/vyos/system/disk.py b/python/vyos/system/disk.py
index 49e6b5c5e..f8e0fd1bf 100644
--- a/python/vyos/system/disk.py
+++ b/python/vyos/system/disk.py
@@ -150,7 +150,7 @@ def filesystem_create(partition: str, fstype: str) -> None:
 def partition_mount(partition: str,
                     path: str,
                     fsype: str = '',
-                    overlay_params: dict[str, str] = {}) -> None:
+                    overlay_params: dict[str, str] = {}) -> bool:
     """Mount a partition into a path
 
     Args:
@@ -159,6 +159,9 @@ def partition_mount(partition: str,
         fsype (str): optionally, set fstype ('squashfs', 'overlay', 'iso9660')
         overlay_params (dict): optionally, set overlay parameters.
         Defaults to None.
+
+    Returns:
+        bool: True on success
     """
     if fsype in ['squashfs', 'iso9660']:
         command: str = f'mount -o loop,ro -t {fsype} {partition} {path}'
@@ -171,7 +174,11 @@ def partition_mount(partition: str,
     else:
         command = f'mount {partition} {path}'
 
-    run(command)
+    rc = run(command)
+    if rc == 0:
+        return True
+
+    return False
 
 
 def partition_umount(partition: str = '', path: str = '') -> None:
diff --git a/src/op_mode/image_installer.py b/src/op_mode/image_installer.py
index cdb84a152..b3e6e518c 100755
--- a/src/op_mode/image_installer.py
+++ b/src/op_mode/image_installer.py
@@ -60,6 +60,8 @@ MSG_INPUT_PASSWORD: str = 'Please enter a password for the "vyos" user'
 MSG_INPUT_ROOT_SIZE_ALL: str = 'Would you like to use all the free space on the drive?'
 MSG_INPUT_ROOT_SIZE_SET: str = 'Please specify the size (in GB) of the root partition (min is 1.5 GB)?'
 MSG_INPUT_CONSOLE_TYPE: str = 'What console should be used by default? (K: KVM, S: Serial, U: USB-Serial)?'
+MSG_INPUT_COPY_DATA: str = 'Would you like to copy data to the new image?'
+MSG_INPUT_CHOOSE_COPY_DATA: str = 'From which image would you like to save config information?'
 MSG_WARN_ISO_SIGN_INVALID: str = 'Signature is not valid. Do you want to continue with installation?'
 MSG_WARN_ISO_SIGN_UNAVAL: str = 'Signature is not available. Do you want to continue with installation?'
 MSG_WARN_ROOT_SIZE_TOOBIG: str = 'The size is too big. Try again.'
@@ -184,6 +186,83 @@ def create_partitions(target_disk: str, target_size: int,
     return disk_details
 
 
+def search_format_selection(image: tuple[str, str]) -> str:
+    """Format a string for selection of image
+
+    Args:
+        image (tuple[str, str]): a tuple of image name and drive
+
+    Returns:
+        str: formatted string
+    """
+    return f'{image[0]} on {image[1]}'
+
+
+def search_previous_installation(disks: list[str]) -> None:
+    """Search disks for previous installation config and SSH keys
+
+    Args:
+        disks (list[str]): a list of available disks
+    """
+    mnt_config = '/mnt/config'
+    mnt_ssh = '/mnt/ssh'
+    mnt_tmp = '/mnt/tmp'
+    rmtree(Path(mnt_config), ignore_errors=True)
+    rmtree(Path(mnt_ssh), ignore_errors=True)
+    Path(mnt_tmp).mkdir(exist_ok=True)
+
+    print('Searching for data from previous installations')
+    image_data = []
+    for disk_name in disks:
+        for partition in disk.partition_list(disk_name):
+            if disk.partition_mount(partition, mnt_tmp):
+                if Path(mnt_tmp + '/boot').exists():
+                    for path in Path(mnt_tmp + '/boot').iterdir():
+                        if path.joinpath('rw/config/.vyatta_config').exists():
+                            image_data.append((path.name, partition))
+
+                disk.partition_umount(partition)
+
+    if len(image_data) == 1:
+        image_name, image_drive = image_data[0]
+        print('Found data from previous installation:')
+        print(f'\t{image_name} on {image_drive}')
+        if not ask_yes_no(MSG_INPUT_COPY_DATA, default=True):
+            return
+
+    elif len(image_data) > 1:
+        print('Found data from previous installations')
+        if not ask_yes_no(MSG_INPUT_COPY_DATA, default=True):
+            return
+
+        image_name, image_drive = select_entry(image_data,
+                                               'Available versions:',
+                                               MSG_INPUT_CHOOSE_COPY_DATA,
+                                               search_format_selection)
+    else:
+        print('No previous installation found')
+        return
+
+    disk.partition_mount(image_drive, mnt_tmp)
+
+    copytree(f'{mnt_tmp}/boot/{image_name}/rw/config', mnt_config)
+    Path(mnt_ssh).mkdir()
+    host_keys: list[str] = glob(f'{mnt_tmp}/boot/{image_name}/rw/etc/ssh/ssh_host*')
+    for host_key in host_keys:
+        copy(host_key, mnt_ssh)
+
+    disk.partition_umount(image_drive)
+
+
+def copy_previous_installation_data(target_dir: str) -> None:
+    if Path('/mnt/config').exists():
+        copytree('/mnt/config', f'{target_dir}/opt/vyatta/etc/config',
+                 dirs_exist_ok=True)
+    if Path('/mnt/ssh').exists():
+        copytree('/mnt/ssh', f'{target_dir}/etc/ssh',
+                 dirs_exist_ok=True)
+
+
 def ask_single_disk(disks_available: dict[str, int]) -> str:
     """Ask user to select a disk for installation
 
@@ -204,6 +283,8 @@ def ask_single_disk(disks_available: dict[str, int]) -> str:
         print(MSG_INFO_INSTALL_EXIT)
         exit()
 
+    search_previous_installation(list(disks_available))
+
     disk_details: disk.DiskDetails = create_partitions(disk_selected,
                                                        disks_available[disk_selected])
 
@@ -260,6 +341,8 @@ def check_raid_install(disks_available: dict[str, int]) -> Union[str, None]:
         print(MSG_INFO_INSTALL_EXIT)
         exit()
 
+    search_previous_installation(list(disks_available))
+
     disks: list[disk.DiskDetails] = []
     for disk_selected in list(disks_selected):
         print(f'Creating partitions on {disk_selected}')
@@ -581,6 +664,10 @@ def install_image() -> None:
         copy(FILE_ROOTFS_SRC,
              f'{DIR_DST_ROOT}/boot/{image_name}/{image_name}.squashfs')
 
+        # copy saved config data and SSH keys
+        # owner restored on copy of config data by chmod_2775, above
+        copy_previous_installation_data(f'{DIR_DST_ROOT}/boot/{image_name}/rw')
+
         if is_raid_install(install_target):
             write_dir: str = f'{DIR_DST_ROOT}/boot/{image_name}/rw'
             raid.update_default(write_dir)