# Copyright 2026 Cloudbase Solutions Srl
# All Rights Reserved.

import json
import os
from unittest import mock

from coriolis import constants
from coriolis import exception
from coriolis.osmorphing.osmount import base
from coriolis.osmorphing.osmount import luks_mixin
from coriolis.tests import test_base


class ConcreteLinuxLUKSMixin(
    luks_mixin.LinuxLUKSMixin, base.BaseSSHOSMountTools
):
    def check_os(self):
        pass

    def mount_os(self):
        pass

    def dismount_os(self):
        pass

    def run_user_script(self, user_script):
        pass


_CONN_INFO = {"ip": "127.0.0.1", "username": "foo", "password": "lish"}
_OS_ROOT_DIR = "/mnt/os"
_DEV = "/dev/sda"
_PASSPHRASE = "dont-dead-open-inside"
_UUID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
_KEYFILE = "/etc/luks/coriolis_sda.key"


class LinuxLUKSMixinTestCase(test_base.CoriolisBaseTestCase):
    @mock.patch.object(base.BaseSSHOSMountTools, "_connect")
    def setUp(self, mock_connect):
        super().setUp()
        self.event_manager = mock.MagicMock()
        self.osmorphing_info = {constants.ENCRYPTED_DISKS_PASS: _PASSPHRASE}
        self.mixin = ConcreteLinuxLUKSMixin(
            _CONN_INFO,
            self.event_manager,
            [],
            30,
            osmorphing_info=self.osmorphing_info,
        )
        self.mixin._ssh = mock.MagicMock()

    def test_validate_config_pcrs(self):
        luks_mixin.CONF.set_override('tpm2_pcrs', [], group='luks')
        self.addCleanup(
            luks_mixin.CONF.clear_override, 'tpm2_pcrs', group='luks')

        self.assertRaises(
            exception.CoriolisException,
            luks_mixin.validate_config)

        luks_mixin.CONF.set_override('tpm2_pcrs', [4, 7], group='luks')

        # Should not raise an exception anymore.
        luks_mixin.validate_config()

        # PCR indexes outside of the valid 0-15 range are rejected by the
        # option's item_type as soon as they're set.
        self.assertRaises(
            ValueError,
            luks_mixin.CONF.set_override, 'tpm2_pcrs', ['16'], group='luks')

        self.assertRaises(
            ValueError,
            luks_mixin.CONF.set_override, 'tpm2_pcrs', ['foo'], group='luks')

    def test_render_tpm2_pcrs(self):
        luks_mixin.CONF.set_override('tpm2_pcrs', ['7', '11'], group='luks')
        self.addCleanup(
            luks_mixin.CONF.clear_override, 'tpm2_pcrs', group='luks')

        dracut_out = luks_mixin._render_tpm2_pcrs(
            luks_mixin._LUKS_FIRSTBOOT_SCRIPTS["dracut"], "dracut")
        self.assertIn("--tpm2-pcrs=7:sha256+11:sha256", dracut_out)

        clevis_out = luks_mixin._render_tpm2_pcrs(
            luks_mixin._LUKS_FIRSTBOOT_SCRIPTS["update-initramfs"],
            "update-initramfs")
        self.assertIn('"pcr_ids":"7,11"', clevis_out)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_unlock_luks_device")
    def test__unlock_luks_devices(self, mock_unlock):
        mock_unlock.return_value = "/dev/mapper/coriolis_sda"
        dev_paths = [_DEV]

        self.mixin._unlock_luks_devices(dev_paths)

        self.assertEqual(dev_paths, ["/dev/mapper/coriolis_sda"])
        self.assertEqual(self.mixin._luks_opened, [("coriolis_sda", _DEV)])
        mock_unlock.assert_called_once_with(_DEV, _PASSPHRASE)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_close_luks_devices")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_unlock_luks_device")
    def test__unlock_luks_devices_cleans_up_on_error(
        self, mock_unlock, mock_close
    ):
        _DEV2 = "/dev/sdb"
        mock_unlock.side_effect = [
            "/dev/mapper/coriolis_sda",
            exception.CoriolisException("bad passphrase"),
        ]

        self.assertRaises(
            exception.CoriolisException,
            self.mixin._unlock_luks_devices,
            [_DEV, _DEV2],
        )

        mock_close.assert_called_once_with()

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_unlock_luks_device")
    def test__unlock_luks_devices_skips_non_luks(self, mock_unlock):
        mock_unlock.return_value = None
        dev_paths = [_DEV]

        self.mixin._unlock_luks_devices(dev_paths)

        self.assertEqual(dev_paths, [_DEV])
        self.assertEqual(self.mixin._luks_opened, [])

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_is_luks")
    def test__unlock_luks_device(self, mock_is_luks):
        # not LUKS, skipped regardless of passphrase.
        mock_is_luks.return_value = False
        self.assertIsNone(self.mixin._unlock_luks_device(_DEV, None))
        self.assertIsNone(self.mixin._unlock_luks_device(_DEV, _PASSPHRASE))

        # is LUKS, no password.
        mock_is_luks.return_value = True
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._unlock_luks_device,
            _DEV,
            None,
        )

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_auth_luks")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_is_luks")
    def test__unlock_luks_device_success(
        self, mock_is_luks, mock_auth_luks, mock_exec_cmd
    ):
        mock_is_luks.return_value = True

        result = self.mixin._unlock_luks_device(_DEV, _PASSPHRASE)

        self.assertEqual(result, "/dev/mapper/coriolis_sda")
        mock_auth_luks.assert_called_once_with(
            _PASSPHRASE, "/tmp/coriolis_sda.key"
        )
        mock_exec_cmd.assert_called_once_with(
            "sudo cryptsetup luksOpen --disable-keyring "
            "--key-file /tmp/coriolis_sda.key %s coriolis_sda" % _DEV
        )

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__is_luks(self, mock_exec_cmd):
        # True.
        mock_exec_cmd.return_value = ""
        self.assertTrue(self.mixin._is_luks(_DEV))
        mock_exec_cmd.assert_called_once_with(
            "sudo cryptsetup isLuks %s" % _DEV
        )

        # False (exit code 1: not a LUKS device).
        mock_exec_cmd.side_effect = exception.SSHCommandFailed(
            "boom goes the dynamite", exit_code=1)
        self.assertFalse(self.mixin._is_luks(_DEV))

        # Unexpected exit code (e.g. 4: device does not exist or access
        # denied), hard failure.
        mock_exec_cmd.side_effect = exception.SSHCommandFailed(
            "boom goes a different dynamite", exit_code=4)
        self.assertRaises(
            exception.SSHCommandFailed, self.mixin._is_luks, _DEV)

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__close_luks_devices(self, mock_exec_cmd):
        self.mixin._luks_opened = [
            ("coriolis_sda", "/dev/sda"),
            ("coriolis_sdb", "/dev/sdb"),
        ]
        self.mixin._close_luks_devices()

        mock_exec_cmd.assert_any_call(
            "sudo cryptsetup luksClose coriolis_sda || true"
        )
        mock_exec_cmd.assert_any_call(
            "sudo cryptsetup luksClose coriolis_sdb || true"
        )
        self.assertEqual(self.mixin._luks_opened, [])

    @mock.patch.object(luks_mixin.utils, "write_ssh_file")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__write_remote_file(self, mock_exec_cmd, mock_write):
        mock_exec_cmd.return_value = "/foo/lish"
        self.mixin._write_remote_file("/bar/tender", "content")

        mock_write.assert_called_once_with(
            self.mixin._ssh, "/foo/lish", b"content"
        )
        mock_exec_cmd.assert_any_call("sudo mv /foo/lish /bar/tender")

        # With mode.
        self.mixin._write_remote_file("/ness/dante", "content", mode="600")

        mock_exec_cmd.assert_any_call(
            "sudo mv /foo/lish /ness/dante && sudo chmod 600 /ness/dante"
        )

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    def test__auth_luks(self, mock_write, mock_exec_cmd):
        with self.assertRaises(Exception):
            with self.mixin._auth_luks(_PASSPHRASE, "/tmp/ting") as key_path:
                self.assertEqual(key_path, "/tmp/ting")
                mock_write.assert_called_once_with("/tmp/ting", _PASSPHRASE)

                # raise an exception, cleanup should be performed.
                raise Exception("al code only sometimes breaks.")

        mock_exec_cmd.assert_called_once_with("sudo rm -f /tmp/ting")

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__get_tpm2_token_info(self, mock_exec_cmd):
        # Exception.
        mock_exec_cmd.return_value = ""

        with self.assertLogs(
            "coriolis.osmorphing.osmount.luks_mixin", level="WARNING"
        ):
            result = self.mixin._get_tpm2_token_info(_DEV)

        self.assertEqual(result, [])

        # Only the last entry is valid / has keyslots.
        header = {
            "tokens": {
                "0": {"type": "systemd-tpm2", "keyslots": []},
                "1": {"type": "luks2-keyring", "keyslots": ["2"]},
                "2": {"type": "systemd-tpm2", "keyslots": ["1"]},
            }
        }
        mock_exec_cmd.return_value = json.dumps(header)

        result = self.mixin._get_tpm2_token_info(_DEV)

        self.assertEqual(result, [("2", "1")])

        # Empty dump.
        mock_exec_cmd.return_value = json.dumps({})

        result = self.mixin._get_tpm2_token_info(_DEV)
        self.assertEqual(result, [])

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_auth_luks")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_get_tpm2_token_info")
    def test__remove_tpm2_tokens(self, mock_info, mock_exec_cmd, mock_auth):
        # No tokens.
        mock_info.return_value = []
        self.mixin._remove_tpm2_tokens(_DEV, _PASSPHRASE)

        # Remove token.
        mock_info.return_value = [("0", "1")]
        self.mixin._remove_tpm2_tokens(_DEV, _PASSPHRASE)

        mock_exec_cmd.assert_any_call(
            "sudo cryptsetup token remove --token-id 0 %s" % _DEV
        )
        mock_exec_cmd.assert_any_call(
            "sudo cryptsetup luksKillSlot --key-file /tmp/coriolis_sda.key "
            "%s 1" % _DEV
        )
        mock_auth.assert_called_once()

        # Token removal failed.
        mock_exec_cmd.side_effect = Exception("toe ken remove failed")
        with self.assertLogs(
            "coriolis.osmorphing.osmount.luks_mixin", level="WARNING"
        ):
            self.mixin._remove_tpm2_tokens(_DEV, _PASSPHRASE)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    @mock.patch.object(luks_mixin.utils, "read_ssh_file")
    @mock.patch.object(luks_mixin.utils, "test_ssh_path")
    def test__transform_crypttab(self, mock_test, mock_read, mock_write):
        # No file.
        mock_test.return_value = False

        result = self.mixin._transform_crypttab(_OS_ROOT_DIR, None)
        self.assertFalse(result)

        # No changes.
        mock_test.return_value = True
        mock_read.return_value = b"# comment\nluks-root UUID=aaa none none\n"
        result = self.mixin._transform_crypttab(_OS_ROOT_DIR, lambda p: None)
        self.assertFalse(result)
        mock_write.assert_not_called()

        # File changed.
        mock_read.return_value = b"luks-root UUID=aaa none none\n"

        def _set_opts(parts):
            parts[3] = 'new-opt'
            return parts

        result = self.mixin._transform_crypttab(_OS_ROOT_DIR, _set_opts)
        self.assertTrue(result)
        mock_write.assert_called_once()

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    @mock.patch.object(luks_mixin.utils, "read_ssh_file")
    @mock.patch.object(luks_mixin.utils, "test_ssh_path", return_value=True)
    def test__remove_tpm2_crypttab_options(
        self, mock_test_path, mock_read, mock_write
    ):
        line = (
            "luks-root UUID=%s none tpm2-device=auto,x-initrd.attach\n" % _UUID
        )
        mock_read.return_value = line.encode("utf-8")

        with self.assertLogs(
            "coriolis.osmorphing.osmount.luks_mixin", level="INFO"
        ):
            self.mixin._remove_tpm2_crypttab_options(_OS_ROOT_DIR)

        new_content = mock_write.call_args[0][1]
        self.assertNotIn("tpm2-", new_content)
        self.assertIn("x-initrd.attach", new_content)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_remove_tpm2_tokens")
    @mock.patch.object(
        luks_mixin.LinuxLUKSMixin, "_remove_tpm2_crypttab_options"
    )
    def test_remove_encryption_artifacts(self, mock_opts, mock_tokens):
        # No LUKS devices open.
        self.mixin._luks_opened = []
        self.mixin.remove_encryption_artifacts(_OS_ROOT_DIR)
        mock_tokens.assert_not_called()
        mock_opts.assert_not_called()

        # Opened LUKS device.
        self.mixin._luks_opened = [("coriolis_sda", _DEV)]
        self.mixin.remove_encryption_artifacts(_OS_ROOT_DIR)

        mock_tokens.assert_called_once_with(_DEV, _PASSPHRASE)
        mock_opts.assert_called_once_with(_OS_ROOT_DIR)
        self.event_manager.progress_update.assert_called()

    def test__get_migration_keyfile_path(self):
        result = self.mixin._get_migration_keyfile_path("/foo/lish")

        self.assertEqual(
            os.path.join(luks_mixin._LUKS_KEYFILE_DIR, "coriolis_lish.key"),
            result,
        )

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__get_luks_uuid(self, mock_exec_cmd):
        mock_exec_cmd.return_value = "  %s  \n" % _UUID

        result = self.mixin._get_luks_uuid(_DEV)

        self.assertEqual(result, _UUID)
        mock_exec_cmd.assert_called_once_with(
            "sudo cryptsetup luksUUID %s" % _DEV
        )

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_detect_initramfs_tool")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_transform_crypttab")
    def test__update_crypttab_keyfile(self, mock_transform, mock_detect_tool):
        mock_transform.return_value = True
        mock_detect_tool.return_value = "update-initramfs"
        self.mixin._update_crypttab_keyfile(_OS_ROOT_DIR, {_UUID: _KEYFILE})
        mock_transform.assert_called_once()

        # Extract and exercise the inner _set_keyfile transform.
        transform = mock_transform.call_args[0][1]

        # too few parts.
        self.assertIsNone(transform(["luks-root"]))

        # no UUID match.
        parts = ["luks-root", "PARTUUID=abc", "none", "none"]
        self.assertIsNone(transform(parts))

        # UUID not in map.
        parts = [
            "luks-root",
            "UUID=00000000-0000-0000-0000-000000000000",
            "none",
            "none",
        ]
        self.assertIsNone(transform(parts))

        # UUID= format match; update-initramfs -> 'initramfs' option added.
        result = transform(["luks-root", "UUID=%s" % _UUID, "none", "none"])
        self.assertEqual(result[2], _KEYFILE)
        self.assertIn("initramfs", result[3].split(","))

        # Case-insensitive UUID match: uppercase UUID in crypttab still matches
        # lowercase key in the map.
        result = transform(
            ["luks-root", "UUID=%s" % _UUID.upper(), "none", "none"])
        self.assertEqual(result[2], _KEYFILE)

        # /by-uuid/ path also matches.
        parts = ["luks-root", "/dev/disk/by-uuid/%s" % _UUID, "none", "none"]
        result = transform(parts)
        self.assertEqual(result[2], _KEYFILE)

        # dracut -> 'initramfs' option NOT added.
        mock_transform.reset_mock()
        mock_detect_tool.return_value = "dracut"
        self.mixin._update_crypttab_keyfile(_OS_ROOT_DIR, {_UUID: _KEYFILE})
        transform_dracut = mock_transform.call_args[0][1]
        result = transform_dracut(
            ["luks-root", "UUID=%s" % _UUID, "none", "none"])
        self.assertEqual(result[2], _KEYFILE)
        self.assertNotIn("initramfs", result[3].split(","))

        # transform returns False -> exception.
        mock_transform.return_value = False
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._update_crypttab_keyfile,
            _OS_ROOT_DIR,
            {_UUID: _KEYFILE},
        )

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_detect_initramfs_tool")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_configure_dracut_keyfiles")
    @mock.patch.object(
        luks_mixin.LinuxLUKSMixin, "_configure_initramfs_tools_keyfiles"
    )
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_update_crypttab_keyfile")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_get_luks_uuid")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__write_migration_keyfiles(
        self,
        mock_exec,
        mock_uuid,
        mock_write_file,
        mock_update_ct,
        mock_cfg_initramfs,
        mock_cfg_dracut,
        mock_detect_tool,
    ):
        # No passphrase / opened devices.
        mock_uuid.return_value = _UUID
        self.mixin._osmorphing_info = {}
        self.mixin._luks_opened = []
        self.mixin._write_migration_keyfiles(_OS_ROOT_DIR)
        mock_exec.assert_not_called()

        # dracut branch.
        self.mixin._osmorphing_info = {
            constants.ENCRYPTED_DISKS_PASS: _PASSPHRASE,
        }
        self.mixin._luks_opened = [("coriolis_sda", _DEV)]
        mock_detect_tool.return_value = "dracut"
        self.mixin._write_migration_keyfiles(_OS_ROOT_DIR)

        keyfile_path = os.path.join(
            luks_mixin._LUKS_KEYFILE_DIR,
            "coriolis_%s.key" % os.path.basename(_DEV),
        )
        expected_abs_path = os.path.join(
            _OS_ROOT_DIR, keyfile_path.lstrip("/")
        )
        mock_write_file.assert_called_once_with(
            expected_abs_path, _PASSPHRASE, mode="400"
        )
        mock_update_ct.assert_called_once_with(_OS_ROOT_DIR, {_UUID: _KEYFILE})
        mock_cfg_dracut.assert_called_once_with(
            _OS_ROOT_DIR, {_UUID: _KEYFILE}
        )
        mock_cfg_initramfs.assert_not_called()

        # update-initramfs branch.
        mock_cfg_dracut.reset_mock()
        mock_update_ct.reset_mock()
        mock_write_file.reset_mock()
        mock_detect_tool.return_value = "update-initramfs"
        self.mixin._write_migration_keyfiles(_OS_ROOT_DIR)

        mock_cfg_initramfs.assert_called_once_with(_OS_ROOT_DIR)
        mock_cfg_dracut.assert_not_called()

        # No tool found: CoriolisException raised.
        mock_detect_tool.return_value = None
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._write_migration_keyfiles,
            _OS_ROOT_DIR,
        )

        # _get_luks_uuid raises: wrapped in CoriolisException.
        mock_detect_tool.return_value = "dracut"
        mock_uuid.side_effect = Exception("cryptsetup went boom")
        self.mixin._luks_opened = [("coriolis_sda", _DEV)]
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._write_migration_keyfiles,
            _OS_ROOT_DIR,
        )

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.utils, "test_ssh_path")
    def test__configure_dracut_keyfiles(
        self, mock_test_path, mock_exec, mock_write
    ):
        plugin_path = luks_mixin._CRYPTSETUP_TPM2_PLUGIN_PATHS[0]
        plugin_abs = os.path.join(_OS_ROOT_DIR, plugin_path.lstrip("/"))
        conf_abs = os.path.join(
            _OS_ROOT_DIR, luks_mixin._DRACUT_LUKS_CONF_PATH.lstrip("/")
        )

        # no TPM2 plugin.
        mock_test_path.return_value = False
        self.mixin._configure_dracut_keyfiles(_OS_ROOT_DIR, {_UUID: _KEYFILE})

        written = mock_write.call_args[0][1]
        self.assertIn(_KEYFILE, written)
        self.assertIn("install_items+=", written)
        self.assertIn('hostonly="no"', written)
        self.assertIn('add_dracutmodules+=" crypt "', written)
        self.assertNotIn(plugin_path, written)
        mock_exec.assert_called_once_with(
            "sudo chown root:root %s && sudo chmod 644 %s"
            % (conf_abs, conf_abs)
        )

        # with TPM2 plugin.
        mock_exec.reset_mock()
        mock_write.reset_mock()
        mock_test_path.side_effect = lambda _ssh, path: path == plugin_abs
        self.mixin._configure_dracut_keyfiles(_OS_ROOT_DIR, {_UUID: _KEYFILE})

        written = mock_write.call_args[0][1]
        self.assertIn(plugin_path, written)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_write_remote_file")
    @mock.patch.object(luks_mixin.utils, "read_ssh_file")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.utils, "test_ssh_path")
    def test__configure_initramfs_tools_keyfiles(
        self, mock_test, mock_exec, mock_read, mock_write
    ):
        hook_abs = os.path.join(
            _OS_ROOT_DIR, "etc/cryptsetup-initramfs/conf-hook"
        )

        # existing conf-hook: content is preserved and KEYFILE_PATTERN
        # appended.
        mock_test.side_effect = lambda _ssh, path: path == hook_abs
        mock_read.return_value = b"# existing\n"

        self.mixin._configure_initramfs_tools_keyfiles(_OS_ROOT_DIR)

        written_content = mock_write.call_args[0][1]
        self.assertIn("# existing", written_content)
        self.assertIn("KEYFILE_PATTERN=", written_content)
        self.assertIn("/etc/luks/coriolis_*.key", written_content)

        # no existing conf-hook: KEYFILE_PATTERN still written.
        mock_test.side_effect = None
        mock_test.return_value = False
        mock_write.reset_mock()

        self.mixin._configure_initramfs_tools_keyfiles(_OS_ROOT_DIR)

        written_content = mock_write.call_args[0][1]
        self.assertIn("KEYFILE_PATTERN=", written_content)

    @mock.patch.object(luks_mixin.utils, "test_ssh_path")
    def test__detect_initramfs_tool(self, mock_test):
        # update-initramfs.
        mock_test.side_effect = lambda _ssh, path: "update-initramfs" in path
        result = self.mixin._detect_initramfs_tool(_OS_ROOT_DIR)
        self.assertEqual(result, "update-initramfs")

        # dracut.
        mock_test.side_effect = lambda _ssh, path: "dracut" in path
        result = self.mixin._detect_initramfs_tool(_OS_ROOT_DIR)
        self.assertEqual(result, "dracut")

        # None.
        mock_test.side_effect = None
        mock_test.return_value = False
        result = self.mixin._detect_initramfs_tool(_OS_ROOT_DIR)
        self.assertIsNone(result)

    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    @mock.patch.object(luks_mixin.utils, "test_ssh_path")
    def test__build_dracut_include_args(self, mock_test, mock_exec):
        # no crypttab.
        mock_test.return_value = False
        mock_exec.return_value = ""
        result = self.mixin._build_dracut_include_args(_OS_ROOT_DIR)
        self.assertEqual(result, [])

        # with crypttab and keyfile.
        crypttab_path = os.path.join(_OS_ROOT_DIR, "etc/crypttab")
        luks_dir = os.path.join(_OS_ROOT_DIR, "etc/luks")
        keyfile_abs = os.path.join(luks_dir, "coriolis_sda.key")

        mock_test.side_effect = lambda _ssh, path: path == crypttab_path
        mock_exec.return_value = keyfile_abs
        result = self.mixin._build_dracut_include_args(_OS_ROOT_DIR)

        expected = [
            "--include",
            "/etc/crypttab",
            "/etc/crypttab",
            "--include",
            "/etc/luks/coriolis_sda.key",
            "/etc/luks/coriolis_sda.key",
        ]
        self.assertEqual(result, expected)

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_rebuild_initramfs_dracut")
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_detect_initramfs_tool")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__rebuild_initramfs(
        self, mock_exec, mock_detect, mock_rebuild_dracut
    ):
        # update-initramfs.
        mock_detect.return_value = "update-initramfs"
        self.mixin._rebuild_initramfs(_OS_ROOT_DIR)
        mock_exec.assert_called_once_with(
            "sudo chroot %s update-initramfs -u -k all" % _OS_ROOT_DIR
        )
        mock_rebuild_dracut.assert_not_called()

        # dracut: delegates to _rebuild_initramfs_dracut.
        mock_exec.reset_mock()
        mock_detect.return_value = "dracut"
        self.mixin._rebuild_initramfs(_OS_ROOT_DIR)

        mock_rebuild_dracut.assert_called_once_with(_OS_ROOT_DIR)
        mock_exec.assert_not_called()

        # no tool found.
        mock_detect.return_value = None
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._rebuild_initramfs,
            _OS_ROOT_DIR,
        )

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, "_build_dracut_include_args")
    @mock.patch.object(base.BaseSSHOSMountTools, "_exec_cmd")
    def test__rebuild_initramfs_dracut(self, mock_exec, mock_include_args):
        # ls raises: treated as empty kernel list -> CoriolisException.
        mock_exec.side_effect = Exception("ls failed")
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._rebuild_initramfs_dracut,
            _OS_ROOT_DIR,
        )
        mock_include_args.assert_called_once_with(_OS_ROOT_DIR)

        # ls succeeds but returns empty output -> CoriolisException.
        mock_exec.reset_mock()
        mock_exec.side_effect = None
        mock_exec.return_value = "  \n\n"
        self.assertRaises(
            exception.CoriolisException,
            self.mixin._rebuild_initramfs_dracut,
            _OS_ROOT_DIR,
        )

        # initramfs rebuilt.
        mock_include_args.return_value = [
            "--include",
            "/etc/crypttab",
            "/etc/crypttab",
        ]

        mock_exec.reset_mock()
        mock_exec.side_effect = [
            "5.15.0-generic\n",  # ls /lib/modules/
            None,                # dracut call
        ]
        self.mixin._rebuild_initramfs_dracut(_OS_ROOT_DIR)

        self.assertEqual(mock_exec.call_count, 2)
        mock_exec.assert_any_call(
            "sudo chroot '%s' dracut -f --kver '5.15.0-generic' "
            "--no-hostonly --add crypt --add-drivers dm-crypt "
            "--include /etc/crypttab /etc/crypttab" % _OS_ROOT_DIR
        )

    @mock.patch.object(luks_mixin, '_render_tpm2_pcrs')
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_detect_initramfs_tool')
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_rebuild_initramfs')
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_fix_grub_luks_root')
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_write_migration_keyfiles')
    def test_install_encryption_firstboot_setup(
        self, mock_write_keyfiles, mock_grub, mock_rebuild, mock_detect_tool,
        mock_render_tpm2_pcrs,
    ):
        mock_morphing_tools = mock.MagicMock()

        # No LUKS opened: early return, nothing called.
        self.mixin._luks_opened = []
        self.mixin.install_encryption_firstboot_setup(
            _OS_ROOT_DIR, mock_morphing_tools)
        mock_write_keyfiles.assert_not_called()
        mock_morphing_tools.register_firstboot_script.assert_not_called()

        # LUKS opened, dracut.
        self.mixin._luks_opened = [("coriolis_sda", _DEV)]
        mock_detect_tool.return_value = "dracut"

        self.mixin.install_encryption_firstboot_setup(
            _OS_ROOT_DIR, mock_morphing_tools)

        mock_write_keyfiles.assert_called_once_with(_OS_ROOT_DIR)
        mock_grub.assert_called_once_with(_OS_ROOT_DIR)
        mock_rebuild.assert_called_once_with(_OS_ROOT_DIR)

        mock_render_tpm2_pcrs.assert_called_once_with(
            luks_mixin._LUKS_FIRSTBOOT_SCRIPTS["dracut"], "dracut")
        mock_morphing_tools.register_firstboot_script.assert_called_once_with(
            mock_render_tpm2_pcrs.return_value,
            user_provided=False,
            script_filename="luks-firstboot.sh",
        )

        # No initramfs tool found: CoriolisException.
        mock_detect_tool.return_value = None
        self.assertRaises(
            exception.CoriolisException,
            self.mixin.install_encryption_firstboot_setup,
            _OS_ROOT_DIR,
            mock_morphing_tools,
        )

    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_write_remote_file')
    @mock.patch.object(luks_mixin.LinuxLUKSMixin, '_get_luks_uuid')
    @mock.patch.object(base.BaseSSHOSMountTools, '_exec_cmd')
    @mock.patch.object(luks_mixin.utils, 'read_ssh_file')
    @mock.patch.object(luks_mixin.utils, 'test_ssh_path')
    def test__fix_grub_luks_root_patches_grub(
        self,
        mock_test_path,
        mock_read_file,
        mock_exec,
        mock_get_luks_uuid,
        mock_write_file,
    ):
        # no opened LUKS devices.
        self.mixin._luks_opened = []
        self.mixin._fix_grub_luks_root(_OS_ROOT_DIR)
        mock_test_path.assert_not_called()

        # no crypttab.
        self.mixin._luks_opened = [("coriolis_sda", _DEV)]
        mock_test_path.return_value = False

        self.mixin._fix_grub_luks_root(_OS_ROOT_DIR)

        crypttab_path = os.path.join(_OS_ROOT_DIR, "etc/crypttab")
        mock_test_path.assert_called_once_with(self.mixin._ssh, crypttab_path)
        mock_read_file.assert_not_called()

        # no uuid_to_crypttab_name.
        mock_test_path.return_value = True
        mock_read_file.return_value = "".encode("utf-8")

        self.mixin._fix_grub_luks_root(_OS_ROOT_DIR)

        mock_read_file.assert_called_once_with(self.mixin._ssh, crypttab_path)
        mock_get_luks_uuid.assert_not_called()

        # no replacements.
        crypttab = 'luks-root UUID=%s none none\n' % _UUID.lower()
        grub_content = (
            'set root=/dev/mapper/coriolis_sda\n'
            'linux /vmlinuz root=/dev/mapper/coriolis_sda\n'
        )
        mock_read_file.return_value = crypttab.encode('utf-8')
        mock_get_luks_uuid.return_value = ""

        self.mixin._fix_grub_luks_root(_OS_ROOT_DIR)

        mock_get_luks_uuid.assert_called_once_with(_DEV)
        mock_exec.assert_not_called()

        # with replacements.
        mock_get_luks_uuid.return_value = _UUID
        grub_path = os.path.join(_OS_ROOT_DIR, 'boot/grub/grub.cfg')
        crypttab_path = os.path.join(_OS_ROOT_DIR, 'etc/crypttab')
        mock_test_path.side_effect = lambda _ssh, path: (
            path in (crypttab_path, grub_path)
        )
        mock_exec.return_value = grub_content

        self.mixin._fix_grub_luks_root(_OS_ROOT_DIR)

        written = mock_write_file.call_args[0][1]
        self.assertIn('/dev/mapper/luks-root', written)
        self.assertNotIn('/dev/mapper/coriolis_sda', written)
