Add support for HPKE key wrapping.

dajiaji · dajiaji · commit fec2e3ac0e0e · 2022-11-19T17:19:25.000+09:00
diff --git a/README.md b/README.md
@@ -623,6 +623,7 @@ Create a COSE-HPKE Encrypt message and decrypt it as follows:
 from cwt import COSE, COSEKey, Recipient
 
 # The sender side:
+enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
 rpk = COSEKey.from_jwk(
     {
         "kty": "EC",
@@ -645,10 +646,13 @@ r = Recipient.new(
         },
     },
 )
-r.apply(recipient_key=rpk)
+r.encode(enc_key.key, recipient_key=rpk)
 sender = COSE.new()
 encoded = sender.encode_and_encrypt(
     b"This is the content.",
+    protected={
+        1: 1,  # alg: "A128GCM"
+    },
     recipients=[r],
 )
 
diff --git a/cwt/cose.py b/cwt/cose.py
@@ -169,8 +169,8 @@ def encode_and_encrypt(
         if not recipients:
             if 1 in p and p[1] == -1:  # HPKE
                 hpke = HPKE(p, u)
-                hpke.apply(recipient_key=key)
-                res = CBORTag(16, hpke.to_list(payload, external_aad, "Encrypt0"))
+                hpke.encode(payload, recipient_key=key, external_aad=external_aad, aad_context="Encrypt0")
+                res = CBORTag(16, hpke.to_list())
                 return res if out == "cbor2/CBORTag" else self._dumps(res)
             if key is None:
                 raise ValueError("key should be set.")
@@ -412,7 +412,7 @@ def decode(
                     try:
                         if not isinstance(protected, bytes) and alg == -1:  # HPKE
                             hpke = HPKE(protected, unprotected, data.value[2])
-                            return hpke.decrypt(k, external_aad=external_aad, aad_context="Encrypt0")
+                            return hpke.decode(k, external_aad=external_aad, aad_context="Encrypt0")
                         return k.decrypt(data.value[2], nonce, aad)
                     except Exception as e:
                         err = e
diff --git a/cwt/recipient_algs/hpke.py b/cwt/recipient_algs/hpke.py
@@ -2,6 +2,7 @@
 
 from pyhpke import AEADId, CipherSuite, KDFId, KEMId, KEMKey, KEMKeyInterface
 
+from ..cose_key import COSEKey
 from ..cose_key_interface import COSEKeyInterface
 from ..exceptions import DecodeError, EncodeError
 from ..recipient_interface import RecipientInterface
@@ -30,62 +31,58 @@ def __init__(
         self._suite = CipherSuite.new(KEMId(unprotected[-4][1]), KDFId(unprotected[-4][2]), AEADId(unprotected[-4][3]))
         return
 
-    def apply(
+    def encode(
         self,
-        key: Optional[COSEKeyInterface] = None,
+        plaintext: bytes,
         recipient_key: Optional[COSEKeyInterface] = None,
         salt: Optional[bytes] = None,
         context: Optional[Union[List[Any], Dict[str, Any]]] = None,
         external_aad: bytes = b"",
         aad_context: str = "Enc_Recipient",
-    ) -> COSEKeyInterface:
-        # if not key:
-        #     raise ValueError("key should be set.")
+    ) -> Optional[COSEKeyInterface]:
         if not recipient_key:
             raise ValueError("recipient_key should be set.")
-        # if recipient_key.kid:
-        #     self._protected[4] = key.kid
         self._recipient_key = recipient_key
         self._kem_key = self._to_kem_key(recipient_key)
-        # enc_structure = ["Enc_Recipient", self._dumps(self._protected), external_aad]
-        # aad = self._dumps(enc_structure)
-        # enc, sender = self._suite.create_sender_context(self._kem_key)
-        # self._unprotected[-4][4] = enc
-        # try:
-        #     self._ciphertext = sender.seal(key.key, aad=aad)
-        # except Exception as err:
-        #     raise EncodeError("Failed to seal.") from err
-        return self._recipient_key
-
-    def to_list(self, payload: bytes = b"", external_aad: bytes = b"", aad_context: str = "Enc_Recipient") -> List[Any]:
         enc_structure = [aad_context, self._dumps(self._protected), external_aad]
         aad = self._dumps(enc_structure)
-        enc, sender = self._suite.create_sender_context(self._kem_key)
-        self._unprotected[-4][4] = enc
         try:
-            self._ciphertext = sender.seal(payload, aad=aad)
-            return super().to_list(payload, external_aad, aad_context)
+            enc, ctx = self._suite.create_sender_context(self._kem_key)
+            self._unprotected[-4][4] = enc
+            self._ciphertext = ctx.seal(plaintext, aad=aad)
         except Exception as err:
             raise EncodeError("Failed to seal.") from err
+        return None
 
-    def decrypt(
+    def decode(
         self,
         key: COSEKeyInterface,
-        alg: Optional[int] = None,
         context: Optional[Union[List[Any], Dict[str, Any]]] = None,
-        payload: bytes = b"",
-        nonce: bytes = b"",
-        aad: bytes = b"",
         external_aad: bytes = b"",
         aad_context: str = "Enc_Recipient",
     ) -> bytes:
         enc_structure = [aad_context, self._dumps(self._protected), external_aad]
         aad = self._dumps(enc_structure)
-        recipient = self._suite.create_recipient_context(self._unprotected[-4][4], self._to_kem_key(key))
         try:
-            return recipient.open(self._ciphertext, aad=aad)
+            ctx = self._suite.create_recipient_context(self._unprotected[-4][4], self._to_kem_key(key))
+            return ctx.open(self._ciphertext, aad=aad)
         except Exception as err:
             raise DecodeError("Failed to open.") from err
 
+    def decrypt(
+        self,
+        key: COSEKeyInterface,
+        alg: Optional[int] = None,
+        context: Optional[Union[List[Any], Dict[str, Any]]] = None,
+        payload: bytes = b"",
+        nonce: bytes = b"",
+        aad: bytes = b"",
+        external_aad: bytes = b"",
+        aad_context: str = "Enc_Recipient",
+    ) -> bytes:
+        alg = alg if isinstance(alg, int) else 0
+        raw = self.decode(key, context, external_aad, aad_context)
+        return COSEKey.from_symmetric_key(raw, alg=alg, kid=self._kid).decrypt(payload, nonce, aad)
+
     def _to_kem_key(self, src: COSEKeyInterface) -> KEMKeyInterface:
         return KEMKey.from_pyca_cryptography_key(src.key)
diff --git a/cwt/recipient_interface.py b/cwt/recipient_interface.py
@@ -154,6 +154,45 @@ def to_list(self, payload: bytes = b"", external_aad: bytes = b"", aad_context:
         res.append(children)
         return res
 
+    def encode(
+        self,
+        plaintext: bytes,
+        recipient_key: Optional[COSEKeyInterface] = None,
+        salt: Optional[bytes] = None,
+        context: Optional[Union[List[Any], Dict[str, Any]]] = None,
+        external_aad: bytes = b"",
+        aad_context: str = "Enc_Recipient",
+    ) -> Optional[COSEKeyInterface]:
+        """
+        Encodes a specified plaintext to a ciphertext with the recipient-specific
+        method (e.g., key wrapping, key agreement, or the combination of them)
+        and sets up the related information (context information or ciphertext)
+        in the recipient structure.
+        Therefore, it will be used by the sender of the recipient information
+        before calling COSE.encode_* functions with the Recipient object. The
+        key generated through this function will be set to ``key`` parameter
+        of COSE.encode_* functions.
+
+        Args:
+            plaintext (bytes): A plaing text to be encrypted. In most of the cases,
+                the plaintext is a byte string of a content encryption key.
+            recipient_key (Optional[COSEKeyInterface]): The external public
+                key provided by the recipient used for ECDH key agreement, HPKE, etc.
+            salt (Optional[bytes]): A salt used for deriving a key.
+            context (Optional[Union[List[Any], Dict[str, Any]]]): Context
+                information structure.
+            external_aad (bytes): External additional authenticated data for AEAD.
+            aad_context (bytes): An additional authenticated data context to build
+                an Enc_structure internally.
+        Returns:
+            Optional[COSEKeyInterface]: A generated key or passed-through key
+                which is used as ``key`` parameter of COSE.encode_* functions.
+        Raises:
+            ValueError: Invalid arguments.
+            EncodeError: Failed to encode(e.g., wrap, derive) the key.
+        """
+        raise NotImplementedError
+
     def apply(
         self,
         key: Optional[COSEKeyInterface] = None,
@@ -164,7 +203,7 @@ def apply(
         aad_context: str = "Enc_Recipient",
     ) -> COSEKeyInterface:
         """
-        Applies a COSEKey as a material to prepare a MAC/encryption key with
+        [DEPRECATED] Applies a COSEKey as a material to prepare a MAC/encryption key with
         the recipient-specific method (e.g., key wrapping, key agreement,
         or the combination of them) and sets up the related information
         (context information or ciphertext) in the recipient structure.
diff --git a/tests/test_cose_sample.py b/tests/test_cose_sample.py
@@ -421,6 +421,7 @@ def test_cose_usage_examples_cose_encrypt(self):
     def test_cose_usage_examples_cose_encrypt_hpke(self):
 
         # The sender side:
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         rpk = COSEKey.from_jwk(
             {
                 "kty": "EC",
@@ -443,10 +444,11 @@ def test_cose_usage_examples_cose_encrypt_hpke(self):
                 },
             },
         )
-        r.apply(recipient_key=rpk)
-        sender = COSE.new()
+        r.encode(enc_key.key, recipient_key=rpk)
+        sender = COSE.new(alg_auto_inclusion=True)
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
+            key=enc_key,
             recipients=[r],
         )
 
@@ -469,6 +471,7 @@ def test_cose_usage_examples_cose_encrypt_hpke(self):
     def test_cose_usage_examples_cose_encrypt_hpke_with_1st_layer_hpke(self):
 
         # The sender side:
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         rpk = COSEKey.from_jwk(
             {
                 "kty": "EC",
@@ -491,7 +494,7 @@ def test_cose_usage_examples_cose_encrypt_hpke_with_1st_layer_hpke(self):
                 },
             },
         )
-        r.apply(recipient_key=rpk)
+        r.encode(enc_key.key, recipient_key=rpk)
         sender = COSE.new()
         with pytest.raises(ValueError) as err:
             sender.encode_and_encrypt(
@@ -515,6 +518,7 @@ def test_cose_usage_examples_cose_encrypt_hpke_with_1st_layer_hpke(self):
     def test_cose_usage_examples_cose_encrypt_hpke_with_nonce(self):
 
         # The sender side:
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         rpk = COSEKey.from_jwk(
             {
                 "kty": "EC",
@@ -537,7 +541,7 @@ def test_cose_usage_examples_cose_encrypt_hpke_with_nonce(self):
                 },
             },
         )
-        r.apply(recipient_key=rpk)
+        r.encode(enc_key.key, recipient_key=rpk)
         sender = COSE.new()
         with pytest.raises(ValueError) as err:
             sender.encode_and_encrypt(
diff --git a/tests/test_recipient.py b/tests/test_recipient.py
@@ -649,6 +649,7 @@ def test_recipients_open_with_empty_recipients(self, rsk1):
         assert "No recipients." in str(err.value)
 
     def test_recipients_open_with_rpk_without_kid(self, rsk1, rsk2):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         r = Recipient.new(protected={1: -1}, unprotected={-4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
         rpk = COSEKey.from_jwk(
             {
@@ -659,19 +660,21 @@ def test_recipients_open_with_rpk_without_kid(self, rsk1, rsk2):
                 "y": "BGU5soLgsu_y7GN2I3EPUXS9EZ7Sw0qif-V70JtInFI",
             }
         )
-        r.apply(recipient_key=rpk)
+        r.encode(enc_key.key, recipient_key=rpk)
         sender = COSE.new()
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
-            # protected={
-            #     1: -1,  # alg: "HPKE"
-            # },
+            enc_key,
+            protected={
+                1: 1,  # alg: "A128GCM"
+            },
             recipients=[r],
         )
         recipient = COSE.new()
         assert b"This is the content." == recipient.decode(encoded, [rsk1, rsk2])
 
     def test_recipients_open_with_verify_kid_and_rpk_without_kid(self, rsk1, rsk2):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         r = Recipient.new(protected={1: -1}, unprotected={-4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
         rpk = COSEKey.from_jwk(
             {
@@ -682,7 +685,7 @@ def test_recipients_open_with_verify_kid_and_rpk_without_kid(self, rsk1, rsk2):
                 "y": "BGU5soLgsu_y7GN2I3EPUXS9EZ7Sw0qif-V70JtInFI",
             }
         )
-        r.apply(recipient_key=rpk)
+        r.encode(enc_key.key, recipient_key=rpk)
         sender = COSE.new()
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
@@ -698,6 +701,7 @@ def test_recipients_open_with_verify_kid_and_rpk_without_kid(self, rsk1, rsk2):
         assert "kid should be specified in recipient." in str(err.value)
 
     def test_recipients_open_failed_with_rpk_without_kid(self, rsk1):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         r = Recipient.new(protected={1: -1}, unprotected={-4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
         rpk = COSEKey.from_jwk(
             {
@@ -708,7 +712,7 @@ def test_recipients_open_failed_with_rpk_without_kid(self, rsk1):
                 "y": "BGU5soLgsu_y7GN2I3EPUXS9EZ7Sw0qif-V70JtInFI",
             }
         )
-        r.apply(recipient_key=rpk)
+        r.encode(enc_key.key, recipient_key=rpk)
         sender = COSE.new()
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
@@ -724,11 +728,13 @@ def test_recipients_open_failed_with_rpk_without_kid(self, rsk1):
         assert "Failed to open." in str(err.value)
 
     def test_recipients_open_with_multiple_rsks(self, rpk2, rsk1, rsk2):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         r = Recipient.new(protected={1: -1}, unprotected={4: b"02", -4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
-        r.apply(recipient_key=rpk2)
-        sender = COSE.new()
+        r.encode(enc_key.key, recipient_key=rpk2)
+        sender = COSE.new(alg_auto_inclusion=True)
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
+            key=enc_key,
             # protected={
             #     1: -1,  # alg: "HPKE"
             # },
@@ -738,8 +744,9 @@ def test_recipients_open_with_multiple_rsks(self, rpk2, rsk1, rsk2):
         assert b"This is the content." == recipient.decode(encoded, [rsk1, rsk2])
 
     def test_recipients_open_with_invalid_rsk(self, rpk1):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         r = Recipient.new(protected={1: -1}, unprotected={4: b"02", -4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
-        r.apply(recipient_key=rpk1)
+        r.encode(enc_key.key, recipient_key=rpk1)
         sender = COSE.new()
         encoded = sender.encode_and_encrypt(
             b"This is the content.",
diff --git a/tests/test_recipient_algs_hpke.py b/tests/test_recipient_algs_hpke.py
@@ -4,6 +4,7 @@
 
 import pytest
 
+from cwt import COSEKey
 from cwt.recipient_algs.hpke import HPKE
 
 
@@ -18,9 +19,10 @@ def test_recipient_algs_hpke(self):
         assert ctx.alg == -1
 
     def test_recipient_algs_hpke_apply_without_recipient_key(self):
+        enc_key = COSEKey.from_symmetric_key(alg="A128GCM")
         ctx = HPKE({1: -1}, {-4: {1: 0x0010, 2: 0x0001, 3: 0x0001}})
         with pytest.raises(ValueError) as err:
-            ctx.apply()
+            ctx.encode(enc_key.key)
             pytest.fail("apply should fail.")
         assert "recipient_key should be set." in str(err.value)
 
