1# Copyright 2022 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"""Provides methods to create keys and keysets in cross language tests. 16""" 17 18import io 19from typing import Any 20from typing import Callable 21 22import tink 23from tink import aead 24from tink import cleartext_keyset_handle 25from tink import daead 26from tink import hybrid 27from tink import jwt 28from tink import mac 29from tink import prf 30from tink import signature 31from tink import streaming_aead 32 33from tink.proto import tink_pb2 34import tink_config 35from util import key_util 36from util.test_keys import _test_keys_container 37from util.test_keys import _test_keys_db 38 39 40_CREATE_NEW_KEY_MESSAGE_TEMPLATE = """ 41Unable to retrieve stored key for template: 42{text_format} 43To create a new key with this template, run: 44blaze test --trim_test_configuration \\ 45 //third_party/tink/testing/cross_language/util:testing_servers_test \\ 46 --test_arg=--force_failure_for_adding_key_to_db \\ 47 --test_arg=--hex_template={hex_template} \\ 48 --test_output=errors 49""".strip() 50 51 52def _use_stored_key(template: tink_pb2.KeyTemplate) -> bool: 53 """Returns true for templates for which we should use _test_keys_db.py.""" 54 # We cannot yet create ChaCha20Poly1305Keys in Python. 55 if (template.type_url == 56 'type.googleapis.com/google.crypto.tink.ChaCha20Poly1305Key'): 57 return True 58 # Creating RSA Keys is very slow. 59 if (template.type_url == 60 'type.googleapis.com/google.crypto.tink.RsaSsaPkcs1PrivateKey'): 61 return True 62 # Creating RSA Keys is very slow. 63 if (template.type_url == 64 'type.googleapis.com/google.crypto.tink.RsaSsaPssPrivateKey'): 65 return True 66 # Creating RSA Keys is very slow. 67 if (template.type_url == 68 'type.googleapis.com/google.crypto.tink.JwtRsaSsaPkcs1PrivateKey'): 69 return True 70 # Creating RSA Keys is very slow. 71 if (template.type_url == 72 'type.googleapis.com/google.crypto.tink.JwtRsaSsaPssPrivateKey'): 73 return True 74 return False 75 76 77def new_or_stored_key( 78 template: tink_pb2.KeyTemplate, 79 container: _test_keys_container.TestKeysContainer = _test_keys_db.db, 80 use_stored_key: Callable[[tink_pb2.KeyTemplate], bool] = _use_stored_key 81) -> tink_pb2.Keyset.Key: 82 """Returns either a new key or one which is stored in the passed in db. 83 84 The arguments 'container' and 'use_stored_key' are for testing and typically 85 do not need to be used. 86 87 Args: 88 template: the template for which to get a key 89 container: the container with test keys, per default the container defined 90 globally in _test_keys_db 91 use_stored_key: a function which returns for a given template whether we 92 should use a precomputed key, defaults to an internal function 93 """ 94 95 if not use_stored_key(template): 96 handle = tink.new_keyset_handle(template) 97 buf = io.BytesIO() 98 writer = tink.BinaryKeysetWriter(buf) 99 cleartext_keyset_handle.write(writer, handle) 100 keyset = tink_pb2.Keyset.FromString(buf.getvalue()) 101 return keyset.key[0] 102 103 try: 104 return container.get_key(template) 105 except KeyError: 106 raise ValueError( 107 _CREATE_NEW_KEY_MESSAGE_TEMPLATE.format( 108 text_format=key_util.text_format(template), 109 hex_template=template.SerializeToString().hex())) from None 110 111 112def new_or_stored_keyset( 113 template: tink_pb2.KeyTemplate, 114 container: _test_keys_container.TestKeysContainer = _test_keys_db.db, 115 use_stored_key: Callable[[tink_pb2.KeyTemplate], bool] = _use_stored_key 116) -> bytes: 117 """Returns a new keyset with a single new or stored key. 118 119 The arguments 'container' and 'use_stored_key' are for testing and typically 120 do not need to be used. 121 122 Args: 123 template: the template for which to get a key 124 container: the container with test keys, per default the container defined 125 globally in _test_keys_db 126 use_stored_key: a function which returns for a given template whether we 127 should use a precomputed key, defaults to an internal function 128 """ 129 key = new_or_stored_key(template, container, use_stored_key) 130 keyset = tink_pb2.Keyset(key=[key], primary_key_id=key.key_id) 131 return keyset.SerializeToString() 132 133 134def _some_template_for_primitive(primitive: Any) -> tink_pb2.KeyTemplate: 135 """Returns an arbitrary template for the given primitive.""" 136 if primitive == aead.Aead: 137 return aead.aead_key_templates.AES128_GCM 138 if primitive == daead.DeterministicAead: 139 return daead.deterministic_aead_key_templates.AES256_SIV 140 if primitive == streaming_aead.StreamingAead: 141 return streaming_aead.streaming_aead_key_templates.AES256_CTR_HMAC_SHA256_1MB 142 if primitive == hybrid.HybridDecrypt: 143 return hybrid.hybrid_key_templates.DHKEM_X25519_HKDF_SHA256_HKDF_SHA256_AES_256_GCM_RAW 144 if primitive == mac.Mac: 145 return mac.mac_key_templates.HMAC_SHA256_256BITTAG 146 if primitive == signature.PublicKeySign: 147 return signature.signature_key_templates.RSA_SSA_PKCS1_4096_SHA512_F4 148 if primitive == prf.PrfSet: 149 return prf.prf_key_templates.HKDF_SHA256 150 if primitive == jwt.JwtMac: 151 return jwt.jwt_hs256_template() 152 if primitive == jwt.JwtPublicKeySign: 153 return jwt.jwt_ps512_4096_f4_template() 154 raise ValueError('Unknown primitive in _some_template_for_primitive') 155 156 157def _get_public_keyset(private_keyset: bytes) -> bytes: 158 reader = tink.BinaryKeysetReader(private_keyset) 159 keyset_handle = cleartext_keyset_handle.read(reader) 160 public_keyset_handle = keyset_handle.public_keyset_handle() 161 public_keyset = io.BytesIO() 162 cleartext_keyset_handle.write( 163 tink.BinaryKeysetWriter(public_keyset), public_keyset_handle) 164 return public_keyset.getvalue() 165 166 167def some_keyset_for_primitive(primitive: Any) -> bytes: 168 """Returns an arbitrary keyset for the given primitive.""" 169 if not tink_config.is_asymmetric_public_key_primitive(primitive): 170 return new_or_stored_keyset(_some_template_for_primitive(primitive)) 171 172 private_key_primitive = tink_config.get_private_key_primitive(primitive) 173 private_keyset = new_or_stored_keyset( 174 _some_template_for_primitive(private_key_primitive)) 175 176 return _get_public_keyset(private_keyset) 177