| # Copyright 2019 Google LLC. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """Tests for tink.python.aead_wrapper.""" |
| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| from absl.testing import absltest |
| from tink.proto import tink_pb2 |
| from tink.python.aead import aead |
| from tink.python.aead import aead_wrapper |
| from tink.python.core import primitive_set |
| from tink.python.core import tink_error |
| from tink.python.testing import helper |
| |
| |
| class AeadWrapperTest(absltest.TestCase): |
| |
| def new_primitive_key_pair(self, key_id, output_prefix_type): |
| fake_key = helper.fake_key(key_id=key_id, |
| output_prefix_type=output_prefix_type) |
| fake_aead = helper.FakeAead('fakeAead {}'.format(key_id)) |
| return fake_aead, fake_key |
| |
| def test_encrypt_decrypt(self): |
| primitive, key = self.new_primitive_key_pair(1234, tink_pb2.TINK) |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| entry = pset.add_primitive(primitive, key) |
| pset.set_primary(entry) |
| |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| |
| plaintext = b'plaintext' |
| associated_data = b'associated_data' |
| ciphertext = wrapped_aead.encrypt(plaintext, associated_data) |
| self.assertEqual( |
| wrapped_aead.decrypt(ciphertext, associated_data), plaintext) |
| |
| def test_encrypt_decrypt_with_key_rotation(self): |
| primitive, key = self.new_primitive_key_pair(1234, tink_pb2.TINK) |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| entry = pset.add_primitive(primitive, key) |
| pset.set_primary(entry) |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| ciphertext = wrapped_aead.encrypt(b'plaintext', b'associated_data') |
| |
| new_primitive, new_key = self.new_primitive_key_pair(5678, tink_pb2.TINK) |
| new_entry = pset.add_primitive(new_primitive, new_key) |
| pset.set_primary(new_entry) |
| new_ciphertext = wrapped_aead.encrypt(b'new_plaintext', |
| b'new_associated_data') |
| |
| self.assertEqual( |
| wrapped_aead.decrypt(ciphertext, b'associated_data'), |
| b'plaintext') |
| self.assertEqual( |
| wrapped_aead.decrypt(new_ciphertext, b'new_associated_data'), |
| b'new_plaintext') |
| |
| def test_encrypt_decrypt_with_key_rotation_from_raw(self): |
| primitive, raw_key = self.new_primitive_key_pair(1234, tink_pb2.RAW) |
| old_raw_ciphertext = primitive.encrypt(b'plaintext', b'associated_data') |
| |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| pset.add_primitive(primitive, raw_key) |
| new_primitive, new_key = self.new_primitive_key_pair(5678, tink_pb2.TINK) |
| new_entry = pset.add_primitive(new_primitive, new_key) |
| pset.set_primary(new_entry) |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| new_ciphertext = wrapped_aead.encrypt(b'new_plaintext', |
| b'new_associated_data') |
| |
| self.assertEqual( |
| wrapped_aead.decrypt(old_raw_ciphertext, b'associated_data'), |
| b'plaintext') |
| self.assertEqual( |
| wrapped_aead.decrypt(new_ciphertext, b'new_associated_data'), |
| b'new_plaintext') |
| |
| def test_encrypt_decrypt_two_raw_keys(self): |
| primitive1, raw_key1 = self.new_primitive_key_pair(1234, tink_pb2.RAW) |
| primitive2, raw_key2 = self.new_primitive_key_pair(5678, tink_pb2.RAW) |
| raw_ciphertext1 = primitive1.encrypt(b'plaintext1', b'associated_data1') |
| raw_ciphertext2 = primitive2.encrypt(b'plaintext2', b'associated_data2') |
| |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| pset.add_primitive(primitive1, raw_key1) |
| pset.set_primary( |
| pset.add_primitive(primitive2, raw_key2)) |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| |
| self.assertEqual( |
| wrapped_aead.decrypt(raw_ciphertext1, b'associated_data1'), |
| b'plaintext1') |
| self.assertEqual( |
| wrapped_aead.decrypt(raw_ciphertext2, b'associated_data2'), |
| b'plaintext2') |
| self.assertEqual( |
| wrapped_aead.decrypt( |
| wrapped_aead.encrypt(b'plaintext', b'associated_data'), |
| b'associated_data'), |
| b'plaintext') |
| |
| def test_decrypt_unknown_ciphertext_fails(self): |
| unknown_primitive = helper.FakeAead('unknownFakeAead') |
| unknown_ciphertext = unknown_primitive.encrypt(b'plaintext', |
| b'associated_data') |
| |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| primitive, raw_key = self.new_primitive_key_pair(1234, tink_pb2.RAW) |
| new_primitive, new_key = self.new_primitive_key_pair(5678, tink_pb2.TINK) |
| pset.add_primitive(primitive, raw_key) |
| new_entry = pset.add_primitive(new_primitive, new_key) |
| pset.set_primary(new_entry) |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| |
| with self.assertRaisesRegex(tink_error.TinkError, 'Decryption failed'): |
| wrapped_aead.decrypt(unknown_ciphertext, b'associated_data') |
| |
| def test_decrypt_wrong_associated_data_fails(self): |
| primitive, key = self.new_primitive_key_pair(1234, tink_pb2.TINK) |
| pset = primitive_set.new_primitive_set(aead.Aead) |
| entry = pset.add_primitive(primitive, key) |
| pset.set_primary(entry) |
| wrapped_aead = aead_wrapper.AeadWrapper().wrap(pset) |
| |
| ciphertext = wrapped_aead.encrypt(b'plaintext', b'associated_data') |
| with self.assertRaisesRegex(tink_error.TinkError, 'Decryption failed'): |
| wrapped_aead.decrypt(ciphertext, b'wrong_associated_data') |
| |
| |
| if __name__ == '__main__': |
| absltest.main() |