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

"""Public Key Sign wrapper."""

from typing import Optional, Type

from tink.proto import tink_pb2
from tink import _monitoring
from tink import core
from tink.signature import _public_key_sign
from tink.signature import _public_key_verify


class _WrappedPublicKeySign(_public_key_sign.PublicKeySign):
  """Implements PublicKeySign for a set of PublicKeySign primitives."""

  def __init__(
      self,
      primitives_set: core.PrimitiveSet,
      monitor: Optional[_monitoring.KeyUsageMonitor] = None,
  ):
    self._primitive_set = primitives_set
    self._monitor = monitor

  def sign(self, data: bytes) -> bytes:
    """Computes the signature for data using the primary primitive.

    Args:
      data: The input data.

    Returns:
      The signature.
    """
    primary = self._primitive_set.primary()

    if not primary:
      raise core.TinkError('primary primitive not set')

    sign_data = data
    if primary.output_prefix_type == tink_pb2.LEGACY:
      sign_data = sign_data + b'\x00'

    result = primary.identifier + primary.primitive.sign(sign_data)

    if self._monitor:
      self._monitor.log(primary.key_id, len(sign_data))

    return result


class PublicKeySignWrapper(
    core.PrimitiveWrapper[_public_key_sign.PublicKeySign,
                          _public_key_sign.PublicKeySign]):
  """A PrimitiveWrapper for the PublicKeySign primitive.

  The returned primitive works with a keyset (rather than a single key). To sign
  a message, it uses the primary key in the keyset, and prepends to the
  signature a certain prefix associated with the primary key.
  """

  def wrap(self, primitives_set: core.PrimitiveSet) -> _WrappedPublicKeySign:
    return _WrappedPublicKeySign(primitives_set)

  def primitive_class(self) -> Type[_public_key_sign.PublicKeySign]:
    return _public_key_sign.PublicKeySign

  def input_primitive_class(self) -> Type[_public_key_sign.PublicKeySign]:
    return _public_key_sign.PublicKeySign

  def _wrap_with_monitoring_info(
      self,
      primitives_set: core.PrimitiveSet,
      monitoring_keyset_info: _monitoring.MonitoringKeySetInfo,
  ) -> _public_key_sign.PublicKeySign:
    key_usage_monitor = _monitoring.get_key_usage_monitor_or_none(
        _monitoring.MonitoringContext(
            primitive='public_key_sign',
            api_function='sign',
            keyset_info=monitoring_keyset_info,
        )
    )
    return _WrappedPublicKeySign(primitives_set, key_usage_monitor)


class _WrappedPublicKeyVerify(_public_key_verify.PublicKeyVerify):
  """Implements PublicKeyVerify for a set of PublicKeyVerify primitives."""

  def __init__(
      self,
      primitives_set: core.PrimitiveSet,
      monitor: Optional[_monitoring.KeyUsageMonitor] = None,
  ):
    self._primitive_set = primitives_set
    self._monitor = monitor

  def verify(self, signature: bytes, data: bytes):
    """Verifies that signature is a digital signature for data.

    Args:
      signature: The signature bytes to be checked.
      data: The data bytes to be checked.

    Raises:
      tink_error.TinkError if the verification fails.
    """
    if len(signature) <= core.crypto_format.NON_RAW_PREFIX_SIZE:
      if self._monitor:
        self._monitor.log_failure()
      # This also rejects raw signatures with size of 4 bytes or fewer.
      # We're not aware of any schemes that output signatures that small.
      raise core.TinkError('signature too short')

    key_id = signature[:core.crypto_format.NON_RAW_PREFIX_SIZE]
    raw_sig = signature[core.crypto_format.NON_RAW_PREFIX_SIZE:]

    for entry in self._primitive_set.primitive_from_identifier(key_id):
      try:
        if entry.output_prefix_type == tink_pb2.LEGACY:
          entry.primitive.verify(raw_sig, data + b'\x00')
        else:
          entry.primitive.verify(raw_sig, data)
        # Signature is valid, we can return
        if self._monitor:
          self._monitor.log(entry.key_id, len(signature))
        return
      except core.TinkError:
        pass

    # No matching key succeeded with verification, try all RAW keys
    for entry in self._primitive_set.raw_primitives():
      try:
        entry.primitive.verify(signature, data)
        # Signature is valid, we can return
        if self._monitor:
          self._monitor.log(entry.key_id, len(signature))
        return
      except core.TinkError:
        pass

    # nothing works.
    if self._monitor:
      self._monitor.log_failure()

    raise core.TinkError('invalid signature')


class PublicKeyVerifyWrapper(
    core.PrimitiveWrapper[_public_key_verify.PublicKeyVerify,
                          _public_key_verify.PublicKeyVerify]):
  """WrappedPublicKeyVerify is the PrimitiveWrapper for PublicKeyVerify.

  The returned primitive works with a keyset (rather than a single key). To sign
  a message, it uses the primary key in the keyset, and prepends to the
  signature a certain prefix associated with the primary key.

  The returned primitive works with a keyset (rather than a single key). To
  verify a signature, the primitive uses the prefix of the signature to
  efficiently select the right key in the set. If there is no key associated
  with the prefix or if the keys associated with the prefix do not work, the
  primitive tries all keys with tink_pb2.OutputPrefixType = tink_pb2.RAW.
  """

  def wrap(self, primitives_set: core.PrimitiveSet) -> _WrappedPublicKeyVerify:
    return _WrappedPublicKeyVerify(primitives_set)

  def primitive_class(self) -> Type[_public_key_verify.PublicKeyVerify]:
    return _public_key_verify.PublicKeyVerify

  def input_primitive_class(self) -> Type[_public_key_verify.PublicKeyVerify]:
    return _public_key_verify.PublicKeyVerify

  def _wrap_with_monitoring_info(
      self,
      primitives_set: core.PrimitiveSet,
      monitoring_keyset_info: _monitoring.MonitoringKeySetInfo,
  ) -> _public_key_verify.PublicKeyVerify:
    key_usage_monitor = _monitoring.get_key_usage_monitor_or_none(
        _monitoring.MonitoringContext(
            primitive='public_key_verify',
            api_function='verify',
            keyset_info=monitoring_keyset_info,
        )
    )
    return _WrappedPublicKeyVerify(primitives_set, key_usage_monitor)
