Update protected_header type from bytes to dict.

cwt/recipient.py

@@ -11,7 +11,7 @@ class Recipient(CBORProcessor):
 
     def __init__(
         self,
-        protected: Dict[int, Any] = {},
+        protected: Union[bytes, Dict[int, Any]] = {},
         unprotected: Dict[int, Any] = {},
         ciphertext: bytes = b"",
         recipients: List[Any] = [],
@@ -27,7 +27,12 @@ def __init__(
                     raise ValueError("ciphertext should be zero-length bytes.")
                 if len(recipients) != 0:
                     raise ValueError("recipients should be absent.")
-        self._protected = protected
+        if protected == b"":
+            self._protected = {}
+        elif isinstance(protected, bytes):
+            self._protected = self._loads(protected)
+        else:
+            self._protected = protected
         self._unprotected = unprotected
         self._ciphertext = ciphertext
 

tests/test_recipient.py

@@ -6,6 +6,7 @@
 """
 Tests for Recipient.
 """
+import cbor2
 import pytest
 
 from cwt import Recipient, cose_key
@@ -40,9 +41,13 @@ def test_recipient_constructor(self):
 
     def test_recipient_constructor_with_args(self):
         child = Recipient(unprotected={1: -6, 4: b"our-secret"})
-        r = Recipient(unprotected={1: -1, 4: b"our-secret"}, recipients=[child])
-        assert isinstance(r, Recipient)
-        assert r.protected == {}
+        r = Recipient(
+            protected={"foo": "bar"},
+            unprotected={1: -1, 4: b"our-secret"},
+            recipients=[child],
+        )
+        assert isinstance(r.protected, dict)
+        assert r.protected["foo"] == "bar"
         assert isinstance(r.unprotected, dict)
         assert r.kid == b"our-secret"
         assert r.alg == -1
@@ -55,6 +60,11 @@ def test_recipient_constructor_with_args(self):
         assert isinstance(res[3][0], list)
         assert len(res[3][0]) == 3
 
+    def test_recipient_constructor_with_protected_bytes(self):
+        r = Recipient(protected=cbor2.dumps({"foo": "bar"}))
+        assert isinstance(r.protected, dict)
+        assert r.protected["foo"] == "bar"
+
     def test_recipient_constructor_with_empty_recipients(self):
         r = Recipient(unprotected={1: -6, 4: b"our-secret"}, recipients=[])
         assert isinstance(r, Recipient)