# 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.core.primitive_set."""

from absl.testing import absltest
from tink.proto import tink_pb2
from tink import aead
from tink import core
from tink import mac
from tink.testing import helper


MAC_TEMPLATE = mac.mac_key_templates.HMAC_SHA256_128BITTAG


def setUpModule():
  mac.register()


def new_key(
    template: tink_pb2.KeyTemplate,
    key_id: int,
    output_prefix_type: tink_pb2.OutputPrefixType = tink_pb2.TINK,
    status: tink_pb2.KeyStatusType = tink_pb2.ENABLED
) -> tink_pb2.Keyset.Key:
  return tink_pb2.Keyset.Key(
      key_data=core.Registry.new_key_data(template),
      key_id=key_id,
      status=status,
      output_prefix_type=output_prefix_type)


class PrimitiveSetTest(absltest.TestCase):

  def test_primitive_returns_entry(self):
    key = new_key(MAC_TEMPLATE, key_id=1)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    primitive_set = core.new_primitive_set(mac.Mac)
    primitive_set.add_primitive(primitive, key)
    entries = primitive_set.primitive(key)
    self.assertLen(entries, 1)
    entry = entries[0]
    self.assertEqual(primitive, entry.primitive)
    self.assertEqual(tink_pb2.ENABLED, entry.status)
    self.assertEqual(core.crypto_format.output_prefix(key), entry.identifier)
    self.assertEqual(1, entry.key_id)

  def test_unknown_key_returns_empty_list(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    unknown_key = new_key(MAC_TEMPLATE, key_id=1)
    self.assertEqual(primitive_set.primitive(unknown_key), [])

  def test_primitive_from_identifier_returns_entry(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key = new_key(MAC_TEMPLATE, key_id=1)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    primitive_set.add_primitive(primitive, key)

    ident = core.crypto_format.output_prefix(key)
    entries = primitive_set.primitive_from_identifier(ident)
    self.assertLen(entries, 1)
    entry = entries[0]
    self.assertEqual(primitive, entry.primitive)
    self.assertEqual(tink_pb2.ENABLED, entry.status)
    self.assertEqual(ident, entry.identifier)
    self.assertEqual(1, entry.key_id)

  def test_list_of_entries_can_be_modified(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key = new_key(MAC_TEMPLATE, key_id=1)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    primitive_set.add_primitive(primitive, key)
    entries = primitive_set.primitive(key)
    entries.append('Something')
    self.assertLen(primitive_set.primitive(key), 1)

  def test_primary_returns_primary(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key = new_key(MAC_TEMPLATE, key_id=1)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    entry = primitive_set.add_primitive(primitive, key)
    primitive_set.set_primary(entry)

    entry = primitive_set.primary()
    self.assertEqual(primitive, entry.primitive)
    self.assertEqual(tink_pb2.ENABLED, entry.status)
    self.assertEqual(core.crypto_format.output_prefix(key), entry.identifier)
    self.assertEqual(1, entry.key_id)

  def test_primary_not_set_primary_fails(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key = new_key(MAC_TEMPLATE, key_id=1)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    primitive_set.add_primitive(primitive, key)
    with self.assertRaises(core.TinkError):
      primitive_set.primary()

  def test_same_key_id_and_prefix_type(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key1 = new_key(MAC_TEMPLATE, key_id=1)
    primitive1 = core.Registry.primitive(key1.key_data, mac.Mac)
    primitive_set.add_primitive(primitive1, key1)
    key2 = new_key(MAC_TEMPLATE, key_id=1, status=tink_pb2.DISABLED)
    primitive2 = core.Registry.primitive(key2.key_data, mac.Mac)
    primitive_set.add_primitive(primitive2, key2)

    expected_ident = core.crypto_format.output_prefix(key1)
    entries = primitive_set.primitive(key1)
    self.assertLen(entries, 2)
    self.assertEqual(primitive1, entries[0].primitive)
    self.assertEqual(primitive2, entries[1].primitive)
    self.assertEqual(tink_pb2.ENABLED, entries[0].status)
    self.assertEqual(tink_pb2.DISABLED, entries[1].status)
    self.assertEqual(expected_ident, entries[0].identifier)
    self.assertEqual(expected_ident, entries[1].identifier)
    self.assertEqual(1, entries[0].key_id)
    self.assertEqual(1, entries[1].key_id)
    self.assertLen(primitive_set.primitive(key2), 2)

  def test_same_key_id_but_different_prefix_type(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key1 = new_key(MAC_TEMPLATE, key_id=1, output_prefix_type=tink_pb2.TINK)
    primitive1 = core.Registry.primitive(key1.key_data, mac.Mac)
    primitive_set.add_primitive(primitive1, key1)
    key2 = new_key(MAC_TEMPLATE, key_id=1, output_prefix_type=tink_pb2.LEGACY)
    primitive2 = core.Registry.primitive(key2.key_data, mac.Mac)
    primitive_set.add_primitive(primitive2, key2)

    entries1 = primitive_set.primitive(key1)
    self.assertLen(entries1, 1)
    self.assertEqual(primitive1, entries1[0].primitive)
    self.assertEqual(tink_pb2.ENABLED, entries1[0].status)
    self.assertEqual(core.crypto_format.output_prefix(key1),
                     entries1[0].identifier)
    self.assertEqual(1, entries1[0].key_id)

    entries2 = primitive_set.primitive(key2)
    self.assertLen(entries2, 1)
    self.assertEqual(primitive2, entries2[0].primitive)
    self.assertEqual(tink_pb2.ENABLED, entries2[0].status)
    self.assertEqual(core.crypto_format.output_prefix(key2),
                     entries2[0].identifier)
    self.assertEqual(1, entries2[0].key_id)

  def test_add_invalid_key_fails(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key = new_key(MAC_TEMPLATE, key_id=1)
    key.ClearField('output_prefix_type')
    with self.assertRaises(core.TinkError):
      primitive_set.add_primitive(helper.FakeMac(), key)

  def test_add_wrong_primitive_fails(self):
    primitive_set = core.new_primitive_set(aead.Aead)
    key = new_key(MAC_TEMPLATE, key_id=1, output_prefix_type=tink_pb2.TINK)
    primitive = core.Registry.primitive(key.key_data, mac.Mac)
    with self.assertRaises(core.TinkError):
      primitive_set.add_primitive(primitive, key)

  def test_primitive_class(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    self.assertEqual(primitive_set.primitive_class(), mac.Mac)

  def test_raw_primitives(self):
    primitive_set = core.new_primitive_set(mac.Mac)
    key1 = new_key(MAC_TEMPLATE, key_id=1, output_prefix_type=tink_pb2.TINK)
    primitive1 = core.Registry.primitive(key1.key_data, mac.Mac)
    primitive_set.add_primitive(primitive1, key1)
    key2 = new_key(MAC_TEMPLATE, key_id=1, output_prefix_type=tink_pb2.RAW)
    primitive2 = core.Registry.primitive(key2.key_data, mac.Mac)
    primitive_set.add_primitive(primitive2, key2)
    key3 = new_key(
        MAC_TEMPLATE,
        key_id=3,
        output_prefix_type=tink_pb2.RAW,
        status=tink_pb2.DISABLED)
    primitive3 = core.Registry.primitive(key3.key_data, mac.Mac)
    primitive_set.add_primitive(primitive3, key3)

    entries = primitive_set.raw_primitives()
    self.assertLen(entries, 2)
    self.assertEqual(primitive2, entries[0].primitive)
    self.assertEqual(tink_pb2.ENABLED, entries[0].status)
    self.assertEqual(core.crypto_format.RAW_PREFIX,
                     entries[0].identifier)
    self.assertEqual(1, entries[0].key_id)
    self.assertEqual(primitive3, entries[1].primitive)
    self.assertEqual(tink_pb2.DISABLED, entries[1].status)
    self.assertEqual(core.crypto_format.RAW_PREFIX,
                     entries[1].identifier)
    self.assertEqual(3, entries[1].key_id)

  def test_all_primitives(self):
    primitive_set = core.new_primitive_set(mac.Mac)

    key0 = new_key(MAC_TEMPLATE, key_id=88, output_prefix_type=tink_pb2.TINK)
    primitive0 = core.Registry.primitive(key0.key_data, mac.Mac)
    primitive_set.add_primitive(primitive0, key0)

    key1 = new_key(MAC_TEMPLATE, key_id=88, output_prefix_type=tink_pb2.LEGACY)
    primitive1 = core.Registry.primitive(key1.key_data, mac.Mac)
    primitive_set.add_primitive(primitive1, key1)

    key2 = new_key(MAC_TEMPLATE, key_id=88, output_prefix_type=tink_pb2.RAW)
    primitive2 = core.Registry.primitive(key2.key_data, mac.Mac)
    primitive_set.add_primitive(primitive2, key2)

    key3 = new_key(
        MAC_TEMPLATE,
        key_id=89,
        output_prefix_type=tink_pb2.RAW,
        status=tink_pb2.DISABLED)
    primitive3 = core.Registry.primitive(key3.key_data, mac.Mac)
    primitive_set.add_primitive(primitive3, key3)

    key4 = new_key(MAC_TEMPLATE, key_id=88, output_prefix_type=tink_pb2.TINK)
    primitive0 = core.Registry.primitive(key4.key_data, mac.Mac)
    primitive_set.add_primitive(primitive0, key4)

    list_of_entries = primitive_set.all()

    v = []
    for entries in list_of_entries:
      v.append(
          sorted([
              (e.identifier, e.output_prefix_type, e.key_id) for e in entries
          ]))
    self.assertCountEqual(v, [
        [(b'', tink_pb2.RAW, 88), (b'', tink_pb2.RAW, 89)],
        [(b'\x01\x00\x00\x00X', tink_pb2.TINK, 88),
         (b'\x01\x00\x00\x00X', tink_pb2.TINK, 88)],
        [(b'\x00\x00\x00\x00X', tink_pb2.LEGACY, 88)],
    ])

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