Ensure that Bitlocker is resumed in case of failures

We need to resume BitLocker if the os-morphing process fails,
otherwise the disks will remain publicly open.

"install_encryption_firstboot_setup" is the last method called
during os-morphing, we can suspend Bitlocker there and resume
it in case of failures.

While at it, we'll move "_unlock_encrypted_volumes" next to
"_unlock_encrypted_volume".

coriolis/osmorphing/osmount/windows.py

@@ -242,20 +242,6 @@ class WindowsMountTools(base.BaseOSMountTools):
             f'manage-bde -unlock "{volume_id}" '
             f'-RecoveryPassword "{recovery_password}"')
 
-    def _suspend_bitlocker(self, volume_id: str):
-        """Suspend BitLocker until the next reboot for a given volume.
-
-        It doesn't decrypt the device, it just adds a publicly accessible
-        BitLocker protector that automatically unlocks the volume.
-
-        When the replica instance boots, the TPM protector will be reconfigured
-        automatically. Unfortunately the '-RebootCount' parameter isn't
-        honored, perhaps due to the fact that the disks are attached to a
-        separate VM. For this reason, we'll use a first-boot script to resume
-        BitLocker explicitly.
-        """
-        self._conn.exec_ps_command(f'Suspend-BitLocker "{volume_id}"')
-
     def _unlock_encrypted_volumes(self):
         recovery_password = self._osmorphing_info.get(
             constants.ENCRYPTED_DISKS_PASS)
@@ -285,12 +271,6 @@ class WindowsMountTools(base.BaseOSMountTools):
                     "recovery password.",
                     encrypted_volume_id)
                 continue
-
-            # Suspend BitLocker until the replica boots.
-            #
-            # We'll intentionally propagate the failure if we managed to
-            # unlock the volume but failed to suspend BitLocker.
-            self._suspend_bitlocker(encrypted_volume_id)
             self._unlocked_volumes.append(encrypted_volume_id)
 
         if not unlocked:
@@ -298,6 +278,23 @@ class WindowsMountTools(base.BaseOSMountTools):
                 "Could not unlock any volume using the specified "
                 "BitLocker recovery password.")
 
+    def _suspend_bitlocker(self, volume_id: str):
+        """Suspend BitLocker until the next reboot for a given volume.
+
+        It doesn't decrypt the device, it just adds a publicly accessible
+        BitLocker protector that automatically unlocks the volume.
+
+        When the replica instance boots, the TPM protector will be reconfigured
+        automatically. Unfortunately the '-RebootCount' parameter isn't
+        honored, perhaps due to the fact that the disks are attached to a
+        separate VM. For this reason, we'll use a first-boot script to resume
+        BitLocker explicitly.
+        """
+        self._conn.exec_ps_command(f'Suspend-BitLocker "{volume_id}"')
+
+    def _resume_bitlocker(self, volume_id: str):
+        self._conn.exec_ps_command(f'Resume-BitLocker "{volume_id}"')
+
     def install_encryption_firstboot_setup(
         self,
         os_root_dir,
@@ -313,18 +310,36 @@ class WindowsMountTools(base.BaseOSMountTools):
         # isn't honored, perhaps due to the fact that the disks are attached
         # to a different VM.
         script_content = ""
-        for encrypted_volume_id in self._unlocked_volumes:
-            LOG.info(
-                "Resuming BitLocker after first boot, volume: %s",
-                encrypted_volume_id)
-            script_content += f'Resume-BitLocker "{encrypted_volume_id}"\r\n'
-
-        # Resume BitLocker after bringing the disks online, which has a script
-        # priority of 10.
-        os_morphing_tools.register_firstboot_script(
-            script_content,
-            user_provided=False,
-            script_filename="11-bitlocker-firstboot.ps1")
+
+        try:
+            for encrypted_volume_id in self._unlocked_volumes:
+                LOG.info(
+                    "Suspending BitLocker for volume %s, scheduling it to be "
+                    "resumed after first-boot.",
+                    encrypted_volume_id)
+                # Suspend BitLocker until the replica boots.
+                self._suspend_bitlocker(encrypted_volume_id)
+                # Add a resume command to the first-boot script.
+                script_content += (
+                    f'Resume-BitLocker "{encrypted_volume_id}"\r\n')
+
+            # Resume BitLocker after bringing the disks online,
+            # which has a script priority of 10.
+            os_morphing_tools.register_firstboot_script(
+                script_content,
+                user_provided=False,
+                script_filename="11-bitlocker-firstboot.ps1")
+        except Exception:
+            LOG.exception("First-boot preparation failed, attempting to "
+                          "resume BitLocker.")
+            for encrypted_volume_id in self._unlocked_volumes:
+                try:
+                    self._resume_bitlocker(encrypted_volume_id)
+                except Exception:
+                    LOG.exception(
+                        "Unable to resume BitLocker for volume: %s" %
+                        encrypted_volume_id)
+            raise
 
     def mount_os(self):
         self._set_basic_disks_rw_mode()

coriolis/tests/osmorphing/osmount/test_windows.py

@@ -381,7 +381,7 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
         exp_cmd = 'manage-bde -unlock "%s" -RecoveryPassword "%s"' % (
             vol, '***')
 
-        self.assertEqual(exp_cmd, strutils.mask_password(exp_cmd))
+        self.assertEqual(exp_cmd, strutils.mask_password(cmd))
 
     def test_suspend_bitlocker(self):
         vol = "\\\\?\\Volume{2750d574-b333-4e7b-a0a2-d739279d39e9}\\"
@@ -392,6 +392,15 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
         self.tools._conn.exec_ps_command.assert_called_once_with(
             exp_cmd)
 
+    def test_resume_bitlocker(self):
+        vol = "\\\\?\\Volume{2750d574-b333-4e7b-a0a2-d739279d39e9}\\"
+
+        self.tools._resume_bitlocker(vol)
+
+        exp_cmd = 'Resume-BitLocker "%s"' % (vol)
+        self.tools._conn.exec_ps_command.assert_called_once_with(
+            exp_cmd)
+
     @mock.patch.object(windows.WindowsMountTools, "_get_encrypted_volume_ids")
     def test_unlock_encrypted_volumes_no_password(
         self,
@@ -417,10 +426,8 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
 
     @mock.patch.object(windows.WindowsMountTools, "_get_encrypted_volume_ids")
     @mock.patch.object(windows.WindowsMountTools, "_unlock_encrypted_volume")
-    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
     def test_unlock_encrypted_volumes_all_failed(
         self,
-        mock_suspend_bitlocker,
         mock_unlock_encrypted_volume,
         mock_get_encrypted_volume_ids,
     ):
@@ -440,10 +447,8 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
 
     @mock.patch.object(windows.WindowsMountTools, "_get_encrypted_volume_ids")
     @mock.patch.object(windows.WindowsMountTools, "_unlock_encrypted_volume")
-    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
     def test_unlock_encrypted_volumes_one_failed(
         self,
-        mock_suspend_bitlocker,
         mock_unlock_encrypted_volume,
         mock_get_encrypted_volume_ids,
     ):
@@ -461,43 +466,10 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
 
         mock_unlock_encrypted_volume.assert_has_calls(
             [mock.call(vol_id, fake_pass) for vol_id in encrypted_volume_ids])
-        mock_suspend_bitlocker.assert_called_once_with(
-            mock.sentinel.volume1)
 
         self.assertEqual(
             self.tools._unlocked_volumes, [mock.sentinel.volume1])
 
-    @mock.patch.object(windows.WindowsMountTools, "_get_encrypted_volume_ids")
-    @mock.patch.object(windows.WindowsMountTools, "_unlock_encrypted_volume")
-    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
-    def test_unlock_encrypted_volumes_suspend_failed(
-        self,
-        mock_suspend_bitlocker,
-        mock_unlock_encrypted_volume,
-        mock_get_encrypted_volume_ids,
-    ):
-        fake_pass = "fake-recovery-password"
-        self.tools._osmorphing_info[constants.ENCRYPTED_DISKS_PASS] = fake_pass
-
-        encrypted_volume_ids = [
-            mock.sentinel.volume0,
-            mock.sentinel.volume1,
-            mock.sentinel.volume2
-        ]
-        mock_get_encrypted_volume_ids.return_value = encrypted_volume_ids
-        mock_unlock_encrypted_volume.side_effect = [ValueError, None, None]
-        mock_suspend_bitlocker.side_effect = [IOError, None]
-
-        # It should error out immediately when failing to suspend BitLocker.
-        self.assertRaises(
-            IOError,
-            self.tools._unlock_encrypted_volumes,
-        )
-        mock_unlock_encrypted_volume.assert_has_calls(
-            [mock.call(mock.sentinel.volume0, fake_pass),
-             mock.call(mock.sentinel.volume1, fake_pass)]
-        )
-
     def test_install_encryption_firstboot_setup_noop(self):
         # No unlocked volumes, nothing to do.
         mock_morphing_tools = mock.Mock()
@@ -506,7 +478,8 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
             mock_morphing_tools)
         mock_morphing_tools.register_firstboot_script.assert_not_called()
 
-    def test_install_encryption_firstboot_setup(self):
+    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
+    def test_install_encryption_firstboot_setup(self, mock_suspend_bitlocker):
         self.tools._unlocked_volumes = ["vol1", "vol2"]
         mock_morphing_tools = mock.Mock()
         self.tools.install_encryption_firstboot_setup(
@@ -519,3 +492,50 @@ class WindowsMountToolsTestCase(test_base.CoriolisBaseTestCase):
             expected_script,
             user_provided=False,
             script_filename="11-bitlocker-firstboot.ps1")
+        mock_suspend_bitlocker.assert_has_calls(
+            [mock.call(volume) for volume in self.tools._unlocked_volumes])
+
+    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
+    @mock.patch.object(windows.WindowsMountTools, "_resume_bitlocker")
+    def test_install_encryption_firstboot_setup_register_failure(
+        self,
+        mock_resume_bitlocker,
+        mock_suspend_bitlocker
+    ):
+        self.tools._unlocked_volumes = ["vol1", "vol2"]
+        mock_morphing_tools = mock.Mock()
+        mock_morphing_tools.register_firstboot_script.side_effect = IOError
+        self.assertRaises(
+            IOError,
+            self.tools.install_encryption_firstboot_setup,
+            mock.sentinel.os_root_dir,
+            mock_morphing_tools)
+
+        mock_suspend_bitlocker.assert_has_calls(
+            [mock.call(volume) for volume in self.tools._unlocked_volumes])
+        mock_resume_bitlocker.assert_has_calls(
+            [mock.call(volume) for volume in self.tools._unlocked_volumes])
+
+    @mock.patch.object(windows.WindowsMountTools, "_suspend_bitlocker")
+    @mock.patch.object(windows.WindowsMountTools, "_resume_bitlocker")
+    def test_install_encryption_firstboot_setup_suspend_failure(
+        self,
+        mock_resume_bitlocker,
+        mock_suspend_bitlocker
+    ):
+        self.tools._unlocked_volumes = ["vol1", "vol2"]
+        mock_morphing_tools = mock.Mock()
+        mock_suspend_bitlocker.side_effect = [None, IOError]
+        # Let resume-bitlocker fail to ensure that we're still trying to
+        # cover all volumes.
+        mock_resume_bitlocker.side_effect = ValueError
+        self.assertRaises(
+            IOError,
+            self.tools.install_encryption_firstboot_setup,
+            mock.sentinel.os_root_dir,
+            mock_morphing_tools)
+
+        mock_suspend_bitlocker.assert_has_calls(
+            [mock.call(volume) for volume in self.tools._unlocked_volumes])
+        mock_resume_bitlocker.assert_has_calls(
+            [mock.call(volume) for volume in self.tools._unlocked_volumes])