diff --git a/jwcrypto/jwe.py b/jwcrypto/jwe.py index 34c2b5d..3d6942d 100644 --- a/jwcrypto/jwe.py +++ b/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") diff --git a/jwcrypto/jws.py b/jwcrypto/jws.py index 38c9175..b0ec073 100644 --- a/jwcrypto/jws.py +++ b/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: diff --git a/jwcrypto/tests.py b/jwcrypto/tests.py index e3f7529..d384de2 100644 --- a/jwcrypto/tests.py +++ b/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)