1# Copyright 2020 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Tests for tink.python.tink.aead_wrapper.""" 16 17from absl.testing import absltest 18import tink 19from tink import prf 20from tink.testing import keyset_builder 21 22 23TEMPLATE = prf.prf_key_templates.HMAC_SHA256 24 25 26def setUpModule(): 27 prf.register() 28 29 30class PrfSetWrapperTest(absltest.TestCase): 31 32 def test_wrapped_output_is_equal(self): 33 keyset_handle = tink.new_keyset_handle(TEMPLATE) 34 primitive = keyset_handle.primitive(prf.PrfSet) 35 output = primitive.primary().compute(b'input', output_length=31) 36 key_id = primitive.primary_id() 37 prfs = primitive.all() 38 self.assertLen(prfs, 1) 39 self.assertEqual(prfs[key_id].compute(b'input', output_length=31), output) 40 41 def test_invalid_length_fails(self): 42 keyset_handle = tink.new_keyset_handle(TEMPLATE) 43 primitive = keyset_handle.primitive(prf.PrfSet) 44 with self.assertRaises(tink.TinkError): 45 _ = primitive.primary().compute(b'input', output_length=1234567) 46 prfs = primitive.all() 47 self.assertLen(prfs, 1) 48 with self.assertRaises(tink.TinkError): 49 _ = prfs[primitive.primary_id()].compute(b'input', output_length=1234567) 50 51 def test_wrap_three_with_one_disabled(self): 52 builder = keyset_builder.new_keyset_builder() 53 id1 = builder.add_new_key(TEMPLATE) 54 id2 = builder.add_new_key(TEMPLATE) 55 disabled_id = builder.add_new_key(TEMPLATE) 56 builder.disable_key(disabled_id) 57 builder.set_primary_key(id1) 58 prf_set1 = builder.keyset_handle().primitive(prf.PrfSet) 59 builder.set_primary_key(id2) 60 prf_set2 = builder.keyset_handle().primitive(prf.PrfSet) 61 self.assertNotEqual(id1, id2) 62 self.assertEqual(prf_set1.primary_id(), id1) 63 self.assertEqual(prf_set2.primary_id(), id2) 64 65 output1 = prf_set1.primary().compute(b'input', output_length=31) 66 output2 = prf_set2.primary().compute(b'input', output_length=31) 67 self.assertNotEqual(output1, output2) 68 prfs = prf_set1.all() 69 self.assertLen(prfs, 2) 70 self.assertEqual(prfs[id1].compute(b'input', output_length=31), 71 output1) 72 self.assertEqual(prfs[id2].compute(b'input', output_length=31), 73 output2) 74 75 76if __name__ == '__main__': 77 absltest.main() 78