# 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.
"""MAC wrapper."""

from typing import Optional, Type

from tink.proto import tink_pb2
from tink import _monitoring
from tink import core
from tink.mac import _mac


class _WrappedMac(_mac.Mac):
  """Implements Mac for a set of Mac primitives."""

  def __init__(
      self,
      pset: core.PrimitiveSet,
      compute_monitor: Optional[_monitoring.KeyUsageMonitor] = None,
      verify_monitor: Optional[_monitoring.KeyUsageMonitor] = None,
  ):
    self._primitive_set = pset
    self._compute_monitor = compute_monitor
    self._verify_monitor = verify_monitor

  def compute_mac(self, data: bytes) -> bytes:
    primary = self._primitive_set.primary()
    if primary.output_prefix_type == tink_pb2.LEGACY:
      result = primary.identifier + primary.primitive.compute_mac(
          data + core.crypto_format.LEGACY_START_BYTE
      )
    else:
      result = primary.identifier + primary.primitive.compute_mac(data)

    if self._compute_monitor:
      self._compute_monitor.log(primary.key_id, len(data))

    return result

  def verify_mac(self, mac_value: bytes, data: bytes) -> None:
    if len(mac_value) <= core.crypto_format.NON_RAW_PREFIX_SIZE:
      if self._verify_monitor:
        self._verify_monitor.log_failure()
      # This also rejects raw MAC with size of 4 bytes or fewer. Those MACs are
      # clearly insecure, thus should be discouraged.
      raise core.TinkError('tag too short')
    prefix = mac_value[: core.crypto_format.NON_RAW_PREFIX_SIZE]
    mac_no_prefix = mac_value[core.crypto_format.NON_RAW_PREFIX_SIZE :]
    for entry in self._primitive_set.primitive_from_identifier(prefix):
      try:
        if entry.output_prefix_type == tink_pb2.LEGACY:
          entry.primitive.verify_mac(mac_no_prefix, data + b'\x00')
        else:
          entry.primitive.verify_mac(mac_no_prefix, data)

        # If there is no exception, the MAC is valid and we can return.
        if self._verify_monitor:
          self._verify_monitor.log(entry.key_id, len(mac_no_prefix))
        return
      except core.TinkError:
        pass

    # No 'non-raw' key matched, so let's try the raw keys (if any exist).
    for entry in self._primitive_set.raw_primitives():
      try:
        entry.primitive.verify_mac(mac_value, data)

        # If there is no exception, the MAC is valid and we can return.
        if self._verify_monitor:
          self._verify_monitor.log(entry.key_id, len(mac_value))
        return
      except core.TinkError:
        pass

    if self._verify_monitor:
      self._verify_monitor.log_failure()

    raise core.TinkError('invalid MAC')


class MacWrapper(core.PrimitiveWrapper[_mac.Mac, _mac.Mac]):
  """MacWrapper is the implementation of PrimitiveWrapper for the Mac primitive.

  The returned primitive works with a keyset (rather than a single key). To
  compute a MAC tag, it uses the primary key in the keyset, and prepends to the
  tag a certain prefix associated with the primary key. To verify a tag, the
  primitive uses the prefix of the tag to efficiently select the right key in
  the set. If the keys associated with the prefix do not validate the tag, the
  primitive tries all keys with tink_pb2.OutputPrefixType = tink_pb2.RAW.
  """

  def wrap(self, pset: core.PrimitiveSet) -> _mac.Mac:
    return _WrappedMac(pset)

  def primitive_class(self) -> Type[_mac.Mac]:
    return _mac.Mac

  def input_primitive_class(self) -> Type[_mac.Mac]:
    return _mac.Mac

  def _wrap_with_monitoring_info(
      self,
      pset: core.PrimitiveSet,
      monitoring_keyset_info: _monitoring.MonitoringKeySetInfo,
  ) -> _mac.Mac:
    compute_key_usage_monitor = _monitoring.get_key_usage_monitor_or_none(
        _monitoring.MonitoringContext(
            primitive='mac',
            api_function='compute_mac',
            keyset_info=monitoring_keyset_info,
        )
    )
    verify_key_usage_monitor = _monitoring.get_key_usage_monitor_or_none(
        _monitoring.MonitoringContext(
            primitive='mac',
            api_function='verify_mac',
            keyset_info=monitoring_keyset_info,
        )
    )
    return _WrappedMac(
        pset, compute_key_usage_monitor, verify_key_usage_monitor
    )
