From 252f14b792805125008c9b18f5f99b2a5da7e40b Mon Sep 17 00:00:00 2001
From: Vojtech Trefny <vtrefny@redhat.com>
Date: Fri, 9 Apr 2021 10:44:33 +0200
Subject: [PATCH] Allow unlocking locked Stratis Pool devices

---
 blivetgui/blivet_utils.py    | 37 +++++++++++++++++++++++++++++++++++-
 blivetgui/blivetgui.py       |  2 +-
 blivetgui/list_partitions.py |  3 ++-
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/blivetgui/blivet_utils.py b/blivetgui/blivet_utils.py
index 8652a2fd..f2939ec3 100644
--- a/blivetgui/blivet_utils.py
+++ b/blivetgui/blivet_utils.py
@@ -1492,7 +1492,18 @@ def create_disk_label(self, blivet_device, label_type):
 
         return ProxyDataContainer(success=True, actions=actions, message=None, exception=None, traceback=None)
 
-    def luks_decrypt(self, blivet_device, passphrase):
+    def unlock_device(self, blivet_device, passphrase):
+        """ Unlock/open this LUKS/dm-crypt encrypted device
+        """
+
+        if blivet_device.format.type == "luks":
+            return self._luks_unlock(blivet_device, passphrase)
+        elif blivet_device.format.type == "stratis":
+            return self._stratis_unlock(blivet_device, passphrase)
+        else:
+            return False
+
+    def _luks_unlock(self, blivet_device, passphrase):
         """ Decrypt selected luks device
 
             :param blivet_device: device to decrypt
@@ -1523,6 +1534,30 @@ def luks_decrypt(self, blivet_device, passphrase):
             self.storage.devicetree.populate()
             return True
 
+    def _stratis_unlock(self, blivet_device, passphrase):
+        """ Unlock Stratis pool on this device
+
+            :param blivet_device: stratis blockdev with a locked pool
+            :type blivet_device: StorageDevice
+            :param passphrase: passphrase
+            :type passphrase: str
+
+        """
+
+        log_msg = "Opening Stratis device '%s':\n" % blivet_device
+        log_utils_call(log=self.log, message=log_msg,
+                       user_input={"device": blivet_device})
+
+        blivet_device.format.passphrase = passphrase
+
+        try:
+            blivet_device.format.unlock_pool()
+        except blivet.errors.StratisError:
+            return False
+        else:
+            self.storage.devicetree.populate()
+            return True
+
     def blivet_cancel_actions(self, actions):
         """ Cancel scheduled actions
         """
diff --git a/blivetgui/blivetgui.py b/blivetgui/blivetgui.py
index 50175024..4d856e54 100644
--- a/blivetgui/blivetgui.py
+++ b/blivetgui/blivetgui.py
@@ -714,7 +714,7 @@ def decrypt_device(self, _widget=None):
         response = self.run_dialog(dialog)
 
         if response:
-            ret = self.client.remote_call("luks_decrypt", self.list_partitions.selected_partition[0], response)
+            ret = self.client.remote_call("unlock_device", self.list_partitions.selected_partition[0], response)
 
             if not ret:
                 msg = _("Unlocking failed. Are you sure provided password is correct?")
diff --git a/blivetgui/list_partitions.py b/blivetgui/list_partitions.py
index ee22440e..9c14f61b 100644
--- a/blivetgui/list_partitions.py
+++ b/blivetgui/list_partitions.py
@@ -326,7 +326,8 @@ def activate_action_buttons(self, selected_device):
         if device.format:
             if device.format.type == "luks" and not device.format.status and device.format.exists:
                 self.blivet_gui.activate_device_actions(["decrypt"])
-
+            elif device.format.type == "stratis" and device.format.locked_pool and not device.children:
+                self.blivet_gui.activate_device_actions(["decrypt"])
             elif device.format.mountable and device.format.system_mountpoint:
                 self.blivet_gui.activate_device_actions(["unmount"])