| # 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.tink.core.primitive_set.""" |
| |
| 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 import aead |
| from tink import core |
| from tink import mac |
| from tink.testing import helper |
| |
| |
| class PrimitiveSetTest(absltest.TestCase): |
| |
| def test_primitive_returns_entry(self): |
| key = helper.fake_key(key_id=1) |
| fake_mac = helper.FakeMac('FakeMac') |
| primitive_set = core.new_primitive_set(mac.Mac) |
| primitive_set.add_primitive(fake_mac, key) |
| entries = primitive_set.primitive(key) |
| self.assertLen(entries, 1) |
| entry = entries[0] |
| self.assertEqual(fake_mac, entry.primitive) |
| self.assertEqual(tink_pb2.ENABLED, entry.status) |
| self.assertEqual(core.crypto_format.output_prefix(key), entry.identifier) |
| |
| def test_unknown_key_returns_empty_list(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| unknown_key = helper.fake_key(key_id=1) |
| self.assertEqual(primitive_set.primitive(unknown_key), []) |
| |
| def test_primitive_from_identifier_returns_entry(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key = helper.fake_key(key_id=1) |
| fake_mac = helper.FakeMac('FakeMac') |
| primitive_set.add_primitive(fake_mac, key) |
| |
| ident = core.crypto_format.output_prefix(key) |
| entries = primitive_set.primitive_from_identifier(ident) |
| self.assertLen(entries, 1) |
| entry = entries[0] |
| self.assertEqual(fake_mac, entry.primitive) |
| self.assertEqual(tink_pb2.ENABLED, entry.status) |
| self.assertEqual(ident, entry.identifier) |
| |
| def test_list_of_entries_can_be_modified(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key = helper.fake_key(key_id=1) |
| primitive_set.add_primitive(helper.FakeMac('FakeMac'), key) |
| entries = primitive_set.primitive(key) |
| entries.append('Something') |
| self.assertLen(primitive_set.primitive(key), 1) |
| |
| def test_primary_returns_primary(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key = helper.fake_key(key_id=1) |
| fake_mac = helper.FakeMac('FakeMac') |
| entry = primitive_set.add_primitive(fake_mac, key) |
| primitive_set.set_primary(entry) |
| |
| entry = primitive_set.primary() |
| self.assertEqual(fake_mac, entry.primitive) |
| self.assertEqual(tink_pb2.ENABLED, entry.status) |
| self.assertEqual(core.crypto_format.output_prefix(key), entry.identifier) |
| |
| def test_primary_returns_none(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| primitive_set.add_primitive( |
| helper.FakeMac('FakeMac'), helper.fake_key(key_id=1)) |
| self.assertEqual(primitive_set.primary(), None) |
| |
| def test_same_key_id_and_prefix_type(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key1 = helper.fake_key(key_id=1, status=tink_pb2.ENABLED) |
| fake_mac1 = helper.FakeMac('FakeMac1') |
| primitive_set.add_primitive(fake_mac1, key1) |
| key2 = helper.fake_key(key_id=1, status=tink_pb2.DISABLED) |
| fake_mac2 = helper.FakeMac('FakeMac2') |
| primitive_set.add_primitive(fake_mac2, key2) |
| |
| expected_ident = core.crypto_format.output_prefix(key1) |
| entries = primitive_set.primitive(key1) |
| self.assertLen(entries, 2) |
| self.assertEqual(fake_mac1, entries[0].primitive) |
| self.assertEqual(fake_mac2, entries[1].primitive) |
| self.assertEqual(tink_pb2.ENABLED, entries[0].status) |
| self.assertEqual(tink_pb2.DISABLED, entries[1].status) |
| self.assertEqual(expected_ident, entries[0].identifier) |
| self.assertEqual(expected_ident, entries[1].identifier) |
| self.assertLen(primitive_set.primitive(key2), 2) |
| |
| def test_same_key_id_but_different_prefix_type(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key1 = helper.fake_key(key_id=1, output_prefix_type=tink_pb2.TINK) |
| fake_mac1 = helper.FakeMac('FakeMac1') |
| primitive_set.add_primitive(fake_mac1, key1) |
| key2 = helper.fake_key(key_id=1, output_prefix_type=tink_pb2.LEGACY) |
| fake_mac2 = helper.FakeMac('FakeMac2') |
| primitive_set.add_primitive(fake_mac2, key2) |
| |
| entries1 = primitive_set.primitive(key1) |
| self.assertLen(entries1, 1) |
| self.assertEqual(fake_mac1, entries1[0].primitive) |
| self.assertEqual(tink_pb2.ENABLED, entries1[0].status) |
| self.assertEqual(core.crypto_format.output_prefix(key1), |
| entries1[0].identifier) |
| |
| entries2 = primitive_set.primitive(key2) |
| self.assertLen(entries2, 1) |
| self.assertEqual(fake_mac2, entries2[0].primitive) |
| self.assertEqual(tink_pb2.ENABLED, entries2[0].status) |
| self.assertEqual(core.crypto_format.output_prefix(key2), |
| entries2[0].identifier) |
| |
| def test_add_invalid_key_fails(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| key = helper.fake_key() |
| key.ClearField('output_prefix_type') |
| with self.assertRaisesRegex(core.TinkError, 'invalid OutputPrefixType'): |
| primitive_set.add_primitive(helper.FakeMac(), key) |
| |
| def test_add_wrong_primitive_fails(self): |
| primitive_set = core.new_primitive_set(aead.Aead) |
| with self.assertRaisesRegex(core.TinkError, |
| 'The primitive is not an instance of '): |
| primitive_set.add_primitive(helper.FakeMac(), helper.fake_key()) |
| |
| def test_primitive_class(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| self.assertEqual(primitive_set.primitive_class(), mac.Mac) |
| |
| def test_raw_primitives(self): |
| primitive_set = core.new_primitive_set(mac.Mac) |
| primitive_set.add_primitive( |
| helper.FakeMac('FakeMac1'), helper.fake_key(key_id=1)) |
| key2 = helper.fake_key(key_id=1, output_prefix_type=tink_pb2.RAW) |
| fake_mac2 = helper.FakeMac('FakeMac2') |
| primitive_set.add_primitive(fake_mac2, key2) |
| key3 = helper.fake_key( |
| key_id=3, status=tink_pb2.DISABLED, output_prefix_type=tink_pb2.RAW) |
| fake_mac3 = helper.FakeMac('FakeMac3') |
| primitive_set.add_primitive(fake_mac3, key3) |
| |
| entries = primitive_set.raw_primitives() |
| self.assertLen(entries, 2) |
| self.assertEqual(fake_mac2, entries[0].primitive) |
| self.assertEqual(tink_pb2.ENABLED, entries[0].status) |
| self.assertEqual(core.crypto_format.RAW_PREFIX, |
| entries[0].identifier) |
| self.assertEqual(fake_mac3, entries[1].primitive) |
| self.assertEqual(tink_pb2.DISABLED, entries[1].status) |
| self.assertEqual(core.crypto_format.RAW_PREFIX, |
| entries[1].identifier) |
| |
| |
| if __name__ == '__main__': |
| absltest.main() |