# Copyright 2019 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.

"""Pre-generated KeyTemplate for Aead.

One can use these templates to generate a new tink_pb2.Keyset with
tink_pb2.KeysetHandle. To generate a new keyset that contains a single
aes_eax_pb2.AesEaxKey, one can do:
handle = keyset_handle.KeysetHandle(aead_key_templates.AES128_EAX).
"""

import warnings

from tink.proto import aes_ctr_hmac_aead_pb2
from tink.proto import aes_eax_pb2
from tink.proto import aes_gcm_pb2
from tink.proto import aes_gcm_siv_pb2
from tink.proto import common_pb2
from tink.proto import kms_aead_pb2
from tink.proto import kms_envelope_pb2
from tink.proto import tink_pb2
from tink.proto import x_aes_gcm_pb2


_AES_EAX_KEY_TYPE_URL = 'type.googleapis.com/google.crypto.tink.AesEaxKey'
_AES_GCM_KEY_TYPE_URL = 'type.googleapis.com/google.crypto.tink.AesGcmKey'
_XAES_GCM_KEY_TYPE_URL = 'type.googleapis.com/google.crypto.tink.XAesGcmKey'
_AES_GCM_SIV_KEY_TYPE_URL = (
    'type.googleapis.com/google.crypto.tink.AesGcmSivKey')
_AES_CTR_HMAC_AEAD_KEY_TYPE_URL = (
    'type.googleapis.com/google.crypto.tink.AesCtrHmacAeadKey')
_XCHACHA20_POLY1305_KEY_TYPE_URL = (
    'type.googleapis.com/google.crypto.tink.XChaCha20Poly1305Key')
_KMS_AEAD_KEY_TYPE_URL = (
    'type.googleapis.com/google.crypto.tink.KmsAeadKey')
_KMS_ENVELOPE_AEAD_KEY_TYPE_URL = (
    'type.googleapis.com/google.crypto.tink.KmsEnvelopeAeadKey')


def _create_aes_eax_key_template(
    key_size: int,
    iv_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  """Creates an AES EAX KeyTemplate, and fills in its values."""
  key_format = aes_eax_pb2.AesEaxKeyFormat()
  key_format.params.iv_size = iv_size
  key_format.key_size = key_size
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_AES_EAX_KEY_TYPE_URL,
      output_prefix_type=output_prefix_type,
  )
  return key_template


def _create_aes_gcm_key_template(
    key_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  """Creates an AES GCM KeyTemplate, and fills in its values."""
  key_format = aes_gcm_pb2.AesGcmKeyFormat(key_size=key_size)
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_AES_GCM_KEY_TYPE_URL,
      output_prefix_type=output_prefix_type,
  )
  return key_template


def _create_xaes_gcm_key_template(
    salt_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK,
) -> tink_pb2.KeyTemplate:
  """Creates an XAES GCM KeyTemplate, and fills in its values."""
  key_format = x_aes_gcm_pb2.XAesGcmKeyFormat(
      params=x_aes_gcm_pb2.XAesGcmParams(salt_size=salt_size)
  )
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_XAES_GCM_KEY_TYPE_URL,
      output_prefix_type=output_prefix_type,
  )
  return key_template


def _create_aes_gcm_siv_key_template(
    key_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  """Creates an AES GCM SIV KeyTemplate, and fills in its values."""
  key_format = aes_gcm_siv_pb2.AesGcmSivKeyFormat(
      key_size=key_size,
  )
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_AES_GCM_SIV_KEY_TYPE_URL,
      output_prefix_type=output_prefix_type,
  )
  return key_template


def _create_aes_ctr_hmac_aead_key_template(
    aes_key_size: int,
    iv_size: int,
    hmac_key_size: int,
    tag_size: int,
    hash_type: common_pb2.HashType,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  """Creates an AES CTR HMAC AEAD KeyTemplate, and fills in its values."""
  key_format = aes_ctr_hmac_aead_pb2.AesCtrHmacAeadKeyFormat()
  key_format.aes_ctr_key_format.params.iv_size = iv_size
  key_format.aes_ctr_key_format.key_size = aes_key_size
  key_format.hmac_key_format.params.hash = hash_type
  key_format.hmac_key_format.params.tag_size = tag_size
  key_format.hmac_key_format.key_size = hmac_key_size
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_AES_CTR_HMAC_AEAD_KEY_TYPE_URL,
      output_prefix_type=output_prefix_type,
  )
  return key_template


def create_kms_aead_key_template(key_uri: str) -> tink_pb2.KeyTemplate:
  """Creates a KMS AEAD KeyTemplate from a key URI.

  Keys generated by this key template uses RAW output prefix to make them
  compatible with the remote KMS' encrypt/decrypt operations. Unlike other
  templates, when you generate new keys with this template, Tink does not
  generate new key material, but only creates a reference to the remote KEK.

  It requires that a KmsClient that can handle key_uri is registered. Avoid
  registering it more than once.

  It is often not necessary to use this function. Instead of registering a
  KmsClient, and creating an Aead using tink.new_keyset_handle(
  aead.create_kms_aead_key_template(key_uri)).primitive(aead.Aead), you can
  create the Aead directly using kmsClient.getAead(kekUri), without registering
  the kmsClient.

  Args:
    key_uri: Text. The remote key URI.

  Returns:
    A KMS Aead KeyTemplate.
  """
  key_format = kms_aead_pb2.KmsAeadKeyFormat(key_uri=key_uri)
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_KMS_AEAD_KEY_TYPE_URL,
      output_prefix_type=tink_pb2.RAW,
  )
  return key_template


def create_kms_envelope_aead_key_template(
    kek_uri: str, dek_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyTemplate:
  """Creates a KMS Envelope AEAD key template from a KEK URI and a DEK template.

  Keys generated by this key template uses RAW output prefix to make them
  compatible with the remote KMS' encrypt/decrypt operations. Unlike other
  templates, when you generate new keys with this template, Tink does not
  generate new key material, but only creates a reference to the remote KEK.

  It requires that a KmsClient that can handle key_uri is registered. Avoid
  registering it more than once.

  It is often not necessary to use this function. Instead of registering a
  KmsClient, and creating an AEAD using tink.new_keyset_handle(
  aead.create_kms_envelope_aead_key_template(key_uri, dek_template)
  ).primitive(aead.Aead), you can create the envelope Aead directly using
  aead.KmsEnvelopeAead(dek_template, kmsClient.getAead(kekUri)),
  without registering the kmsClient.

  Args:
      kek_uri: Text. The URI of the KEK that resides in an external KMS.
      dek_template: tink_pb2.KeyTemplate. The template of the DEK.

  Returns:
      the resulting key template
  """
  key_format = kms_envelope_pb2.KmsEnvelopeAeadKeyFormat(
      kek_uri=kek_uri, dek_template=dek_template
  )
  key_template = tink_pb2.KeyTemplate(
      value=key_format.SerializeToString(),
      type_url=_KMS_ENVELOPE_AEAD_KEY_TYPE_URL,
      output_prefix_type=tink_pb2.RAW,
  )
  return key_template


AES128_EAX = _create_aes_eax_key_template(key_size=16, iv_size=16)
AES128_EAX_RAW = _create_aes_eax_key_template(
    key_size=16, iv_size=16, output_prefix_type=tink_pb2.RAW)
AES256_EAX = _create_aes_eax_key_template(key_size=32, iv_size=16)
AES256_EAX_RAW = _create_aes_eax_key_template(
    key_size=32, iv_size=16, output_prefix_type=tink_pb2.RAW)
AES128_GCM = _create_aes_gcm_key_template(key_size=16)
AES128_GCM_RAW = _create_aes_gcm_key_template(
    key_size=16, output_prefix_type=tink_pb2.RAW)
AES256_GCM = _create_aes_gcm_key_template(key_size=32)
AES256_GCM_RAW = _create_aes_gcm_key_template(
    key_size=32, output_prefix_type=tink_pb2.RAW)
AES128_GCM_SIV = _create_aes_gcm_siv_key_template(key_size=16)
AES128_GCM_SIV_RAW = _create_aes_gcm_siv_key_template(
    key_size=16, output_prefix_type=tink_pb2.RAW)
AES256_GCM_SIV = _create_aes_gcm_siv_key_template(key_size=32)
AES256_GCM_SIV_RAW = _create_aes_gcm_siv_key_template(
    key_size=32, output_prefix_type=tink_pb2.RAW)
AES128_CTR_HMAC_SHA256 = _create_aes_ctr_hmac_aead_key_template(
    aes_key_size=16,
    iv_size=16,
    hmac_key_size=32,
    tag_size=16,
    hash_type=common_pb2.SHA256)
AES128_CTR_HMAC_SHA256_RAW = _create_aes_ctr_hmac_aead_key_template(
    aes_key_size=16,
    iv_size=16,
    hmac_key_size=32,
    tag_size=16,
    hash_type=common_pb2.SHA256,
    output_prefix_type=tink_pb2.RAW)
AES256_CTR_HMAC_SHA256 = _create_aes_ctr_hmac_aead_key_template(
    aes_key_size=32,
    iv_size=16,
    hmac_key_size=32,
    tag_size=32,
    hash_type=common_pb2.SHA256)
AES256_CTR_HMAC_SHA256_RAW = _create_aes_ctr_hmac_aead_key_template(
    aes_key_size=32,
    iv_size=16,
    hmac_key_size=32,
    tag_size=32,
    hash_type=common_pb2.SHA256,
    output_prefix_type=tink_pb2.RAW)
XCHACHA20_POLY1305 = tink_pb2.KeyTemplate(
    type_url=_XCHACHA20_POLY1305_KEY_TYPE_URL,
    output_prefix_type=tink_pb2.TINK)
XCHACHA20_POLY1305_RAW = tink_pb2.KeyTemplate(
    type_url=_XCHACHA20_POLY1305_KEY_TYPE_URL, output_prefix_type=tink_pb2.RAW)
XAES_256_GCM_192_BIT_NONCE = _create_xaes_gcm_key_template(
    salt_size=12, output_prefix_type=tink_pb2.OutputPrefixType.TINK
)
XAES_256_GCM_192_BIT_NONCE_NO_PREFIX = _create_xaes_gcm_key_template(
    salt_size=12, output_prefix_type=tink_pb2.OutputPrefixType.RAW
)
XAES_256_GCM_160_BIT_NONCE_NO_PREFIX = _create_xaes_gcm_key_template(
    salt_size=8, output_prefix_type=tink_pb2.OutputPrefixType.RAW
)


# Deprecated. Use the predefined constant templates above instead.
def create_aes_eax_key_template(
    key_size: int,
    iv_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  warnings.warn('The "create_aes_eax_key_template" function is deprecated.',
                DeprecationWarning, 2)
  return _create_aes_eax_key_template(key_size, iv_size, output_prefix_type)


# Deprecated. Use the predefined constant templates above instead.
def create_aes_gcm_key_template(
    key_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  warnings.warn('The "create_aes_gcm_key_template" function is deprecated.',
                DeprecationWarning, 2)
  return _create_aes_gcm_key_template(key_size, output_prefix_type)


# Deprecated. Use the predefined constant templates above instead.
def create_aes_gcm_siv_key_template(
    key_size: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  warnings.warn('The "create_aes_gcm_siv_key_template" function is deprecated.',
                DeprecationWarning, 2)
  return _create_aes_gcm_siv_key_template(key_size, output_prefix_type)


# Deprecated. Use the predefined constant templates above instead.
def create_aes_ctr_hmac_aead_key_template(
    aes_key_size: int,
    iv_size: int,
    hmac_key_size: int,
    tag_size: int,
    hash_type: common_pb2.HashType,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK
) -> tink_pb2.KeyTemplate:
  """Creates AesCtrHmacAeadKey template."""
  warnings.warn(
      'The "create_aes_ctr_hmac_aead_key_template" function is deprecated.',
      DeprecationWarning, 2)
  return _create_aes_ctr_hmac_aead_key_template(aes_key_size, iv_size,
                                                hmac_key_size, tag_size,
                                                hash_type, output_prefix_type)
