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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import absltest
from tink.proto import tink_pb2
from tink.python import core
from tink.python.hybrid import hybrid_decrypt
from tink.python.hybrid import hybrid_decrypt_wrapper
from tink.python.hybrid import hybrid_encrypt
from tink.python.hybrid import hybrid_encrypt_wrapper
from tink.python.testing import helper


def new_primitives_and_keys(key_id, output_prefix_type):
  fake_dec_key = helper.fake_key(
      key_material_type=tink_pb2.KeyData.ASYMMETRIC_PRIVATE,
      key_id=key_id,
      output_prefix_type=output_prefix_type)
  fake_enc_key = helper.fake_key(
      key_material_type=tink_pb2.KeyData.ASYMMETRIC_PUBLIC,
      key_id=key_id, output_prefix_type=output_prefix_type)
  fake_hybrid_decrypt = helper.FakeHybridDecrypt(
      'fakeHybrid {}'.format(key_id))
  fake_hybrid_encrypt = helper.FakeHybridEncrypt(
      'fakeHybrid {}'.format(key_id))
  return fake_hybrid_decrypt, fake_hybrid_encrypt, fake_dec_key, fake_enc_key


class HybridWrapperTest(absltest.TestCase):

  def test_encrypt_decrypt(self):
    dec, enc, dec_key, enc_key = new_primitives_and_keys(1234, tink_pb2.TINK)
    dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)
    dec_pset.set_primary(dec_pset.add_primitive(dec, dec_key))
    wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(dec_pset)

    enc_pset = core.new_primitive_set(hybrid_encrypt.HybridEncrypt)
    enc_pset.set_primary(enc_pset.add_primitive(enc, enc_key))
    wrapped_enc = hybrid_encrypt_wrapper.HybridEncryptWrapper().wrap(enc_pset)

    ciphertext = wrapped_enc.encrypt(b'plaintext', b'context_info')
    self.assertEqual(
        wrapped_dec.decrypt(ciphertext, b'context_info'), b'plaintext')

  def test_encrypt_decrypt_with_key_rotation(self):
    dec, enc, dec_key, enc_key = new_primitives_and_keys(1234, tink_pb2.TINK)
    enc_pset = core.new_primitive_set(hybrid_encrypt.HybridEncrypt)
    enc_pset.set_primary(enc_pset.add_primitive(enc, enc_key))
    wrapped_enc = hybrid_encrypt_wrapper.HybridEncryptWrapper().wrap(enc_pset)
    ciphertext = wrapped_enc.encrypt(b'plaintext', b'context_info')

    new_dec, new_enc, new_dec_key, new_enc_key = new_primitives_and_keys(
        5678, tink_pb2.TINK)
    new_enc_pset = core.new_primitive_set(hybrid_encrypt.HybridEncrypt)
    new_enc_pset.set_primary(new_enc_pset.add_primitive(new_enc, new_enc_key))
    new_wrapped_enc = hybrid_encrypt_wrapper.HybridEncryptWrapper().wrap(
        new_enc_pset)

    new_dec, new_enc, new_dec_key, new_enc_key = new_primitives_and_keys(
        5678, tink_pb2.TINK)
    new_dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)
    new_dec_pset.add_primitive(dec, dec_key)
    new_dec_pset.set_primary(new_dec_pset.add_primitive(new_dec, new_dec_key))
    new_wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(
        new_dec_pset)

    new_ciphertext = new_wrapped_enc.encrypt(b'new_plaintext',
                                             b'new_context_info')
    self.assertEqual(
        new_wrapped_dec.decrypt(ciphertext, b'context_info'),
        b'plaintext')
    self.assertEqual(
        new_wrapped_dec.decrypt(new_ciphertext, b'new_context_info'),
        b'new_plaintext')

  def test_encrypt_decrypt_with_key_rotation_from_raw(self):
    raw_dec, raw_enc, raw_dec_key, raw_enc_key = new_primitives_and_keys(
        1234, tink_pb2.RAW)
    old_raw_ciphertext = raw_enc.encrypt(b'old_raw_ciphertext', b'context_info')

    new_dec, new_enc, new_dec_key, new_enc_key = new_primitives_and_keys(
        5678, tink_pb2.TINK)
    enc_pset = core.new_primitive_set(hybrid_encrypt.HybridEncrypt)
    enc_pset.add_primitive(raw_enc, raw_enc_key)
    enc_pset.set_primary(enc_pset.add_primitive(new_enc, new_enc_key))
    wrapped_enc = hybrid_encrypt_wrapper.HybridEncryptWrapper().wrap(
        enc_pset)

    dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)
    dec_pset.add_primitive(raw_dec, raw_dec_key)
    dec_pset.set_primary(dec_pset.add_primitive(new_dec, new_dec_key))
    wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(dec_pset)

    new_ciphertext = wrapped_enc.encrypt(b'new_plaintext', b'new_context_info')
    self.assertEqual(
        wrapped_dec.decrypt(old_raw_ciphertext, b'context_info'),
        b'old_raw_ciphertext')
    self.assertEqual(
        wrapped_dec.decrypt(new_ciphertext, b'new_context_info'),
        b'new_plaintext')

  def test_encrypt_decrypt_two_raw_keys(self):
    dec1, enc1, dec1_key, _ = new_primitives_and_keys(
        1234, tink_pb2.RAW)
    raw_ciphertext1 = enc1.encrypt(b'plaintext1', b'context_info1')
    dec2, enc2, dec2_key, _ = new_primitives_and_keys(
        1234, tink_pb2.RAW)
    raw_ciphertext2 = enc2.encrypt(b'plaintext2', b'context_info2')

    dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)
    dec_pset.add_primitive(dec1, dec1_key)
    dec_pset.set_primary(dec_pset.add_primitive(dec2, dec2_key))
    wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(dec_pset)

    self.assertEqual(
        wrapped_dec.decrypt(raw_ciphertext1, b'context_info1'),
        b'plaintext1')
    self.assertEqual(
        wrapped_dec.decrypt(raw_ciphertext2, b'context_info2'),
        b'plaintext2')

  def test_decrypt_unknown_ciphertext_fails(self):
    unknown_enc = helper.FakeHybridEncrypt('unknownHybrid')
    unknown_ciphertext = unknown_enc.encrypt(b'plaintext', b'context_info')

    dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)

    dec1, _, dec1_key, _ = new_primitives_and_keys(1234, tink_pb2.RAW)
    dec2, _, dec2_key, _ = new_primitives_and_keys(5678, tink_pb2.TINK)
    dec_pset.add_primitive(dec1, dec1_key)
    dec_pset.set_primary(dec_pset.add_primitive(dec2, dec2_key))

    wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(dec_pset)

    with self.assertRaisesRegex(core.TinkError, 'Decryption failed'):
      wrapped_dec.decrypt(unknown_ciphertext, b'context_info')

  def test_decrypt_wrong_associated_data_fails(self):
    dec, enc, dec_key, enc_key = new_primitives_and_keys(1234, tink_pb2.TINK)
    dec_pset = core.new_primitive_set(hybrid_decrypt.HybridDecrypt)
    dec_pset.set_primary(dec_pset.add_primitive(dec, dec_key))
    wrapped_dec = hybrid_decrypt_wrapper.HybridDecryptWrapper().wrap(dec_pset)

    enc_pset = core.new_primitive_set(hybrid_encrypt.HybridEncrypt)
    enc_pset.set_primary(enc_pset.add_primitive(enc, enc_key))
    wrapped_enc = hybrid_encrypt_wrapper.HybridEncryptWrapper().wrap(enc_pset)

    ciphertext = wrapped_enc.encrypt(b'plaintext', b'context_info')
    with self.assertRaisesRegex(core.TinkError, 'Decryption failed'):
      wrapped_dec.decrypt(ciphertext, b'wrong_context_info')


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