# Copyright 2020 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.
r"""Custom Text-format functions for Tink Keys, Keysets and Key Templates.

Tink keys contain a serialized proto. Because we don't use any proto, the
text output of the proto library is not helpful. The function

key_util.text_format(msg: message.Message)

is similar to text_format.MessageToString(msg), but additionally output the
parsed serialized proto as a comment, which makes the proto human readable,
but keep them readable by machines. For example, the AES128_EAX template
looks like this:

type_url: "type.googleapis.com/google.crypto.tink.AesEaxKey"
# value: [type.googleapis.com/google.crypto.tink.AesEaxKeyFormat] {
#   params {
#     iv_size: 16
#   }
#   key_size: 16
# }
value: "\n\002\010\020\020\020"
output_prefix_type: TINK


The function

assert_tink_proto_equal(self, a: message.Message, b: message.Message)

can be used in tests to assert that two protos must be equal. If they are
not equal, the function tries to output a meaningfull error message.
"""

from typing import Any, Optional

from google.protobuf import descriptor
from google.protobuf import message
from google.protobuf import text_encoding
from tink.proto import aes_cmac_pb2
from tink.proto import aes_cmac_prf_pb2
from tink.proto import aes_ctr_hmac_aead_pb2
from tink.proto import aes_ctr_hmac_streaming_pb2
from tink.proto import aes_eax_pb2
from tink.proto import aes_gcm_hkdf_streaming_pb2
from tink.proto import aes_gcm_pb2
from tink.proto import aes_gcm_siv_pb2
from tink.proto import aes_siv_pb2
from tink.proto import chacha20_poly1305_pb2
from tink.proto import ecdsa_pb2
from tink.proto import ecies_aead_hkdf_pb2
from tink.proto import ed25519_pb2
from tink.proto import hkdf_prf_pb2
from tink.proto import hmac_pb2
from tink.proto import hmac_prf_pb2
from tink.proto import hpke_pb2
from tink.proto import jwt_ecdsa_pb2
from tink.proto import jwt_hmac_pb2
from tink.proto import jwt_rsa_ssa_pkcs1_pb2
from tink.proto import jwt_rsa_ssa_pss_pb2
from tink.proto import kms_aead_pb2
from tink.proto import kms_envelope_pb2
from tink.proto import rsa_ssa_pkcs1_pb2
from tink.proto import rsa_ssa_pss_pb2
from tink.proto import xchacha20_poly1305_pb2


TYPE_STRING = 9
TYPE_MESSAGE = 11
TYPE_BYTES = 12
TYPE_ENUM = 14
LABEL_REPEATED = 3

TYPE_PREFIX = 'type.googleapis.com/'


class KeyProto:
  """A map from type URLs to key protos and key format protos."""

  _from_url = {}
  _format_from_url = {}

  @classmethod
  def from_url(cls, type_url: str) -> Any:
    return cls._from_url[type_url]

  @classmethod
  def format_from_url(cls, type_url: str) -> Any:
    return cls._format_from_url[type_url]

  @classmethod
  def add_key_type(cls, key_type: Any, key_format_type: Any):
    type_url = TYPE_PREFIX + key_type.DESCRIPTOR.full_name
    cls._from_url[type_url] = key_type
    cls._format_from_url[type_url] = key_format_type


KeyProto.add_key_type(aes_eax_pb2.AesEaxKey, aes_eax_pb2.AesEaxKeyFormat)
KeyProto.add_key_type(aes_gcm_pb2.AesGcmKey, aes_gcm_pb2.AesGcmKeyFormat)
KeyProto.add_key_type(aes_gcm_siv_pb2.AesGcmSivKey,
                      aes_gcm_siv_pb2.AesGcmSivKeyFormat)
KeyProto.add_key_type(aes_ctr_hmac_aead_pb2.AesCtrHmacAeadKey,
                      aes_ctr_hmac_aead_pb2.AesCtrHmacAeadKeyFormat)
KeyProto.add_key_type(chacha20_poly1305_pb2.ChaCha20Poly1305Key,
                      chacha20_poly1305_pb2.ChaCha20Poly1305KeyFormat)
KeyProto.add_key_type(xchacha20_poly1305_pb2.XChaCha20Poly1305Key,
                      xchacha20_poly1305_pb2.XChaCha20Poly1305KeyFormat)
KeyProto.add_key_type(aes_siv_pb2.AesSivKey, aes_siv_pb2.AesSivKeyFormat)
KeyProto.add_key_type(aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKey,
                      aes_ctr_hmac_streaming_pb2.AesCtrHmacStreamingKeyFormat)
KeyProto.add_key_type(aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKey,
                      aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKeyFormat)
KeyProto.add_key_type(ecies_aead_hkdf_pb2.EciesAeadHkdfPrivateKey,
                      ecies_aead_hkdf_pb2.EciesAeadHkdfKeyFormat)
KeyProto.add_key_type(ecies_aead_hkdf_pb2.EciesAeadHkdfPublicKey,
                      ecies_aead_hkdf_pb2.EciesAeadHkdfKeyFormat)
KeyProto.add_key_type(hpke_pb2.HpkePrivateKey, hpke_pb2.HpkeKeyFormat)
KeyProto.add_key_type(hpke_pb2.HpkePublicKey, hpke_pb2.HpkeKeyFormat)
KeyProto.add_key_type(aes_cmac_pb2.AesCmacKey, aes_cmac_pb2.AesCmacKeyFormat)
KeyProto.add_key_type(hmac_pb2.HmacKey, hmac_pb2.HmacKeyFormat)
KeyProto.add_key_type(ecdsa_pb2.EcdsaPrivateKey, ecdsa_pb2.EcdsaKeyFormat)
KeyProto.add_key_type(ecdsa_pb2.EcdsaPublicKey, ecdsa_pb2.EcdsaKeyFormat)
KeyProto.add_key_type(ed25519_pb2.Ed25519PrivateKey,
                      ed25519_pb2.Ed25519KeyFormat)
KeyProto.add_key_type(ed25519_pb2.Ed25519PublicKey,
                      ed25519_pb2.Ed25519KeyFormat)
KeyProto.add_key_type(rsa_ssa_pkcs1_pb2.RsaSsaPkcs1PrivateKey,
                      rsa_ssa_pkcs1_pb2.RsaSsaPkcs1KeyFormat)
KeyProto.add_key_type(rsa_ssa_pkcs1_pb2.RsaSsaPkcs1PublicKey,
                      rsa_ssa_pkcs1_pb2.RsaSsaPkcs1KeyFormat)
KeyProto.add_key_type(rsa_ssa_pss_pb2.RsaSsaPssPrivateKey,
                      rsa_ssa_pss_pb2.RsaSsaPssKeyFormat)
KeyProto.add_key_type(rsa_ssa_pss_pb2.RsaSsaPssPublicKey,
                      rsa_ssa_pss_pb2.RsaSsaPssKeyFormat)
KeyProto.add_key_type(aes_cmac_prf_pb2.AesCmacPrfKey,
                      aes_cmac_prf_pb2.AesCmacPrfKeyFormat)
KeyProto.add_key_type(hmac_prf_pb2.HmacPrfKey, hmac_prf_pb2.HmacPrfKeyFormat)
KeyProto.add_key_type(hkdf_prf_pb2.HkdfPrfKey, hkdf_prf_pb2.HkdfPrfKeyFormat)
KeyProto.add_key_type(jwt_ecdsa_pb2.JwtEcdsaPrivateKey,
                      jwt_ecdsa_pb2.JwtEcdsaKeyFormat)
KeyProto.add_key_type(jwt_ecdsa_pb2.JwtEcdsaPublicKey,
                      jwt_ecdsa_pb2.JwtEcdsaKeyFormat)
KeyProto.add_key_type(jwt_hmac_pb2.JwtHmacKey, jwt_hmac_pb2.JwtHmacKeyFormat)
KeyProto.add_key_type(jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1PrivateKey,
                      jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1KeyFormat)
KeyProto.add_key_type(jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1PublicKey,
                      jwt_rsa_ssa_pkcs1_pb2.JwtRsaSsaPkcs1KeyFormat)
KeyProto.add_key_type(jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssPrivateKey,
                      jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssKeyFormat)
KeyProto.add_key_type(jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssPublicKey,
                      jwt_rsa_ssa_pss_pb2.JwtRsaSsaPssKeyFormat)
KeyProto.add_key_type(kms_aead_pb2.KmsAeadKey, kms_aead_pb2.KmsAeadKeyFormat)
KeyProto.add_key_type(kms_envelope_pb2.KmsEnvelopeAeadKey,
                      kms_envelope_pb2.KmsEnvelopeAeadKeyFormat)


def _text_format_field(value: Any,
                       field: descriptor.FieldDescriptor,
                       indent: str, remove_value: bool) -> str:
  """Returns a text formated proto field."""
  if field.type == TYPE_MESSAGE:
    output = [
        indent + field.name + ' {',
        _text_format_message(value, indent + '  ', remove_value), indent + '}'
    ]
    return '\n'.join(output)
  elif field.type == TYPE_ENUM:
    value_name = field.enum_type.values_by_number[value].name
    return indent + field.name + ': ' + value_name
  elif field.type in [TYPE_STRING, TYPE_BYTES]:
    return (indent + field.name + ': "' + text_encoding.CEscape(value, False) +
            '"')
  else:
    return indent + field.name + ': ' + str(value)


def _text_format_message(msg: message.Message, indent: str,
                         remove_value: bool) -> str:
  """Returns a text formated proto message.

  Args:
    msg: the proto to be formated.
    indent: the indentation prefix of each line in the output.
    remove_value: if True, replaced the value fields of tink's custom any protos
      with '<removed>'. This is useful to compare protos, but should not be used
      otherwise.

  Returns:
    A proto text format output, where serialized fields are deserialized in
    a comment.
  """
  output = []
  fields = msg.DESCRIPTOR.fields
  if (len(fields) >= 2 and fields[0].name == 'type_url' and
      fields[1].name == 'value'):
    # special case for custom 'any' proto.
    if getattr(msg, 'type_url'):
      type_url = getattr(msg, 'type_url')
      output.append(
          _text_format_field(type_url, fields[0], indent, remove_value))
      if getattr(msg, 'value'):
        value = getattr(msg, 'value')
        if msg.DESCRIPTOR.full_name == 'google.crypto.tink.KeyTemplate':
          # In KeyTemplates, type_url does not match the proto type used.
          proto_type = KeyProto.format_from_url(type_url)
        else:
          proto_type = KeyProto.from_url(type_url)
        # parse 'value' and text format the content in a comment.
        field_proto = proto_type.FromString(value)
        output.append(indent + '# value: [' + TYPE_PREFIX +
                      proto_type.DESCRIPTOR.full_name + '] {')
        output.append(
            _text_format_message(field_proto, indent + '#   ', remove_value))
        output.append(indent + '# }')
        if remove_value:
          output.append(
              _text_format_field('<removed>', fields[1], indent, remove_value))
        else:
          output.append(
              _text_format_field(value, fields[1], indent, remove_value))
    fields = fields[2:]
  for field in fields:
    if field.label == LABEL_REPEATED:
      for value in getattr(msg, field.name):
        output.append(_text_format_field(value, field, indent, remove_value))
    else:
      output.append(
          _text_format_field(
              getattr(msg, field.name), field, indent, remove_value))
  return '\n'.join(output)


def text_format(msg: message.Message) -> str:
  return _text_format_message(msg, '', False)


def assert_tink_proto_equal(self,
                            a: message.Message,
                            b: message.Message,
                            msg: Optional[str] = None) -> None:
  """Fails with a useful error if a and b aren't equal."""
  self.assertMultiLineEqual(
      _text_format_message(a, '', True),
      _text_format_message(b, '', True),
      msg=msg)