fix: jwe deserialization issues

Signed-off-by: Daniel Bluhm <[email protected]>

Author: dbluhm

didcomm_messaging/crypto/__init__.py

@@ -6,7 +6,7 @@
 
 from pydid import VerificationMethod
 
-from didcomm_messaging.jwe import JweEnvelope, from_b64url
+from didcomm_messaging.jwe import JweEnvelope
 
 
 class CryptoServiceError(Exception):
@@ -34,6 +34,7 @@ def multikey(self) -> str:
 
 class SecretKey(ABC):
     """Secret Key Type."""
+
     @property
     @abstractmethod
     def kid(self) -> str:

didcomm_messaging/crypto/askar/__init__.py

@@ -1,7 +1,7 @@
 """Askar backend for DIDComm Messaging."""
 from collections import OrderedDict
 import json
-from typing import Mapping, Optional, Sequence, Union
+from typing import Optional, Sequence, Union
 
 from pydid import VerificationMethod
 from didcomm_messaging.crypto import SecretsManager
@@ -406,8 +406,8 @@ async def ecdh_1pu_decrypt(
                 recip.encrypted_key,
                 cc_tag=wrapper.tag,
             )
-        except AskarError:
-            raise CryptoServiceError("Error decrypting content encryption key")
+        except AskarError as err:
+            raise CryptoServiceError("Error decrypting content encryption key") from err
 
         try:
             plaintext = cek.aead_decrypt(

didcomm_messaging/crypto/basic.py

@@ -0,0 +1,20 @@
+"""Basic Crypto Implementations."""
+
+from typing import Optional
+from . import S, SecretsManager
+
+
+class InMemorySecretsManager(SecretsManager[S]):
+    """In Memory Secrets Manager."""
+
+    def __init__(self, secrets: Optional[dict] = None):
+        """Initialize the InMemorySecretsManager."""
+        self.secrets = secrets or {}
+
+    async def get_secret_by_kid(self, kid: str) -> Optional[S]:
+        """Get a secret by its kid."""
+        return self.secrets.get(kid)
+
+    async def add_secret(self, secret: S) -> None:
+        """Add a secret to the secrets manager."""
+        self.secrets[secret.kid] = secret

didcomm_messaging/didcomm.py

@@ -1,20 +1,22 @@
 """Class DIDComm Messaging interface."""
 
 
-from typing import Generic, Literal, NamedTuple, Optional, Sequence, Union
+from dataclasses import dataclass
+from typing import Generic, Literal, Optional, Sequence, Union
 
 from pydid import DIDUrl, VerificationMethod
-from didcomm_messaging.crypto import P, S, CryptoService, SecretKey, SecretsManager
+from didcomm_messaging.crypto import P, S, CryptoService, SecretsManager
 from didcomm_messaging.jwe import JweEnvelope, from_b64url
 from didcomm_messaging.resolver import Resolver
 
 
-class PackedMessageMetadata(NamedTuple):
+@dataclass
+class PackedMessageMetadata(Generic[S]):
     """Unpack result."""
 
     wrapper: JweEnvelope
     method: Literal["ECDH-ES", "ECDH-1PU"]
-    recip_key: SecretKey
+    recip_key: S
     sender_kid: Optional[str]
 
 
@@ -36,7 +38,7 @@ def __init__(
         self.crypto = crypto
         self.secrets = secrets
 
-    async def extract_packed_message_metadata(
+    async def extract_packed_message_metadata(  # noqa: C901
         self, enc_message: Union[str, bytes]
     ) -> PackedMessageMetadata:
         """Extract metadata from a packed DIDComm message."""
@@ -101,9 +103,7 @@ async def unpack(self, enc_message: Union[str, bytes]) -> bytes:
             enc_message, metadata.recip_key, sender_key
         )
 
-    async def recip_for_kid_or_default_for_did(
-        self, kid_or_did: str
-    ) -> P:
+    async def recip_for_kid_or_default_for_did(self, kid_or_did: str) -> P:
         """Resolve a verification method for a kid or return default recip."""
         if "#" in kid_or_did:
             vm = await self.resolver.resolve_and_dereference_verification_method(
@@ -128,6 +128,31 @@ async def recip_for_kid_or_default_for_did(
 
         return self.crypto.verification_method_to_public_key(vm)
 
+    async def default_sender_kid_for_did(self, did: str) -> str:
+        """Determine the kid of the default sender key for a DID."""
+        if "#" in did:
+            return did
+
+        doc = await self.resolver.resolve_and_parse(did)
+        if not doc.key_agreement:
+            raise DIDCommMessagingError(
+                "No key agreement methods found; cannot determine recipient"
+            )
+
+        default = doc.key_agreement[0]
+        if isinstance(default, DIDUrl):
+            vm = doc.dereference(default)
+            if not isinstance(vm, VerificationMethod):
+                raise DIDCommMessagingError(
+                    f"Expected verification method, found: {type(vm)}"
+                )
+        else:
+            vm = default
+
+        if not vm.id.did:
+            return vm.id.as_absolute(vm.controller)
+        return vm.id
+
     async def pack(
         self,
         message: bytes,
@@ -136,33 +161,15 @@ async def pack(
         **options,
     ):
         """Pack a DIDComm message."""
-        recip_keys = [
-            await self.recip_for_kid_or_default_for_did(kid) for kid in to
-        ]
-        sender_key = await self.secrets.get_secret_by_kid(frm) if frm else None
+        recip_keys = [await self.recip_for_kid_or_default_for_did(kid) for kid in to]
+        sender_kid = await self.default_sender_kid_for_did(frm) if frm else None
+        sender_key = (
+            await self.secrets.get_secret_by_kid(sender_kid) if sender_kid else None
+        )
+        if frm and not sender_key:
+            raise DIDCommMessagingError("No sender key found")
 
         if sender_key:
-            return await self.crypto.ecdh_1pu_encrypt(
-                recip_keys, sender_key, message
-            )
+            return await self.crypto.ecdh_1pu_encrypt(recip_keys, sender_key, message)
         else:
             return await self.crypto.ecdh_es_encrypt(recip_keys, message)
-
-
-async def main():
-    from aries_askar import Store
-    from didcomm_messaging.askar import AskarCryptoService, AskarSecretsManager
-    from didcomm_messaging.resolver.peer import Peer2, Peer4
-    from didcomm_messaging.resolver import PrefixResolver
-
-    store = await Store.open("sqlite:///:memory:")
-    kms = AskarSecretsManager(store)
-    crypto = AskarCryptoService()
-    didcomm = DIDCommMessaging(
-        PrefixResolver(
-            {
-                "did:peer:2": Peer2(),
-                "did:peer:4": Peer4()
-            }
-        ), crypto, kms
-    )

didcomm_messaging/jwe.py

@@ -46,11 +46,14 @@ def deserialize(cls, entry: Mapping[str, Any]) -> "JweRecipient":
         if "encrypted_key" not in entry:
             raise ValueError("Invalid JWE recipient: missing encrypted_key")
 
+        encrypyted_key = from_b64url(entry["encrypted_key"])
+
         if "header" in entry:
             if not isinstance(entry["header"], dict):
                 raise ValueError("Invalid JWE recipient: invalid header")
-            return cls(encrypted_key=entry["encrypted_key"], header=entry["header"])
-        return cls(**entry)
+            return cls(encrypted_key=encrypyted_key, header=entry["header"])
+
+        return cls(encrypted_key=encrypyted_key)
 
     def serialize(self) -> dict:
         """Serialize the JWE recipient to a mapping."""
@@ -75,6 +78,7 @@ def __init__(
         self._with_flatten_recipients = with_flatten_recipients
         self._recipients: List[JweRecipient] = []
         self._protected: Optional[OrderedDict] = None
+        self._unprotected: Optional[OrderedDict] = None
         self._protected_b64: Optional[bytes] = None
         self._ciphertext: Optional[bytes] = None
         self._iv: Optional[bytes] = None
@@ -137,9 +141,6 @@ def build(self):
         if not self._tag:
             raise ValueError("Missing tag for JWE")
 
-        if not self._aad:
-            raise ValueError("Missing additional authenticated data for JWE")
-
         if not self._recipients:
             raise ValueError("Missing recipients for JWE")
 
@@ -158,6 +159,8 @@ def build(self):
             aad=self._aad,
             unprotected=self._unprotected,
             recipients=self._recipients,
+            with_flatten_recipients=self._with_flatten_recipients,
+            with_protected_recipients=self._with_protected_recipients,
         )
 
 
@@ -172,7 +175,7 @@ class JweEnvelope:
     iv: bytes
     tag: bytes
     aad: Optional[bytes] = None
-    unprotected: dict = field(default_factory=OrderedDict)
+    unprotected: Optional[dict] = field(default_factory=OrderedDict)
     with_protected_recipients: bool = False
     with_flatten_recipients: bool = True
 
@@ -185,7 +188,7 @@ def from_json(cls, message: Union[bytes, str]) -> "JweEnvelope":
             raise ValueError("Invalid JWE: not JSON")
 
     @classmethod
-    def deserialize(cls, message: Mapping[str, Any]) -> "JweEnvelope":
+    def deserialize(cls, message: Mapping[str, Any]) -> "JweEnvelope":  # noqa: C901
         """Deserialize a JWE envelope from a mapping."""
         # Basic validation
 
@@ -248,7 +251,7 @@ def deserialize(cls, message: Mapping[str, Any]) -> "JweEnvelope":
         return cls._deserialize(message)
 
     @classmethod
-    def _deserialize(cls, parsed: Mapping[str, Any]) -> "JweEnvelope":
+    def _deserialize(cls, parsed: Mapping[str, Any]) -> "JweEnvelope":  # noqa: C901
         protected_b64 = parsed[IDENT_PROTECTED]
         try:
             protected: dict = json.loads(from_b64url(protected_b64))
@@ -281,7 +284,7 @@ def _deserialize(cls, parsed: Mapping[str, Any]) -> "JweEnvelope":
             encrypted_key = parsed[IDENT_ENC_KEY]
             header = parsed.get(IDENT_HEADER)
         else:
-            raise ValueError("Invalid JWE: missing encrypted key")
+            header = None
 
         if recipients:
             if encrypted_key:
@@ -303,22 +306,27 @@ def _deserialize(cls, parsed: Mapping[str, Any]) -> "JweEnvelope":
             if recip.header and recip.header.keys() & all_h:
                 raise ValueError("Invalid JWE: duplicate header")
 
+        ciphertext = from_b64url(parsed["ciphertext"])
+        iv = from_b64url(parsed["iv"])
+        tag = from_b64url(parsed["tag"])
+        aad = from_b64url(parsed["aad"]) if "aad" in parsed else None
+
         inst = cls(
             recipients=recipients,
             protected=protected,
             protected_b64=protected_b64,
             unprotected=unprotected,
-            ciphertext=parsed["ciphertext"],
-            iv=parsed["iv"],
-            tag=parsed["tag"],
-            aad=parsed.get("aad"),
+            ciphertext=ciphertext,
+            iv=iv,
+            tag=tag,
+            aad=aad,
             with_protected_recipients=protected_recipients,
             with_flatten_recipients=flat_recipients,
         )
 
         return inst
 
-    def serialize(self) -> dict:
+    def serialize(self) -> dict:  # noqa: C901
         """Serialize the JWE envelope to a mapping."""
         if self.protected_b64 is None:
             raise ValueError("Missing protected: use set_protected")
@@ -329,7 +337,7 @@ def serialize(self) -> dict:
         if self.tag is None:
             raise ValueError("Missing tag for JWE")
         env = OrderedDict()
-        env["protected"] = self.protected_b64
+        env["protected"] = self.protected_b64.decode("utf-8")
         if self.unprotected:
             env["unprotected"] = self.unprotected.copy()
         if not self.with_protected_recipients:

example.py

@@ -0,0 +1,51 @@
+"""Example of using DIDComm Messaging."""
+
+from aries_askar import Key, KeyAlg
+from didcomm_messaging.crypto.askar import AskarCryptoService, AskarSecretKey
+from didcomm_messaging.crypto.basic import InMemorySecretsManager
+from didcomm_messaging.didcomm import DIDCommMessaging
+from didcomm_messaging.multiformats import multibase
+from didcomm_messaging.multiformats import multicodec
+from didcomm_messaging.resolver.peer import Peer2, Peer4
+from didcomm_messaging.resolver import PrefixResolver
+from did_peer_2 import KeySpec, generate, json
+
+
+async def main():
+    """An example of using DIDComm Messaging."""
+    secrets = InMemorySecretsManager()
+    crypto = AskarCryptoService()
+    didcomm = DIDCommMessaging(
+        PrefixResolver({"did:peer:2": Peer2(), "did:peer:4": Peer4()}), crypto, secrets
+    )
+    verkey = Key.generate(KeyAlg.ED25519)
+    xkey = Key.generate(KeyAlg.X25519)
+    did = generate(
+        [
+            KeySpec.verification(
+                multibase.encode(
+                    multicodec.wrap("ed25519-pub", verkey.get_public_bytes()),
+                    "base58btc",
+                )
+            ),
+            KeySpec.key_agreement(
+                multibase.encode(
+                    multicodec.wrap("x25519-pub", xkey.get_public_bytes()), "base58btc"
+                )
+            ),
+        ],
+        [],
+    )
+    await secrets.add_secret(AskarSecretKey(verkey, f"{did}#key-1"))
+    await secrets.add_secret(AskarSecretKey(xkey, f"{did}#key-2"))
+    print(did)
+    packed = await didcomm.pack(b"hello world", [did], did)
+    print(json.dumps(json.loads(packed), indent=2))
+    unpacked = await didcomm.unpack(packed)
+    print(unpacked)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

pdm.lock

@@ -50,4 +50,5 @@ dev = [
     "pre-commit>=3.5.0",
     "black>=23.10.1",
     "ruff>=0.1.3",
+    "pytest-asyncio>=0.21.1",
 ]