# 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.registry."""
from absl.testing import absltest

from tink.proto import tink_pb2

from tink.python import aead
from tink.python import core
from tink.python import mac
from tink.python.testing import helper


class DummyKeyManager(core.KeyManager):

  def __init__(self, type_url, primitive_class=aead.Aead):
    self._type_url = type_url
    self._primitive_class = primitive_class

  def primitive_class(self):
    return self._primitive_class

  def primitive(self, key_data):
    return helper.FakeAead()

  def key_type(self):
    return self._type_url

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


class DummyPrivateKeyManager(core.PrivateKeyManager):

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

  def primitive_class(self):
    return None

  def primitive(self, key_data):
    return None

  def key_type(self):
    return self._type_url

  def new_key_data(self, key_template):
    return None

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


class DummyMacWrapper(core.PrimitiveWrapper):

  def wrap(self, _):
    return helper.FakeMac()

  def primitive_class(self):
    return mac.Mac


class InconsistentWrapper(core.PrimitiveWrapper):

  def wrap(self, _):
    return helper.FakeAead()

  def primitive_class(self):
    return mac.Mac


def _mac_set(mac_list):
  """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(RegistryTest, self).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_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.assertRaisesRegex(core.TinkError,
                                'uses primitive Aead, and not Mac'):
      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.assertRaisesRegex(core.TinkError,
                                'uses primitive FakeAead, and not Aead'):
      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.assertRaisesRegex(core.TinkError,
                                'No manager for type unknown_type_url'):
      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.assertRaisesRegex(core.TinkError,
                                'does not allow for creation of new keys'):
      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.assertRaisesRegex(core.TinkError,
                                'contains a non-private key'):
      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.assertRaisesRegex(core.TinkError,
                                'is not a PrivateKeyManager'):
      self.reg.public_key_data(key_data)

  def test_wrap_success(self):
    self.reg.register_primitive_wrapper(mac.MacWrapper())
    mac1 = helper.FakeMac('FakeMac1')
    mac2 = helper.FakeMac('FakeMac2')
    wrapped_mac = self.reg.wrap(_mac_set([mac1, mac2]))
    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_unknown_primitive(self):
    with self.assertRaisesRegex(
        core.TinkError,
        'No PrimitiveWrapper registered for primitive Mac.'):
      self.reg.wrap(_mac_set([helper.FakeMac()]))

  def test_primitive_wrapper_reset(self):
    self.reg.register_primitive_wrapper(mac.MacWrapper())
    self.reg.reset()
    with self.assertRaisesRegex(
        core.TinkError,
        'No PrimitiveWrapper registered for primitive Mac.'):
      self.reg.wrap(_mac_set([helper.FakeMac()]))

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

  def test_register_different_primitive_wrappers_twice_fails(self):
    self.reg.register_primitive_wrapper(mac.MacWrapper())
    with self.assertRaisesRegex(
        core.TinkError,
        'A wrapper for primitive Mac has already been added.'):
      self.reg.register_primitive_wrapper(DummyMacWrapper())

  def test_register_inconsistent_wrapper_fails(self):
    with self.assertRaisesRegex(
        core.TinkError,
        'Wrapper for primitive Mac generates incompatible primitive'):
      self.reg.register_primitive_wrapper(InconsistentWrapper())


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