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

"""Tests for tink.python.tink.registry."""

from typing import cast, Mapping, Type, TypeVar
from absl.testing import absltest

from tink.proto import tink_pb2

from tink import aead
from tink import core
from tink import hybrid
from tink import mac
from tink import prf
from tink.mac import _mac_wrapper
from tink.prf import _prf_set_wrapper
from tink.testing import helper


P = TypeVar('P')


class DummyKeyManager(core.KeyManager[P]):

  def __init__(self, type_url: str, primitive_class: Type[P] = aead.Aead):
    self._type_url = type_url
    self._primitive_class = primitive_class

  def primitive_class(self) -> Type[P]:
    return self._primitive_class

  def primitive(self, key_data: tink_pb2.KeyData) -> P:
    return helper.FakeAead()

  def key_type(self) -> str:
    return self._type_url

  def new_key_data(self,
                   key_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyData:
    return tink_pb2.KeyData(type_url=key_template.type_url)

  def does_support(self, type_url: str) -> bool:
    return self.key_type() == type_url


class UnsupportedKeyManager(DummyKeyManager[P]):

  def does_support(self, type_url: str) -> bool:
    return False


class DummyPrivateKeyManager(core.PrivateKeyManager[hybrid.HybridDecrypt]):

  def __init__(self, type_url: str):
    self._type_url = type_url

  def primitive_class(self) -> Type[hybrid.HybridDecrypt]:
    return hybrid.HybridDecrypt

  def primitive(self, key_data: tink_pb2.KeyData) -> hybrid.HybridDecrypt:
    return helper.FakeHybridDecrypt()

  def key_type(self) -> str:
    return self._type_url

  def new_key_data(self,
                   key_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyData:
    return tink_pb2.KeyData()

  def public_key_data(
      self, private_key_data: tink_pb2.KeyData) -> tink_pb2.KeyData:
    return tink_pb2.KeyData(type_url='public_' + private_key_data.type_url)


class DummyMacWrapper(core.PrimitiveWrapper[mac.Mac, mac.Mac]):

  def wrap(self, pset: core.PrimitiveSet) -> mac.Mac:
    _ = pset
    return helper.FakeMac()

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

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


class InconsistentWrapper(core.PrimitiveWrapper[mac.Mac, mac.Mac]):

  def wrap(self, pset: core.PrimitiveSet) -> mac.Mac:
    _ = pset
    # returns a primitive of the wrong type
    return cast(mac.Mac, helper.FakeAead())

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

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


class _WrappedPrfSet(prf.PrfSet):
  """Implements PrfSet for a set of Prf primitives."""

  def __init__(self, primitive_set: core.PrimitiveSet):
    self._primitive_set = primitive_set

  def primary_id(self) -> int:
    return self._primitive_set.primary().key_id

  def all(self) -> Mapping[int, prf.Prf]:
    return {
        entry.key_id: entry.primitive
        for entry in self._primitive_set.raw_primitives()
    }

  def primary(self) -> prf.Prf:
    return self._primitive_set.primary().primitive


class PrfToPrfSetWrapper(core.PrimitiveWrapper[prf.Prf, prf.PrfSet]):
  """A PrimitiveWrapper that wraps Prfs into a PrfSet."""

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

  def primitive_class(self) -> Type[prf.PrfSet]:
    return prf.PrfSet

  def input_primitive_class(self) -> Type[prf.Prf]:
    return prf.Prf


class _InvalidPrf(prf.Prf):

  def compute(self, input_data: bytes, output_length: int) -> bytes:
    raise core.TinkError('Invalid Prf')


class PrfToPrfWrapper(core.PrimitiveWrapper[prf.Prf, prf.Prf]):
  """A PrimitiveWrapper that wraps Prfs into a Prf."""

  def wrap(self, primitives_set: core.PrimitiveSet) -> prf.Prf:
    # This is needed because register_primitive_wrapper calls wrap with
    # an empty PrimitiveSet.
    if not primitives_set.raw_primitives():
      return _InvalidPrf()
    return primitives_set.primary().primitive

  def primitive_class(self) -> Type[prf.Prf]:
    return prf.Prf

  def input_primitive_class(self) -> Type[prf.Prf]:
    return prf.Prf


def _mac_set(mac_list) -> core.PrimitiveSet[mac.Mac]:
  """Converts a List of Mac in a PrimitiveSet and sets the last primary."""
  mac_set = core.new_primitive_set(mac.Mac)
  for i, primitive in enumerate(mac_list):
    mac_set.set_primary(
        mac_set.add_primitive(
            primitive,
            helper.fake_key(key_id=i, output_prefix_type=tink_pb2.RAW)))
  return mac_set


class RegistryTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.reg = core.Registry()
    self.reg.reset()

  def test_key_manager_no_exist(self):
    with self.assertRaises(core.TinkError):
      self.reg.key_manager('invalid')

  def test_key_manager_register(self):
    dummy_key_manager = DummyKeyManager('dummy_type_url')
    self.reg.register_key_manager(dummy_key_manager)
    self.assertEqual(dummy_key_manager, self.reg.key_manager('dummy_type_url'))

  def test_key_manager_reset(self):
    dummy_key_manager = DummyKeyManager('dummy_type_url')
    self.reg.register_key_manager(dummy_key_manager)
    self.reg.reset()
    with self.assertRaises(core.TinkError):
      self.reg.key_manager('dummy_type_url')

  def test_register_same_key_manager_twice(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url', aead.Aead))
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url', aead.Aead))

  def test_register_unsupported_key_manager_fails(self):
    with self.assertRaises(core.TinkError):
      self.reg.register_key_manager(UnsupportedKeyManager('unsupported'))

  def test_key_manager_replace_fails(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url', aead.Aead))
    # Replacing the primitive_class for a type_url not allowed.
    with self.assertRaises(core.TinkError):
      self.reg.register_key_manager(
          DummyKeyManager('dummy_type_url', mac.Mac), new_key_allowed=False)

  def test_key_manager_disable_new_key_enable_fails(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url'))
    # Disable new keys.
    self.reg.register_key_manager(
        DummyKeyManager('dummy_type_url'), new_key_allowed=False)
    # Check new keys can't be enabled again.
    with self.assertRaises(core.TinkError):
      self.reg.register_key_manager(
          DummyKeyManager('dummy_type_url'), new_key_allowed=True)

  def test_primitive_ok(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url', aead.Aead))
    primitive = self.reg.primitive(
        tink_pb2.KeyData(type_url='dummy_type_url'), aead.Aead)
    self.assertIsInstance(primitive, helper.FakeAead)

  def test_primitive_fails_on_wrong_primitive(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url', aead.Aead))
    with self.assertRaises(core.TinkError):
      self.reg.primitive(tink_pb2.KeyData(type_url='dummy_type_url'), mac.Mac)

  def test_primitive_fails_on_subclass(self):
    self.reg.register_key_manager(
        DummyKeyManager('dummy_type_url', helper.FakeAead))
    with self.assertRaises(core.TinkError):
      self.reg.primitive(tink_pb2.KeyData(type_url='dummy_type_url'), aead.Aead)

  def test_new_key_data_success(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url'))
    key_template = tink_pb2.KeyTemplate(type_url='dummy_type_url')
    key_data = self.reg.new_key_data(key_template)
    self.assertEqual(key_data.type_url, 'dummy_type_url')

  def test_new_key_data_wrong_type_url(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url'))
    unknown_key_template = tink_pb2.KeyTemplate(type_url='unknown_type_url')
    with self.assertRaises(core.TinkError):
      self.reg.new_key_data(unknown_key_template)

  def test_new_key_data_no_new_key_allowed(self):
    self.reg.register_key_manager(
        DummyKeyManager('dummy_type_url'), new_key_allowed=False)
    key_template = tink_pb2.KeyTemplate(type_url='dummy_type_url')
    with self.assertRaises(core.TinkError):
      self.reg.new_key_data(key_template)

  def test_public_key_data_success(self):
    self.reg.register_key_manager(DummyPrivateKeyManager('dummy_type_url'))
    key_data = tink_pb2.KeyData(
        type_url='dummy_type_url',
        key_material_type=tink_pb2.KeyData.ASYMMETRIC_PRIVATE)
    public_key_data = self.reg.public_key_data(key_data)
    self.assertEqual(public_key_data.type_url, 'public_dummy_type_url')

  def test_public_key_data_fails_for_non_asymmetric_private_key(self):
    self.reg.register_key_manager(DummyPrivateKeyManager('dummy_type_url'))
    key_data = tink_pb2.KeyData(
        type_url='dummy_type_url',
        key_material_type=tink_pb2.KeyData.ASYMMETRIC_PUBLIC)
    with self.assertRaises(core.TinkError):
      self.reg.public_key_data(key_data)

  def test_public_key_data_fails_for_non_private_key_manager(self):
    self.reg.register_key_manager(DummyKeyManager('dummy_type_url'))
    key_data = tink_pb2.KeyData(
        type_url='dummy_type_url',
        key_material_type=tink_pb2.KeyData.ASYMMETRIC_PRIVATE)
    with self.assertRaises(core.TinkError):
      self.reg.public_key_data(key_data)

  def test_wrap_success(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    mac1 = helper.FakeMac('FakeMac1')
    mac2 = helper.FakeMac('FakeMac2')
    wrapped_mac = self.reg.wrap(_mac_set([mac1, mac2]), mac.Mac)
    wrapped_mac.verify_mac(mac1.compute_mac(b'data1'), b'data1')
    wrapped_mac.verify_mac(mac2.compute_mac(b'data2'), b'data2')
    wrapped_mac.verify_mac(wrapped_mac.compute_mac(b'data'), b'data')

  def test_wrap_with_primitive_class_success(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    mac1 = helper.FakeMac('FakeMac1')
    mac2 = helper.FakeMac('FakeMac2')
    wrapped_mac = self.reg.wrap(_mac_set([mac1, mac2]), mac.Mac)
    wrapped_mac.verify_mac(mac1.compute_mac(b'data1'), b'data1')
    wrapped_mac.verify_mac(mac2.compute_mac(b'data2'), b'data2')
    wrapped_mac.verify_mac(wrapped_mac.compute_mac(b'data'), b'data')

  def test_wrap_with_incompatible_primitive_class_fails(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    pset = core.new_primitive_set(prf.Prf)
    prf1 = helper.FakePrf('FakePrf1')
    pset.set_primary(
        pset.add_primitive(
            prf1,
            helper.fake_key(key_id=1234, output_prefix_type=tink_pb2.RAW)))
    with self.assertRaises(core.TinkError):
      _ = self.reg.wrap(pset, mac.Mac)

  def test_wrap_unknown_primitive(self):
    with self.assertRaises(core.TinkError):
      self.reg.wrap(_mac_set([helper.FakeMac()]), mac.Mac)

  def test_primitive_wrapper_reset(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    self.reg.reset()
    with self.assertRaises(core.TinkError):
      self.reg.wrap(_mac_set([helper.FakeMac()]), mac.Mac)

  def test_register_same_primitive_wrapper_twice(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())

  def test_register_different_primitive_wrappers_twice_fails(self):
    self.reg.register_primitive_wrapper(_mac_wrapper.MacWrapper())
    with self.assertRaises(core.TinkError):
      self.reg.register_primitive_wrapper(DummyMacWrapper())

  def test_register_inconsistent_wrapper_fails(self):
    with self.assertRaises(core.TinkError):
      self.reg.register_primitive_wrapper(InconsistentWrapper())

  def test_register_prf_to_prfset_wrapper_success(self):
    self.reg.register_primitive_wrapper(PrfToPrfSetWrapper())
    pset = core.new_primitive_set(prf.Prf)
    prf1 = helper.FakePrf('FakePrf1')
    prf2 = helper.FakePrf('FakePrf2')
    pset.add_primitive(
        prf1,
        helper.fake_key(key_id=1234, output_prefix_type=tink_pb2.RAW))
    pset.set_primary(
        pset.add_primitive(
            prf2,
            helper.fake_key(key_id=5678, output_prefix_type=tink_pb2.RAW)))
    wrapped_prf_set = self.reg.wrap(pset, prf.PrfSet)

    expected_output1 = prf1.compute(b'input', output_length=31)
    expected_output2 = prf2.compute(b'input', output_length=31)

    self.assertEqual(
        wrapped_prf_set.primary().compute(b'input', output_length=31),
        expected_output2)
    self.assertEqual(wrapped_prf_set.primary_id(), 5678)
    prfs = wrapped_prf_set.all()
    self.assertLen(prfs, 2)
    self.assertEqual(prfs[1234].compute(b'input', output_length=31),
                     expected_output1)
    self.assertEqual(prfs[5678].compute(b'input', output_length=31),
                     expected_output2)

  def test_two_wrappers_with_equal_primitive_class_fails(self):
    self.reg.register_primitive_wrapper(_prf_set_wrapper.PrfSetWrapper())
    # Only one wrapper for a output primitive is allowed.
    with self.assertRaises(core.TinkError):
      self.reg.register_primitive_wrapper(PrfToPrfSetWrapper())

  def test_two_wrappers_with_equal_input_primitive_class_success(self):
    self.reg.register_primitive_wrapper(PrfToPrfSetWrapper())
    # this is allowed, since PrfToPrfWrapper has a different output primitive.
    self.reg.register_primitive_wrapper(PrfToPrfWrapper())
    pset = core.new_primitive_set(prf.Prf)
    prf1 = helper.FakePrf('FakePrf1')
    prf2 = helper.FakePrf('FakePrf2')
    pset.add_primitive(
        prf1,
        helper.fake_key(key_id=1234, output_prefix_type=tink_pb2.RAW))
    pset.set_primary(
        pset.add_primitive(
            prf2,
            helper.fake_key(key_id=5678, output_prefix_type=tink_pb2.RAW)))

    wrapped_prf_set = self.reg.wrap(pset, prf.PrfSet)
    wrapped_prf = self.reg.wrap(pset, prf.Prf)

    expected_output2 = prf2.compute(b'input', output_length=31)

    self.assertEqual(
        wrapped_prf_set.primary().compute(b'input', output_length=31),
        expected_output2)
    self.assertEqual(
        wrapped_prf.compute(b'input', output_length=31),
        expected_output2)

if __name__ == '__main__':
  absltest.main()
