image: T5455: Add migration of SSH `known_hosts` files during image upgrade

During upgrade, the script now checks if any `known_hosts` files exist.
If so, it prompts the user to save these SSH fingerprints, and upon
confirmation, copies the files to the new image persistence directory.

4 files changed, 224 insertions, 29 deletions

diff --git a/src/conf_mode/system_login.py b/src/conf_mode/system_login.py
index 7ddcbff4a..702b04529 100755
--- a/src/conf_mode/system_login.py
+++ b/src/conf_mode/system_login.py
@@ -22,7 +22,6 @@ from copy import deepcopy
 from passlib.hosts import linux_context
 from psutil import users
 from pwd import getpwall
-from pwd import getpwnam
 from pwd import getpwuid
 from sys import exit
 from time import sleep
@@ -39,9 +38,13 @@ from vyos.template import is_ipv4
 from vyos.utils.auth import EPasswdStrength
 from vyos.utils.auth import evaluate_strength
 from vyos.utils.auth import get_current_user
+from vyos.utils.auth import get_local_users
+from vyos.utils.auth import get_user_home_dir
+from vyos.utils.auth import MIN_USER_UID
 from vyos.utils.configfs import delete_cli_node
 from vyos.utils.configfs import add_cli_node
 from vyos.utils.dict import dict_search
+from vyos.utils.file import move_recursive
 from vyos.utils.permission import chown
 from vyos.utils.process import cmd
 from vyos.utils.process import call
@@ -59,11 +62,7 @@ tacacs_nss_config_file = "/etc/tacplus_nss.conf"
 nss_config_file = "/etc/nsswitch.conf"
 login_motd_dsa_warning = r'/run/motd.d/92-vyos-user-dsa-deprecation-warning'
 
-# Minimum UID used when adding system users
-MIN_USER_UID: int = 1000
-# Maximim UID used when adding system users
-MAX_USER_UID: int = 59999
-# LOGIN_TIMEOUT from /etc/loign.defs minus 10 sec0
+# LOGIN_TIMEOUT from /etc/loign.defs minus 10 sec
 MAX_RADIUS_TIMEOUT: int = 50
 # MAX_RADIUS_TIMEOUT divided by 2 sec (minimum recomended timeout)
 MAX_RADIUS_COUNT: int = 8
@@ -71,29 +70,11 @@ MAX_RADIUS_COUNT: int = 8
 MAX_TACACS_COUNT: int = 8
 # Minimum USER id for TACACS users
 MIN_TACACS_UID = 900
-# List of local user accounts that must be preserved
-SYSTEM_USER_SKIP_LIST: list = ['radius_user', 'radius_priv_user', 'tacacs0', 'tacacs1',
-                              'tacacs2', 'tacacs3', 'tacacs4', 'tacacs5', 'tacacs6',
-                              'tacacs7', 'tacacs8', 'tacacs9', 'tacacs10',' tacacs11',
-                              'tacacs12', 'tacacs13', 'tacacs14', 'tacacs15']
 
 # As of OpenSSH 9.8p1 in Debian trixie, DSA keys are no longer supported
 SSH_DSA_DEPRECATION_WARNING: str = f'{SSH_DSA_DEPRECATION_WARNING} '\
 'The following users are using SSH-DSS keys for authentication.'
 
-def get_local_users(min_uid=MIN_USER_UID, max_uid=MAX_USER_UID):
-    """Return list of dynamically allocated users (see Debian Policy Manual)"""
-    local_users = []
-    for s_user in getpwall():
-        if getpwnam(s_user.pw_name).pw_uid < min_uid:
-            continue
-        if getpwnam(s_user.pw_name).pw_uid > max_uid:
-            continue
-        if s_user.pw_name in SYSTEM_USER_SKIP_LIST:
-            continue
-        local_users.append(s_user.pw_name)
-
-    return local_users
 
 def get_shadow_password(username):
     with open('/etc/shadow') as f:
@@ -389,9 +370,10 @@ def apply(login):
             tmp = dict_search('full_name', user_config)
             if tmp: command += f" --comment '{tmp}'"
 
-            tmp = dict_search('home_directory', user_config)
-            if tmp: command += f" --home '{tmp}'"
-            else: command += f" --home '/home/{user}'"
+            home_directory = dict_search('home_directory', user_config)
+            if not home_directory:
+                home_directory = f'/home/{user}'
+            command += f" --home '{home_directory}'"
 
             if 'operator' not in user_config:
                 command += f' --groups frr,frrvty,vyattacfg,sudo,adm,dip,disk,_kea'
@@ -404,7 +386,7 @@ def apply(login):
                 # crazy user will choose username root or any other system user which will fail.
                 #
                 # XXX: Should we deny using root at all?
-                home_dir = getpwnam(user).pw_dir
+                home_dir = get_user_home_dir(user)
                 # always re-render SSH keys with appropriate permissions
                 render(f'{home_dir}/.ssh/authorized_keys', 'login/authorized_keys.j2',
                        user_config, permission=0o600,
@@ -424,6 +406,17 @@ def apply(login):
             except Exception as e:
                 raise ConfigError(f'Adding user "{user}" raised exception: "{e}"')
 
+            # After invoking 'useradd' for each user, if /var/.users_backups/{user} exists, restore the
+            # backed up files to the newly created home directory. This reinstates the user's
+            # SSH environment and avoids loss of access or trust relationships due to the user
+            # creation process, which does not copy such custom files by default.
+            #
+            # More details: https://github.com/vyos/vyos-1x/pull/4678#pullrequestreview-3169648265
+            backup_directory = f"/var/.users_backups/{user}"
+            if command.startswith('useradd') and os.path.exists(backup_directory):
+                move_recursive(backup_directory, home_dir)
+                chown(home_dir, user=user, group='users', recursive=True)
+
             # T5875: ensure UID is properly set on home directory if user is re-added
             # the home directory will always exist, as it's created above by --create-home,
             # retrieve current owner of home directory and adjust on demand
@@ -459,7 +452,7 @@ def apply(login):
                 # Disable user to prevent re-login
                 call(f'usermod -s /sbin/nologin {user}')
 
-                home_dir = getpwnam(user).pw_dir
+                home_dir = get_user_home_dir(user)
                 # Remove SSH authorized keys file
                 authorized_keys_file = f'{home_dir}/.ssh/authorized_keys'
                 if os.path.exists(authorized_keys_file):
diff --git a/src/op_mode/image_installer.py b/src/op_mode/image_installer.py
index a5afd26f4..fa7640ac7 100755
--- a/src/op_mode/image_installer.py
+++ b/src/op_mode/image_installer.py
@@ -60,6 +60,8 @@ from vyos.utils.file import chmod_2775
 from vyos.utils.file import read_file
 from vyos.utils.file import write_file
 from vyos.utils.process import cmd, run, rc_cmd
+from vyos.utils.auth import get_local_users
+from vyos.utils.auth import get_user_home_dir
 from vyos.version import get_version_data
 
 # define text messages
@@ -719,6 +721,20 @@ def copy_ssh_host_keys() -> bool:
     return False
 
 
+def copy_ssh_known_hosts() -> bool:
+    """Ask user to copy SSH `known_hosts` files
+
+    Returns:
+        bool: user's decision
+    """
+    known_hosts_files = get_known_hosts_files()
+    msg = (
+        'Would you like to save the SSH known hosts (fingerprints) '
+        'from your current configuration?'
+    )
+    return known_hosts_files and ask_yes_no(msg, default=True)
+
+
 def console_hint() -> str:
     pid = getppid() if 'SUDO_USER' in environ else getpid()
     try:
@@ -1014,6 +1030,55 @@ def install_image() -> None:
         exit(1)
 
 
+def get_known_hosts_files(for_root=True, for_users=True) -> list:
+    """Collect all existing `known_hosts` files for root and/or users under /home"""
+
+    files = []
+
+    if for_root:
+        base_files = ('/root/.ssh/known_hosts', '/etc/ssh/ssh_known_hosts')
+        for file_path in base_files:
+            root_known_hosts = Path(file_path)
+            if root_known_hosts.exists():
+                files.append(root_known_hosts)
+
+    if for_users:  # for each non-system user
+        for user in get_local_users():
+            home_dir = Path(get_user_home_dir(user))
+            if home_dir.exists():
+                known_hosts = home_dir / '.ssh' / 'known_hosts'
+                if known_hosts.exists():
+                    files.append(known_hosts)
+
+    return files
+
+
+def migrate_known_hosts(target_dir: str):
+    """Copy `known_hosts` for root and all users to the new image directory"""
+
+    def _mkdir_and_copy_file(known_hosts_file, target_known_hosts):
+        target_known_hosts.parent.mkdir(parents=True, exist_ok=True)
+        copy(known_hosts_file, target_known_hosts)
+
+    # Copy root only files using default path
+    known_hosts_files = get_known_hosts_files(for_root=True, for_users=False)
+    for known_hosts_file in known_hosts_files:
+        target_known_hosts = Path(f'{target_dir}{known_hosts_file}')
+        _mkdir_and_copy_file(known_hosts_file, target_known_hosts)
+
+    # During image installation, backup critical user-specific files (e.g., known_hosts)
+    # from each user's home directory into /var/.users_backups/{user}. This ensures that their
+    # SSH configuration and trust relationships are preserved across system re-installations
+    # or provisioning.
+    # More details: https://github.com/vyos/vyos-1x/pull/4678#pullrequestreview-3169648265
+    known_hosts_files = get_known_hosts_files(for_root=False, for_users=True)
+    for known_hosts_file in known_hosts_files:
+        username = known_hosts_file.parent.parent.name
+        base_dir = Path(f'{target_dir}/var/.users_backups/{username}')
+        target_known_hosts = base_dir / '.ssh' / 'known_hosts'
+        _mkdir_and_copy_file(known_hosts_file, target_known_hosts)
+
+
 @compat.grub_cfg_update
 def add_image(image_path: str, vrf: str = None, username: str = '',
               password: str = '', no_prompt: bool = False, force: bool = False) -> None:
@@ -1149,6 +1214,11 @@ def add_image(image_path: str, vrf: str = None, username: str = '',
             for host_key in host_keys:
                 copy(host_key, target_ssh_dir)
 
+        target_ssh_known_hosts_dir: str = f'{root_dir}/boot/{image_name}/rw'
+        if no_prompt or copy_ssh_known_hosts():
+            print('Copying SSH known_hosts files')
+            migrate_known_hosts(target_ssh_known_hosts_dir)
+
         # copy system image and kernel files
         print('Copying system image files')
         for file in Path(f'{DIR_ISO_MOUNT}/live').iterdir():
diff --git a/src/tests/test_utils.py b/src/tests/test_utils.py
index 5c06a4394..12fda84d0 100644
--- a/src/tests/test_utils.py
+++ b/src/tests/test_utils.py
@@ -12,7 +12,12 @@
 # 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 pwd
 from unittest import TestCase
+
+from vyos.utils import auth
+
+
 class TestVyOSUtils(TestCase):
     def test_key_mangling(self):
         from vyos.utils.dict import mangle_dict_keys
@@ -36,3 +41,44 @@ class TestVyOSUtils(TestCase):
         self.assertEqual(list_strip(lst, rsb, right=True), ['a', 'b', 'c'])
         self.assertEqual(list_strip(lst, non), [])
         self.assertEqual(list_strip(sub, lst), [])
+
+class TestVyOSUtilsAuth(TestCase):
+
+    def test_get_local_users_returns_existing_usernames(self):
+        # Returned users exist, skip list is excluded, and UIDs are in range
+
+        all_users = set(s_user.pw_name for s_user in pwd.getpwall())
+        local_users = auth.get_local_users()
+
+        # All returned users must really exist
+        for user in local_users:
+            self.assertIn(user, all_users)
+
+        # Nobody in the skip list
+        for skipped in auth.SYSTEM_USER_SKIP_LIST:
+            self.assertNotIn(skipped, local_users)
+
+        # All are within UID range
+        for s_user in pwd.getpwall():
+            if s_user.pw_name in local_users:
+                self.assertGreaterEqual(s_user.pw_uid, auth.MIN_USER_UID)
+                self.assertLessEqual(s_user.pw_uid, auth.MAX_USER_UID)
+
+    def test_get_user_home_dir_for_real_user(self):
+        # User's homedir is a non-empty string for a valid user
+
+        local_users = auth.get_local_users()
+        if local_users:
+            for user in local_users:
+                home_dir = auth.get_user_home_dir(user)
+                self.assertIsInstance(home_dir, str)
+                self.assertTrue(bool(home_dir))  # Should not be empty
+        else:
+            self.skipTest("No suitable non-system users found on this system")
+
+    def test_get_user_home_dir_invalid_user(self):
+        # Raises KeyError for nonexistent username
+
+        user = "__this_user_does_not_exist__"  # Test using unlikely username
+        with self.assertRaises(KeyError):
+            auth.get_user_home_dir(user)
diff --git a/src/tests/test_utils_file.py b/src/tests/test_utils_file.py
new file mode 100644
index 000000000..4519cf03c
--- /dev/null
+++ b/src/tests/test_utils_file.py
@@ -0,0 +1,86 @@
+# Copyright VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# 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 tempfile
+import unittest
+
+from pathlib import Path
+
+from vyos.utils.file import copy_recursive
+from vyos.utils.file import move_recursive
+
+
+class TestVyOSUtilsFile(unittest.TestCase):
+    def setUp(self):
+        """Create temporary directories for source and destination."""
+        self.tmpdir = tempfile.TemporaryDirectory()
+        self.src = Path(self.tmpdir.name) / 'src'
+        self.dst = Path(self.tmpdir.name) / 'dst'
+
+        # Create test directory structure in `src`
+        (self.src / 'subdir').mkdir(parents=True)
+        (self.dst).mkdir(parents=True)
+
+        # Create files
+        (self.src / 'file1.txt').write_text('hello world')
+        (self.src / 'subdir' / 'file2.txt').write_text('subdir file')
+
+    def tearDown(self):
+        """Cleanup temp directory."""
+        self.tmpdir.cleanup()
+
+    def test_copy_recursive_no_overwrite(self):
+        copy_recursive(str(self.src), str(self.dst), overwrite=False)
+
+        self.assertTrue((self.dst / 'file1.txt').exists())
+        self.assertTrue((self.dst / 'subdir' / 'file2.txt').exists())
+
+    def test_copy_recursive_skip_existing(self):
+        # Create conflicting file in destination
+        (self.dst / 'file1.txt').write_text('different content')
+
+        copy_recursive(str(self.src), str(self.dst), overwrite=False)
+
+        # Destination should remain the same (not overwritten)
+        content = (self.dst / 'file1.txt').read_text()
+        self.assertEqual(content, 'different content')
+
+    def test_copy_recursive_overwrite(self):
+        (self.dst / 'file1.txt').write_text('different content')
+
+        copy_recursive(str(self.src), str(self.dst), overwrite=True)
+
+        # Destination should be overwritten with source content
+        content = (self.dst / 'file1.txt').read_text()
+        self.assertEqual(content, 'hello world')
+
+    def test_move_recursive(self):
+        move_recursive(str(self.src), str(self.dst), overwrite=False)
+
+        # Files should appear in destination
+        self.assertTrue((self.dst / 'file1.txt').exists())
+        self.assertTrue((self.dst / 'subdir' / 'file2.txt').exists())
+
+        # Source should be removed
+        self.assertFalse(self.src.exists())
+
+    def test_move_recursive_overwrite(self):
+        # Prepare conflicting file in destination
+        (self.dst / 'file1.txt').write_text('conflicting')
+
+        move_recursive(str(self.src), str(self.dst), overwrite=True)
+
+        content = (self.dst / 'file1.txt').read_text()
+        self.assertEqual(content, 'hello world')  # overwritten
+        self.assertFalse(self.src.exists())  # source cleaned up