Add rekey feature with blank key support (openwallet-foundation#3125)

* Add rekey feature with blank key support Signed-off-by: jamshale <jamiehalebc@gmail.com> * Refactor / Remove other key exists check / unit tests Signed-off-by: jamshale <jamiehalebc@gmail.com> * Remove no wallet key test from store config Signed-off-by: jamshale <jamiehalebc@gmail.com> --------- Signed-off-by: jamshale <jamiehalebc@gmail.com> Co-authored-by: Ian Costanzo <ian@anon-solutions.ca>
Northern-Block · Sep 5, 2024 · 4ed99bc · 4ed99bc
1 parent 3e125a3
commit 4ed99bc
Show file tree

Hide file tree

Showing 3 changed files with 132 additions and 27 deletions.
diff --git a/aries_cloudagent/askar/store.py b/aries_cloudagent/askar/store.py
@@ -6,7 +6,7 @@
 
 from aries_askar import AskarError, AskarErrorCode, Store
 
-from ..core.error import ProfileError, ProfileDuplicateError, ProfileNotFoundError
+from ..core.error import ProfileDuplicateError, ProfileError, ProfileNotFoundError
 from ..core.profile import Profile
 from ..utils.env import storage_path
 
@@ -36,21 +36,16 @@ def __init__(self, config: dict = None):
             config = {}
         self.auto_recreate = config.get("auto_recreate", False)
         self.auto_remove = config.get("auto_remove", False)
+
         self.key = config.get("key", self.DEFAULT_KEY)
         self.key_derivation_method = (
             config.get("key_derivation_method") or self.DEFAULT_KEY_DERIVATION
         )
 
-        if (
-            self.key_derivation_method.lower() == self.KEY_DERIVATION_RAW.lower()
-            and self.key == ""
-        ):
-            raise ProfileError(
-                f"With key derivation method '{self.KEY_DERIVATION_RAW}',"
-                "key should also be provided"
-            )
-        # self.rekey = config.get("rekey")
-        # self.rekey_derivation_method = config.get("rekey_derivation_method")
+        self.rekey = config.get("rekey")
+        self.rekey_derivation_method = (
+            config.get("rekey_derivation_method") or self.DEFAULT_KEY_DERIVATION
+        )
 
         self.name = config.get("name") or Profile.DEFAULT_NAME
         self.in_memory = self.name == ":memory:"
@@ -133,6 +128,20 @@ async def remove_store(self):
                 )
             raise ProfileError("Error removing store") from err
 
+    def _handle_open_error(self, err: AskarError, retry=False):
+        if err.code == AskarErrorCode.DUPLICATE:
+            raise ProfileDuplicateError(
+                f"Duplicate store '{self.name}'",
+            )
+        if err.code == AskarErrorCode.NOT_FOUND:
+            raise ProfileNotFoundError(
+                f"Store '{self.name}' not found",
+            )
+        if retry and self.rekey:
+            return
+
+        raise ProfileError("Error opening store") from err
+
     async def open_store(self, provision: bool = False) -> "AskarOpenStore":
         """Open a store, removing and/or creating it if so configured.
 
@@ -156,16 +165,27 @@ async def open_store(self, provision: bool = False) -> "AskarOpenStore":
                     self.key_derivation_method,
                     self.key,
                 )
+                if self.rekey:
+                    await Store.rekey(store, self.rekey_derivation_method, self.rekey)
+
         except AskarError as err:
-            if err.code == AskarErrorCode.DUPLICATE:
-                raise ProfileDuplicateError(
-                    f"Duplicate store '{self.name}'",
-                )
-            if err.code == AskarErrorCode.NOT_FOUND:
-                raise ProfileNotFoundError(
-                    f"Store '{self.name}' not found",
-                )
-            raise ProfileError("Error opening store") from err
+            self._handle_open_error(err, retry=True)
+
+            if self.rekey:
+                # Attempt to rekey the store with a default key in the case the key
+                # was created with a blank key before version 0.12.0. This can be removed
+                # in a future version or when 0.11.0 is no longer supported.
+                try:
+                    store = await Store.open(
+                        self.get_uri(),
+                        self.key_derivation_method,
+                        AskarStoreConfig.DEFAULT_KEY,
+                    )
+                except AskarError as err:
+                    self._handle_open_error(err)
+
+                await Store.rekey(store, self.rekey_derivation_method, self.rekey)
+                return AskarOpenStore(self, provision, store)
 
         return AskarOpenStore(self, provision, store)
 

diff --git a/aries_cloudagent/askar/tests/test_store.py b/aries_cloudagent/askar/tests/test_store.py
@@ -1,8 +1,11 @@
 from unittest import IsolatedAsyncioTestCase
 
-from ...core.error import ProfileError
+from aries_askar import AskarError, AskarErrorCode, Store
 
-from ..store import AskarStoreConfig
+from aries_cloudagent.tests import mock
+
+from ...core.error import ProfileDuplicateError, ProfileError, ProfileNotFoundError
+from ..store import AskarOpenStore, AskarStoreConfig
 
 
 class TestStoreConfig(IsolatedAsyncioTestCase):
@@ -23,11 +26,83 @@ async def test_init_success(self):
         assert askar_store.key == self.key
         assert askar_store.storage_type == self.storage_type
 
-    async def test_init_should_fail_when_key_missing(self):
+
+class TestStoreOpen(IsolatedAsyncioTestCase):
+    key_derivation_method = "Raw"
+    key = "key"
+    storage_type = "default"
+
+    @mock.patch.object(Store, "open", autospec=True)
+    async def test_open_store(self, mock_store_open):
+        config = {
+            "key_derivation_method": self.key_derivation_method,
+            "key": self.key,
+            "storage_type": self.storage_type,
+        }
+
+        store = await AskarStoreConfig(config).open_store()
+        assert isinstance(store, AskarOpenStore)
+        assert mock_store_open.called
+
+    @mock.patch.object(Store, "open")
+    async def test_open_store_fails(self, mock_store_open):
         config = {
             "key_derivation_method": self.key_derivation_method,
+            "key": self.key,
             "storage_type": self.storage_type,
         }
 
+        mock_store_open.side_effect = [
+            AskarError(AskarErrorCode.NOT_FOUND, message="testing"),
+            AskarError(AskarErrorCode.DUPLICATE, message="testing"),
+            AskarError(AskarErrorCode.ENCRYPTION, message="testing"),
+        ]
+
+        with self.assertRaises(ProfileNotFoundError):
+            await AskarStoreConfig(config).open_store()
+        with self.assertRaises(ProfileDuplicateError):
+            await AskarStoreConfig(config).open_store()
         with self.assertRaises(ProfileError):
-            askar_store = AskarStoreConfig(config)
+            await AskarStoreConfig(config).open_store()
+
+    @mock.patch.object(Store, "open")
+    @mock.patch.object(Store, "rekey")
+    async def test_open_store_fail_retry_with_rekey(self, mock_store_open, mock_rekey):
+        config = {
+            "key_derivation_method": self.key_derivation_method,
+            "key": self.key,
+            "storage_type": self.storage_type,
+            "rekey": "rekey",
+        }
+
+        mock_store_open.side_effect = [
+            AskarError(AskarErrorCode.ENCRYPTION, message="testing"),
+            mock.AsyncMock(auto_spec=True),
+        ]
+
+        store = await AskarStoreConfig(config).open_store()
+
+        assert isinstance(store, AskarOpenStore)
+        assert mock_rekey.called
+
+    @mock.patch.object(Store, "open")
+    @mock.patch.object(Store, "rekey")
+    async def test_open_store_fail_retry_with_rekey_fails(
+        self, mock_store_open, mock_rekey
+    ):
+        config = {
+            "key_derivation_method": self.key_derivation_method,
+            "key": self.key,
+            "storage_type": self.storage_type,
+            "rekey": "rekey",
+        }
+
+        mock_store_open.side_effect = [
+            AskarError(AskarErrorCode.ENCRYPTION, message="testing"),
+            mock.AsyncMock(auto_spec=True),
+        ]
+
+        store = await AskarStoreConfig(config).open_store()
+
+        assert isinstance(store, AskarOpenStore)
+        assert mock_rekey.called
diff --git a/aries_cloudagent/config/argparse.py b/aries_cloudagent/config/argparse.py
@@ -1632,10 +1632,16 @@ def add_arguments(self, parser: ArgumentParser):
             type=str,
             metavar="<key-derivation-method>",
             env_var="ACAPY_WALLET_KEY_DERIVATION_METHOD",
+            help=("Specifies the key derivation method used for wallet encryption."),
+        )
+        parser.add_argument(
+            "--wallet-rekey-derivation-method",
+            type=str,
+            metavar="<rekey-derivation-method>",
+            env_var="ACAPY_WALLET_REKEY_DERIVATION_METHOD",
             help=(
-                "Specifies the key derivation method used for wallet encryption."
-                "If RAW key derivation method is used, also --wallet-key parameter"
-                " is expected."
+                "Specifies the key derivation method used for the replacement"
+                "rekey encryption."
             ),
         )
         parser.add_argument(
@@ -1694,6 +1700,10 @@ def get_settings(self, args: Namespace) -> dict:
             settings["wallet.type"] = args.wallet_type
         if args.wallet_key_derivation_method:
             settings["wallet.key_derivation_method"] = args.wallet_key_derivation_method
+        if args.wallet_rekey_derivation_method:
+            settings["wallet.rekey_derivation_method"] = (
+                args.wallet_rekey_derivation_method
+            )
         if args.wallet_storage_config:
             settings["wallet.storage_config"] = args.wallet_storage_config
         if args.wallet_storage_creds: