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"""Tests for tink.python.tink.testing.keyset_builder.""" 15 16from absl.testing import absltest 17 18from tink.proto import tink_pb2 19import tink 20from tink import aead 21from tink import hybrid 22from tink.testing import keyset_builder 23 24 25def setUpModule(): 26 aead.register() 27 hybrid.register() 28 29 30class KeysetBuilderTest(absltest.TestCase): 31 32 def test_legacy_template(self): 33 template = aead.aead_key_templates.AES128_GCM 34 legacy_template = keyset_builder.legacy_template(template) 35 self.assertEqual(legacy_template.output_prefix_type, tink_pb2.LEGACY) 36 self.assertEqual(legacy_template.type_url, template.type_url) 37 self.assertEqual(legacy_template.value, template.value) 38 # check that generating legacy_template did not change template. 39 self.assertNotEqual(template.output_prefix_type, tink_pb2.LEGACY) 40 41 def test_raw_template(self): 42 template = aead.aead_key_templates.AES128_GCM 43 raw_template = keyset_builder.raw_template(template) 44 self.assertEqual(raw_template.output_prefix_type, tink_pb2.RAW) 45 self.assertEqual(raw_template.type_url, template.type_url) 46 self.assertEqual(raw_template.value, template.value) 47 # check that generating raw_template did not change template. 48 self.assertNotEqual(template.output_prefix_type, tink_pb2.RAW) 49 50 def test_keyset_handle_conversion(self): 51 keyset_handle1 = tink.new_keyset_handle(aead.aead_key_templates.AES128_GCM) 52 p1 = keyset_handle1.primitive(aead.Aead) 53 builder = keyset_builder.from_keyset_handle(keyset_handle1) 54 keyset_handle2 = builder.keyset_handle() 55 p2 = keyset_handle2.primitive(aead.Aead) 56 ciphertext = p1.encrypt(b'plaintext', b'ad') 57 self.assertEqual(p2.decrypt(ciphertext, b'ad'), b'plaintext') 58 59 def test_keyset_conversion(self): 60 builder1 = keyset_builder.new_keyset_builder() 61 new_key_id = builder1.add_new_key(aead.aead_key_templates.AES128_GCM) 62 builder1.set_primary_key(new_key_id) 63 keyset = builder1.keyset() 64 keyset_handle1 = builder1.keyset_handle() 65 p1 = keyset_handle1.primitive(aead.Aead) 66 builder2 = keyset_builder.from_keyset(keyset) 67 keyset_handle2 = builder2.keyset_handle() 68 p2 = keyset_handle2.primitive(aead.Aead) 69 ciphertext = p1.encrypt(b'plaintext', b'ad') 70 self.assertEqual(p2.decrypt(ciphertext, b'ad'), b'plaintext') 71 72 def test_asymmetric_keyset_conversion(self): 73 builder = keyset_builder.new_keyset_builder() 74 new_key_id = builder.add_new_key( 75 hybrid.hybrid_key_templates.ECIES_P256_HKDF_HMAC_SHA256_AES128_GCM) 76 builder.set_primary_key(new_key_id) 77 private_keyset = builder.keyset() 78 public_keyset = builder.public_keyset() 79 private_handle = keyset_builder.from_keyset(private_keyset).keyset_handle() 80 dec = private_handle.primitive(hybrid.HybridDecrypt) 81 public_handle = keyset_builder.from_keyset(public_keyset).keyset_handle() 82 enc = public_handle.primitive(hybrid.HybridEncrypt) 83 ciphertext = enc.encrypt(b'plaintext', b'context') 84 self.assertEqual(dec.decrypt(ciphertext, b'context'), b'plaintext') 85 86 def test_add_new_key_new_id(self): 87 builder = keyset_builder.new_keyset_builder() 88 key_id1 = builder.add_new_key(aead.aead_key_templates.AES128_GCM) 89 key_id2 = builder.add_new_key(aead.aead_key_templates.AES128_GCM) 90 self.assertNotEqual(key_id1, key_id2) 91 92 def test_set_primary_success(self): 93 builder = keyset_builder.new_keyset_builder() 94 secondary_key_id = builder.add_new_key(aead.aead_key_templates.AES128_GCM) 95 builder.set_primary_key(secondary_key_id) 96 97 def test_operation_on_unknown_key_fails(self): 98 builder = keyset_builder.new_keyset_builder() 99 key_id = builder.add_new_key( 100 aead.aead_key_templates.AES128_GCM) 101 unknown_key_id = key_id + 1 102 with self.assertRaises(tink.TinkError): 103 builder.set_primary_key(unknown_key_id) 104 with self.assertRaises(tink.TinkError): 105 builder.enable_key(unknown_key_id) 106 with self.assertRaises(tink.TinkError): 107 builder.disable_key(unknown_key_id) 108 with self.assertRaises(tink.TinkError): 109 builder.delete_key(unknown_key_id) 110 111 def test_key_rotation(self): 112 builder = keyset_builder.new_keyset_builder() 113 older_key_id = builder.add_new_key(aead.aead_key_templates.AES128_GCM) 114 builder.set_primary_key(older_key_id) 115 p1 = builder.keyset_handle().primitive(aead.Aead) 116 117 newer_key_id = builder.add_new_key(aead.aead_key_templates.AES128_GCM) 118 p2 = builder.keyset_handle().primitive(aead.Aead) 119 120 builder.set_primary_key(newer_key_id) 121 p3 = builder.keyset_handle().primitive(aead.Aead) 122 123 builder.disable_key(older_key_id) 124 p4 = builder.keyset_handle().primitive(aead.Aead) 125 126 self.assertNotEqual(older_key_id, newer_key_id) 127 # p1 encrypts with the older key. So p1, p2 and p3 can decrypt it, 128 # but not p4. 129 ciphertext1 = p1.encrypt(b'plaintext', b'ad') 130 self.assertEqual(p1.decrypt(ciphertext1, b'ad'), b'plaintext') 131 self.assertEqual(p2.decrypt(ciphertext1, b'ad'), b'plaintext') 132 self.assertEqual(p3.decrypt(ciphertext1, b'ad'), b'plaintext') 133 with self.assertRaises(tink.TinkError): 134 _ = p4.decrypt(ciphertext1, b'ad') 135 136 # p2 encrypts with the older key. So p1, p2 and p3 can decrypt it, 137 # but not p4. 138 ciphertext2 = p2.encrypt(b'plaintext', b'ad') 139 self.assertEqual(p1.decrypt(ciphertext2, b'ad'), b'plaintext') 140 self.assertEqual(p2.decrypt(ciphertext2, b'ad'), b'plaintext') 141 self.assertEqual(p3.decrypt(ciphertext2, b'ad'), b'plaintext') 142 with self.assertRaises(tink.TinkError): 143 _ = p4.decrypt(ciphertext2, b'ad') 144 145 # p3 encrypts with the newer key. So p2, p3 and p4 can decrypt it, 146 # but not p1. 147 ciphertext3 = p3.encrypt(b'plaintext', b'ad') 148 with self.assertRaises(tink.TinkError): 149 _ = p1.decrypt(ciphertext3, b'ad') 150 self.assertEqual(p2.decrypt(ciphertext3, b'ad'), b'plaintext') 151 self.assertEqual(p3.decrypt(ciphertext3, b'ad'), b'plaintext') 152 self.assertEqual(p4.decrypt(ciphertext3, b'ad'), b'plaintext') 153 154 # p4 encrypts with the newer key. So p2, p3 and p4 can decrypt it, 155 # but not p1. 156 ciphertext4 = p4.encrypt(b'plaintext', b'ad') 157 with self.assertRaises(tink.TinkError): 158 _ = p1.decrypt(ciphertext4, b'ad') 159 self.assertEqual(p2.decrypt(ciphertext4, b'ad'), b'plaintext') 160 self.assertEqual(p3.decrypt(ciphertext4, b'ad'), b'plaintext') 161 self.assertEqual(p4.decrypt(ciphertext4, b'ad'), b'plaintext') 162 163 164if __name__ == '__main__': 165 absltest.main() 166