fix(jwe): auto add kid to recipient when kid exists

Author: lepture

src/joserfc/_rfc7516/json.py

@@ -36,10 +36,10 @@ def represent_general_json(obj: GeneralJSONEncryption) -> GeneralJSONSerializati
     recipients = []
     for recipient in obj.recipients:
         item: JSONRecipientDict = {}
-        if recipient.header:
-            item["header"] = recipient.header
-        if recipient.encrypted_key:
-            item["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
+        assert recipient.header is not None
+        assert recipient.encrypted_key is not None
+        item["header"] = recipient.header
+        item["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
         recipients.append(item)
     data["recipients"] = recipients
     return data
@@ -49,10 +49,10 @@ def represent_flattened_json(obj: FlattenedJSONEncryption) -> FlattenedJSONSeria
     data: FlattenedJSONSerialization = __represent_json_serialization(obj)
     recipient = obj.recipients[0]
     assert recipient is not None
-    if recipient.header:
-        data["header"] = recipient.header
-    if recipient.encrypted_key:
-        data["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
+    assert recipient.header is not None
+    assert recipient.encrypted_key is not None
+    data["header"] = recipient.header
+    data["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
     return data
 
 

src/joserfc/_rfc7516/models.py

@@ -156,9 +156,8 @@ class GeneralJSONEncryption(BaseJSONEncryption):
 
     def add_recipient(self, header: Header | None = None, key: Key | None = None) -> None:
         recipient = Recipient(self, header, key)
-        if key:
-            key.ensure_kid()
-            recipient.set_kid(t.cast(str, key.kid))
+        if key and key.kid:
+            recipient.set_kid(key.kid)
         self.recipients.append(recipient)
 
 
@@ -180,7 +179,10 @@ class FlattenedJSONEncryption(BaseJSONEncryption):
     flattened = True
 
     def add_recipient(self, header: Header | None = None, key: Key | None = None) -> None:
-        self.recipients = [Recipient(self, header, key)]
+        recipient = Recipient(self, header, key)
+        if key and key.kid:
+            recipient.set_kid(key.kid)
+        self.recipients = [recipient]
 
 
 class JWEEncModel(metaclass=ABCMeta):

tests/jwe/test_json.py

@@ -1,6 +1,6 @@
 from unittest import TestCase
 from joserfc import jwe
-from joserfc.jwe import GeneralJSONEncryption
+from joserfc.jwe import GeneralJSONEncryption, FlattenedJSONEncryption
 from joserfc.jwk import KeySet, RSAKey, ECKey, OctKey
 from joserfc.errors import (
     DecodeError,
@@ -72,3 +72,19 @@ def test_decode_multiple_recipients(self):
             key3,
             registry=registry,
         )
+
+    def test_flattened_encryption(self):
+        key = OctKey.generate_key(128)
+        protected = {"enc": "A128CBC-HS256"}
+        plaintext = b"hello world"
+        obj0 = FlattenedJSONEncryption(protected, plaintext)
+        obj0.add_recipient({"alg": "A128KW"})
+        value = jwe.encrypt_json(obj0, key)
+        obj1 = jwe.decrypt_json(value, key)
+        self.assertEqual(obj1.plaintext, plaintext)
+
+        obj2 = FlattenedJSONEncryption(protected, plaintext)
+        obj2.add_recipient({"alg": "A128KW"}, key)
+        value = jwe.encrypt_json(obj0, None)
+        obj3 = jwe.decrypt_json(value, key)
+        self.assertEqual(obj3.plaintext, plaintext)