Fix an issue with parsing headers with 'kid'

jwcrypto/jwe.py

@@ -162,7 +162,10 @@ def _get_jose_header(self, header=None):
             uh = json_decode(self.objects['unprotected'])
             jh = self._merge_headers(jh, uh)
         if header:
-            rh = json_decode(header)
+            if isinstance(header, dict):
+                rh = header
+            else:
+                rh = json_decode(header)
             jh = self._merge_headers(jh, rh)
         return jh
 
@@ -382,11 +385,11 @@ def _decrypt(self, key, ppe, max_plaintext=default_max_plaintext_size):
 
         if isinstance(key, JWKSet):
             keys = key
-            if 'kid' in self.jose_header:
-                kid_keys = key.get_keys(self.jose_header['kid'])
+            if 'kid' in jh:
+                kid_keys = key.get_keys(jh['kid'])
                 if not kid_keys:
                     raise JWKeyNotFound('Key ID {} not in key set'.format(
-                                        self.jose_header['kid']))
+                                        jh['kid']))
                 keys = kid_keys
 
             for k in keys:
@@ -565,6 +568,11 @@ def payload(self):
 
     @property
     def jose_header(self):
+        if 'recipients' in self.objects:
+            jhl = []
+            for rec in self.objects['recipients']:
+                jhl.append(self._get_jose_header(rec.get('header')))
+            return jhl
         jh = self._get_jose_header(self.objects.get('header'))
         if len(jh) == 0:
             raise InvalidJWEOperation("JOSE Header not available")

jwcrypto/jws.py

@@ -299,11 +299,11 @@ def _verify(self, alg, key, payload, signature, protected, header=None):
             self.verifylog.append("Success")
         elif isinstance(key, JWKSet):
             keys = key
-            if 'kid' in self.jose_header:
-                kid_keys = key.get_keys(self.jose_header['kid'])
+            if 'kid' in chk_hdrs:
+                kid_keys = key.get_keys(chk_hdrs['kid'])
                 if not kid_keys:
                     raise JWKeyNotFound('Key ID {} not in key set'.format(
-                                        self.jose_header['kid']))
+                                        chk_hdrs['kid']))
                 keys = kid_keys
 
             for k in keys:

jwcrypto/tests.py

@@ -1224,6 +1224,42 @@ def test_decrypt_keyset(self):
         with self.assertRaises(JWKeyNotFound):
             s4.deserialize(s3.serialize(), ks)
 
+    def test_multiple_signatures_keyset(self):
+        ks = jwk.JWKSet()
+        key1 = jwk.JWK.generate(kty='oct', alg='HS256', kid='key1')
+        key2 = jwk.JWK.generate(kty='oct', alg='HS256', kid='key2')
+        ks.add(key1)
+        ks.add(key2)
+
+        payload = b'JSON Serialization JWS with multiple signatures'
+        s = jws.JWS(payload=payload)
+        s.add_signature(key1, protected={'alg': 'HS256', 'kid': 'key1'})
+        s.add_signature(key2, protected={'alg': 'HS256'},
+                        header={'kid': 'key2'})
+
+        serialized = s.serialize()
+        decoded = json_decode(serialized)
+        self.assertIn('signatures', decoded)
+        self.assertEqual(len(decoded['signatures']), 2)
+
+        s2 = jws.JWS()
+        s2.deserialize(serialized, ks)
+        self.assertTrue(s2.is_valid)
+        self.assertEqual(s2.payload, payload)
+
+        # Test with keys swapped
+        key1_dict = key1.export(as_dict=True)
+        key1_dict['kid'] = 'key2'
+        key2_dict = key2.export(as_dict=True)
+        key2_dict['kid'] = 'key1'
+        ks_swapped = jwk.JWKSet()
+        ks_swapped.add(jwk.JWK(**key1_dict))
+        ks_swapped.add(jwk.JWK(**key2_dict))
+
+        s3 = jws.JWS()
+        with self.assertRaises(JWKeyNotFound):
+            s3.deserialize(serialized, ks_swapped)
+
 
 E_A1_plaintext = \
     [84, 104, 101, 32, 116, 114, 117, 101, 32, 115, 105, 103, 110, 32,
@@ -1603,6 +1639,48 @@ def test_decrypt_keyset(self):
         with self.assertRaises(JWKeyNotFound):
             e4.deserialize(e3.serialize(), ks)
 
+    def test_multiple_recipients_keyset(self):
+        ks = jwk.JWKSet()
+        key1 = jwk.JWK.generate(kty='oct', alg='A128KW', kid='key1')
+        key2 = jwk.JWK.generate(kty='oct', alg='A256KW', kid='key2')
+        ks.add(key1)
+        ks.add(key2)
+
+        payload = b'JSON Serialization JWE with multiple recipients'
+        e = jwe.JWE(plaintext=payload,
+                    protected={'enc': 'A128CBC-HS256'})
+        e.add_recipient(key1, header={'alg': 'A128KW', 'kid': 'key1'})
+        e.add_recipient(key2, header={'alg': 'A256KW', 'kid': 'key2'})
+
+        serialized = e.serialize()
+        decoded = json_decode(serialized)
+        self.assertIn('recipients', decoded)
+        self.assertEqual(len(decoded['recipients']), 2)
+
+        e2 = jwe.JWE()
+        e2.deserialize(serialized, ks)
+        self.assertEqual(e2.payload, payload)
+        self.assertEqual(
+            e2.jose_header,
+            [
+                {'enc': 'A128CBC-HS256', 'alg': 'A128KW', 'kid': 'key1'},
+                {'enc': 'A128CBC-HS256', 'alg': 'A256KW', 'kid': 'key2'}
+            ]
+        )
+
+        # Test with keys swapped
+        key1_dict = key1.export(as_dict=True)
+        key1_dict['kid'] = 'key2'
+        key2_dict = key2.export(as_dict=True)
+        key2_dict['kid'] = 'key1'
+        ks_swapped = jwk.JWKSet()
+        ks_swapped.add(jwk.JWK(**key1_dict))
+        ks_swapped.add(jwk.JWK(**key2_dict))
+
+        e3 = jwe.JWE()
+        with self.assertRaises(JWKeyNotFound):
+            e3.deserialize(serialized, ks_swapped)
+
     def test_serialize_not_flattened(self):
         # JWE with flattened=False adds recipients in objects and in serialized
         e = jwe.JWE(E_A1_ex['plaintext'], flattened=False)