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. 14r"""Custom Text-format functions for Tink Keys, Keysets and Key Templates. 15 16Tink keys contain a serialized proto. Because we don't use any proto, the 17text output of the proto library is not helpful. The function 18 19key_util.text_format(msg: message.Message) 20 21is similar to text_format.MessageToString(msg), but additionally output the 22parsed serialized proto as a comment, which makes the proto human readable, 23but keep them readable by machines. For example, the AES128_EAX template 24looks like this: 25 26type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey" 27# value: [type.googleapis.com/google.crypto.tink.AesEaxKeyFormat] { 28# params { 29# iv_size: 16 30# } 31# key_size: 16 32# } 33value: "\n\002\010\020\020\020" 34output_prefix_type: TINK 35 36 37The function 38 39assert_tink_proto_equal(self, a: message.Message, b: message.Message) 40 41can be used in tests to assert that two protos must be equal. If they are 42not equal, the function tries to output a meaningfull error message. 43""" 44 45import copy 46from typing import Any, Optional 47 48# copybara:tink_placeholder(encoder) 49from google.protobuf import descriptor 50from google.protobuf import message 51from google.protobuf import text_encoding 52from google.protobuf import text_format as proto_text_format 53from tink.proto import aes_cmac_pb2 54from tink.proto import aes_cmac_prf_pb2 55from tink.proto import aes_ctr_hmac_aead_pb2 56from tink.proto import aes_ctr_hmac_streaming_pb2 57from tink.proto import aes_eax_pb2 58from tink.proto import aes_gcm_hkdf_streaming_pb2 59from tink.proto import aes_gcm_pb2 60from tink.proto import aes_gcm_siv_pb2 61from tink.proto import aes_siv_pb2 62from tink.proto import chacha20_poly1305_pb2 63from tink.proto import ecdsa_pb2 64from tink.proto import ecies_aead_hkdf_pb2 65from tink.proto import ed25519_pb2 66from tink.proto import hkdf_prf_pb2 67from tink.proto import hmac_pb2 68from tink.proto import hmac_prf_pb2 69from tink.proto import hpke_pb2 70from tink.proto import jwt_ecdsa_pb2 71from tink.proto import jwt_hmac_pb2 72from tink.proto import jwt_rsa_ssa_pkcs1_pb2 73from tink.proto import jwt_rsa_ssa_pss_pb2 74from tink.proto import kms_aead_pb2 75from tink.proto import kms_envelope_pb2 76from tink.proto import rsa_ssa_pkcs1_pb2 77from tink.proto import rsa_ssa_pss_pb2 78from tink.proto import xchacha20_poly1305_pb2 79 80 81TYPE_STRING = 9 82TYPE_MESSAGE = 11 83TYPE_BYTES = 12 84TYPE_ENUM = 14 85LABEL_REPEATED = 3 86 87TYPE_PREFIX = 'type.googleapis.com/' 88 89 90class KeyProto: 91 """A map from type URLs to key protos and key format protos.""" 92 93 _from_url = {} 94 _format_from_url = {} 95 96 @classmethod 97 def from_url(cls, type_url: str) -> Any: 98 return cls._from_url[type_url] 99 100 @classmethod 101 def format_from_url(cls, type_url: str) -> Any: 102 return cls._format_from_url[type_url] 103 104 @classmethod 105 def add_key_type(cls, key_type: Any, key_format_type: Any): 106 type_url = TYPE_PREFIX + key_type.DESCRIPTOR.full_name 107 cls._from_url[type_url] = key_type 108 cls._format_from_url[type_url] = key_format_type 109 110 111KeyProto.add_key_type(aes_eax_pb2.AesEaxKey, aes_eax_pb2.AesEaxKeyFormat) 112KeyProto.add_key_type(aes_gcm_pb2.AesGcmKey, aes_gcm_pb2.AesGcmKeyFormat) 113KeyProto.add_key_type(aes_gcm_siv_pb2.AesGcmSivKey, 114 aes_gcm_siv_pb2.AesGcmSivKeyFormat) 115KeyProto.add_key_type(aes_ctr_hmac_aead_pb2.AesCtrHmacAeadKey, 116 aes_ctr_hmac_aead_pb2.AesCtrHmacAeadKeyFormat) 117KeyProto.add_key_type(chacha20_poly1305_pb2.ChaCha20Poly1305Key, 118 chacha20_poly1305_pb2.ChaCha20Poly1305KeyFormat) 119KeyProto.add_key_type(xchacha20_poly1305_pb2.XChaCha20Poly1305Key, 120 xchacha20_poly1305_pb2.XChaCha20Poly1305KeyFormat) 121KeyProto.add_key_type(aes_siv_pb2.AesSivKey, aes_siv_pb2.AesSivKeyFormat) 122KeyProto.add_key_type(aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKey, 123 aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat) 124KeyProto.add_key_type(aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKey, 125 aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat) 126KeyProto.add_key_type(ecies_aead_hkdf_pb2.EciesAeadHkdfPrivateKey, 127 ecies_aead_hkdf_pb2.EciesAeadHkdfKeyFormat) 128KeyProto.add_key_type(ecies_aead_hkdf_pb2.EciesAeadHkdfPublicKey, 129 ecies_aead_hkdf_pb2.EciesAeadHkdfKeyFormat) 130KeyProto.add_key_type(hpke_pb2.HpkePrivateKey, hpke_pb2.HpkeKeyFormat) 131KeyProto.add_key_type(hpke_pb2.HpkePublicKey, hpke_pb2.HpkeKeyFormat) 132KeyProto.add_key_type(aes_cmac_pb2.AesCmacKey, aes_cmac_pb2.AesCmacKeyFormat) 133KeyProto.add_key_type(hmac_pb2.HmacKey, hmac_pb2.HmacKeyFormat) 134KeyProto.add_key_type(ecdsa_pb2.EcdsaPrivateKey, ecdsa_pb2.EcdsaKeyFormat) 135KeyProto.add_key_type(ecdsa_pb2.EcdsaPublicKey, ecdsa_pb2.EcdsaKeyFormat) 136KeyProto.add_key_type(ed25519_pb2.Ed25519PrivateKey, 137 ed25519_pb2.Ed25519KeyFormat) 138KeyProto.add_key_type(ed25519_pb2.Ed25519PublicKey, 139 ed25519_pb2.Ed25519KeyFormat) 140KeyProto.add_key_type(rsa_ssa_pkcs1_pb2.RsaSsaPkcs1PrivateKey, 141 rsa_ssa_pkcs1_pb2.RsaSsaPkcs1KeyFormat) 142KeyProto.add_key_type(rsa_ssa_pkcs1_pb2.RsaSsaPkcs1PublicKey, 143 rsa_ssa_pkcs1_pb2.RsaSsaPkcs1KeyFormat) 144KeyProto.add_key_type(rsa_ssa_pss_pb2.RsaSsaPssPrivateKey, 145 rsa_ssa_pss_pb2.RsaSsaPssKeyFormat) 146KeyProto.add_key_type(rsa_ssa_pss_pb2.RsaSsaPssPublicKey, 147 rsa_ssa_pss_pb2.RsaSsaPssKeyFormat) 148KeyProto.add_key_type(aes_cmac_prf_pb2.AesCmacPrfKey, 149 aes_cmac_prf_pb2.AesCmacPrfKeyFormat) 150KeyProto.add_key_type(hmac_prf_pb2.HmacPrfKey, hmac_prf_pb2.HmacPrfKeyFormat) 151KeyProto.add_key_type(hkdf_prf_pb2.HkdfPrfKey, hkdf_prf_pb2.HkdfPrfKeyFormat) 152KeyProto.add_key_type(jwt_ecdsa_pb2.JwtEcdsaPrivateKey, 153 jwt_ecdsa_pb2.JwtEcdsaKeyFormat) 154KeyProto.add_key_type(jwt_ecdsa_pb2.JwtEcdsaPublicKey, 155 jwt_ecdsa_pb2.JwtEcdsaKeyFormat) 156KeyProto.add_key_type(jwt_hmac_pb2.JwtHmacKey, jwt_hmac_pb2.JwtHmacKeyFormat) 157KeyProto.add_key_type(jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1PrivateKey, 158 jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1KeyFormat) 159KeyProto.add_key_type(jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1PublicKey, 160 jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1KeyFormat) 161KeyProto.add_key_type(jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssPrivateKey, 162 jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssKeyFormat) 163KeyProto.add_key_type(jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssPublicKey, 164 jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssKeyFormat) 165KeyProto.add_key_type(kms_aead_pb2.KmsAeadKey, kms_aead_pb2.KmsAeadKeyFormat) 166KeyProto.add_key_type(kms_envelope_pb2.KmsEnvelopeAeadKey, 167 kms_envelope_pb2.KmsEnvelopeAeadKeyFormat) 168 169 170def _text_format_field(value: Any, 171 field: descriptor.FieldDescriptor, 172 indent: str) -> str: 173 """Returns a text formated proto field.""" 174 if field.type == TYPE_MESSAGE: 175 output = [ 176 indent + field.name + ' {', 177 _normalize_and_text_format_message(value, indent + ' '), 178 indent + '}' 179 ] 180 return '\n'.join(output) 181 elif field.type == TYPE_ENUM: 182 value_name = field.enum_type.values_by_number[value].name 183 return indent + field.name + ': ' + value_name 184 elif field.type in [TYPE_STRING, TYPE_BYTES]: 185 return (indent + field.name + ': "' + text_encoding.CEscape(value, False) + 186 '"') 187 else: 188 return indent + field.name + ': ' + str(value) 189 190 191def _normalize_and_text_format_message(msg: message.Message, 192 indent: str) -> str: 193 """Returns a text formated proto message and changes msg to be canonical. 194 195 Args: 196 msg: the proto to be formated. 197 indent: the indentation prefix of each line in the output. 198 199 Returns: 200 A proto text format output, where serialized fields are deserialized in 201 a comment. 202 """ 203 output = [] 204 fields = list(msg.DESCRIPTOR.fields) 205 # special case for Tinks custom 'any' proto. 206 if (msg.DESCRIPTOR.full_name == 'google.crypto.tink.KeyTemplate' or 207 msg.DESCRIPTOR.full_name == 'google.crypto.tink.KeyData'): 208 type_url = getattr(msg, 'type_url') # Pytype requires to use getattr 209 output.append( 210 _text_format_field(type_url, fields[0], indent)) 211 value = getattr(msg, 'value') 212 if msg.DESCRIPTOR.full_name == 'google.crypto.tink.KeyTemplate': 213 # In KeyTemplates, type_url does not match the proto type used. 214 proto_type = KeyProto.format_from_url(type_url) 215 else: 216 proto_type = KeyProto.from_url(type_url) 217 # parse 'value' and text format the content in a comment. 218 field_proto = proto_type.FromString(value) 219 output.append(indent + '# value: [' + TYPE_PREFIX + 220 proto_type.DESCRIPTOR.full_name + '] {') 221 formatted_message = _normalize_and_text_format_message( 222 field_proto, indent + '# ') 223 if formatted_message: 224 output.append(formatted_message) 225 output.append(indent + '# }') 226 # Serialize message again so it is canonicalized 227 # We require here that proto serialization is in increasing field order 228 # (Tink protos are basically unchangeable, so we don't need to worry about 229 # unknown fields). This is not guaranteed by proto, but is currently the 230 # case. If this ever changes we either hopefully have already a better 231 # solution in Tink, or else the proto team provides us with a reflection 232 # based API to do this (as they do in C++.) In this case, we simply use the 233 # slow API here. 234 value = field_proto.SerializeToString(deterministic = True) 235 setattr(msg, 'value', value) 236 output.append( 237 _text_format_field(value, fields[1], indent)) 238 fields = fields[2:] 239 for field in fields: 240 if field.label == LABEL_REPEATED: 241 for value in getattr(msg, field.name): 242 output.append(_text_format_field(value, field, indent)) 243 else: 244 output.append( 245 _text_format_field( 246 getattr(msg, field.name), field, indent)) 247 return '\n'.join(output) 248 249 250def text_format(msg: message.Message) -> str: 251 msgcopy = copy.deepcopy(msg) 252 return _normalize_and_text_format_message(msgcopy, '') 253 254 255def parse_text_format(serialized: str, msg: message.Message) -> None: 256 # Different from binary parsing, text_format.Parse does not Clear the message. 257 msg.Clear() 258 proto_text_format.Parse(serialized, msg) 259 serialized_copy = text_format(msg) 260 assert serialized_copy == serialized, serialized_copy 261 262 263def assert_tink_proto_equal(self, 264 a: message.Message, 265 b: message.Message, 266 msg: Optional[str] = None) -> None: 267 """Fails with a useful error if a and b aren't equal.""" 268 a_copy = copy.deepcopy(a) 269 b_copy = copy.deepcopy(b) 270 271 self.assertMultiLineEqual( 272 _normalize_and_text_format_message(a_copy, ''), 273 _normalize_and_text_format_message(b_copy, ''), 274 msg=msg) 275