# 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.
"""Tests for tink.python.tink.streaming_aead._streaming_aead_wrapper."""

import io
from typing import BinaryIO, cast

from absl.testing import absltest
from absl.testing import parameterized
from tink.proto import aes_gcm_hkdf_streaming_pb2
from tink.proto import common_pb2
from tink.proto import tink_pb2
import tink
from tink import secret_key_access
from tink import streaming_aead
from tink.testing import bytes_io
from tink.testing import keyset_builder


TEMPLATE = streaming_aead.streaming_aead_key_templates.AES128_GCM_HKDF_4KB
TYPE_URL = 'type.googleapis.com/google.crypto.tink.AesGcmHkdfStreamingKey'


def setUpModule():
  streaming_aead.register()


def _encrypt(primitive: streaming_aead.StreamingAead, plaintext: bytes,
             associated_data: bytes) -> bytes:
  ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
  with primitive.new_encrypting_stream(ciphertext_dest, associated_data) as es:
    es.write(plaintext)
  return ciphertext_dest.value_after_close()


class StreamingAeadWrapperTest(parameterized.TestCase):

  @parameterized.parameters(
      [b'plaintext', b'', b'smile \xf0\x9f\x98\x80', b'\xf0\x9f\x98'])
  def test_encrypt_decrypt_success(self, plaintext):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)

    aad = b'associated_data'
    ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
    with primitive.new_encrypting_stream(ciphertext_dest, aad) as es:
      self.assertLen(plaintext, es.write(plaintext))
    self.assertTrue(ciphertext_dest.closed)

    ciphertext_src = io.BytesIO(ciphertext_dest.value_after_close())
    with primitive.new_decrypting_stream(ciphertext_src, aad) as ds:
      output = ds.read()
    self.assertTrue(ciphertext_src.closed)
    self.assertEqual(output, plaintext)

  def test_long_plaintext_encrypt_decrypt_success(self):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)

    long_plaintext = b' '.join(b'%d' % i for i in range(10 * 1000 * 1000))
    aad = b'associated_data'
    ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
    with primitive.new_encrypting_stream(ciphertext_dest, aad) as es:
      self.assertLen(long_plaintext, es.write(long_plaintext))
    self.assertTrue(ciphertext_dest.closed)

    ciphertext_src = io.BytesIO(ciphertext_dest.value_after_close())
    with primitive.new_decrypting_stream(ciphertext_src, aad) as ds:
      output = ds.read()
    self.assertTrue(ciphertext_src.closed)
    self.assertEqual(output, long_plaintext)

  def test_encrypt_decrypt_with_slow_bytes_io_success(self):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)
    plaintext = b' '.join(b'%d' % i for i in range(10 * 1000))
    aad = b'associated_data'

    # Even if the ciphertext destination only writes small data chunks and
    # sometimes None, calling write() should write the whole plaintext.
    ciphertext_dest = bytes_io.SlowBytesIO()
    with primitive.new_encrypting_stream(ciphertext_dest, aad) as es:
      es.write(plaintext)
    ciphertext = ciphertext_dest.value_after_close()

    # Even if the ciphertext source only returns small data chunks and sometimes
    # None, calling read() should return the whole ciphertext.
    ciphertext_src = cast(BinaryIO, bytes_io.SlowBytesIO(ciphertext))
    with primitive.new_decrypting_stream(ciphertext_src, aad) as ds:
      output = ds.read()
    self.assertTrue(ciphertext_src.closed)
    self.assertEqual(output, plaintext)

  def test_encrypt_decrypt_with_bad_write_bytes_io_success(self):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)
    plaintext = b'plaintext'
    aad = b'associated_data'

    ciphertext_dest = bytes_io.BadWriteBytesIO()
    with self.assertRaises(tink.TinkError):
      es = primitive.new_encrypting_stream(ciphertext_dest, aad)
      es.write(plaintext)
      es.close()

  def test_encrypt_decrypt_with_slow_readable_raw_bytes_success(self):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)
    plaintext = b' '.join(b'%d' % i for i in range(10 * 1000))
    aad = b'associated_data'
    ciphertext = _encrypt(primitive, plaintext, aad)

    # Even if the ciphertext source only returns small data chunks and sometimes
    # None, calling read() should return the whole ciphertext.
    ciphertext_src = cast(BinaryIO, bytes_io.SlowReadableRawBytes(ciphertext))
    with primitive.new_decrypting_stream(ciphertext_src, aad) as ds:
      output = ds.read()
    self.assertTrue(ciphertext_src.closed)
    self.assertEqual(output, plaintext)

  def test_encrypt_decrypt_bad_aad(self):
    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)

    plaintext = b'plaintext'
    aad = b'associated_data'

    ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
    with primitive.new_encrypting_stream(ciphertext_dest, aad) as es:
      self.assertLen(plaintext, es.write(plaintext))
    self.assertTrue(ciphertext_dest.closed)

    ciphertext_src = io.BytesIO(ciphertext_dest.value_after_close())
    with primitive.new_decrypting_stream(ciphertext_src, b'bad aad') as ds:
      with self.assertRaises(tink.TinkError):
        _ = ds.read()

  def test_decrypt_unknown_key_fails(self):
    plaintext = b'plaintext'
    aad = b'associated_data'

    unknown_keyset_handle = tink.new_keyset_handle(TEMPLATE)
    unknown_primitive = unknown_keyset_handle.primitive(
        streaming_aead.StreamingAead)
    unknown_ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
    with unknown_primitive.new_encrypting_stream(unknown_ciphertext_dest,
                                                 aad) as es:
      es.write(plaintext)

    keyset_handle = tink.new_keyset_handle(TEMPLATE)
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)
    ciphertext_src = io.BytesIO(unknown_ciphertext_dest.value_after_close())
    with primitive.new_decrypting_stream(ciphertext_src, aad) as ds:
      with self.assertRaises(tink.TinkError):
        _ = ds.read()

  @parameterized.parameters(
      [io.BytesIO, bytes_io.SlowBytesIO, bytes_io.SlowReadableRawBytes])
  def test_encrypt_decrypt_with_key_rotation(self, input_stream_factory):
    builder = keyset_builder.new_keyset_builder()
    older_key_id = builder.add_new_key(TEMPLATE)
    builder.set_primary_key(older_key_id)
    p1 = builder.keyset_handle().primitive(streaming_aead.StreamingAead)

    newer_key_id = builder.add_new_key(TEMPLATE)
    p2 = builder.keyset_handle().primitive(streaming_aead.StreamingAead)

    builder.set_primary_key(newer_key_id)
    p3 = builder.keyset_handle().primitive(streaming_aead.StreamingAead)

    builder.disable_key(older_key_id)
    p4 = builder.keyset_handle().primitive(streaming_aead.StreamingAead)

    self.assertNotEqual(older_key_id, newer_key_id)

    # p1 encrypts with the older key. So p1, p2 and p3 can decrypt it,
    # but not p4.
    plaintext1 = b' '.join(b'%d' % i for i in range(100 * 101))
    ciphertext1 = _encrypt(p1, plaintext1, b'aad1')
    with p1.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext1)), b'aad1') as ds:
      self.assertEqual(ds.read(), plaintext1)
    with p2.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext1)), b'aad1') as ds:
      self.assertEqual(ds.read(), plaintext1)
    with p3.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext1)), b'aad1') as ds:
      self.assertEqual(ds.read(), plaintext1)
    with p4.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext1)), b'aad1') as ds:
      with self.assertRaises(tink.TinkError):
        ds.read()

    # p2 encrypts with the older key. So p1, p2 and p3 can decrypt it,
    # but not p4.
    plaintext2 = b' '.join(b'%d' % i for i in range(100 * 102))
    ciphertext2 = _encrypt(p2, plaintext2, b'aad2')
    with p1.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext2)), b'aad2') as ds:
      self.assertEqual(ds.read(), plaintext2)
    with p2.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext2)), b'aad2') as ds:
      self.assertEqual(ds.read(), plaintext2)
    with p3.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext2)), b'aad2') as ds:
      self.assertEqual(ds.read(), plaintext2)
    with p4.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext2)), b'aad2') as ds:
      with self.assertRaises(tink.TinkError):
        ds.read()

    # p3 encrypts with the newer key. So p2, p3 and p4 can decrypt it,
    # but not p1.
    plaintext3 = b' '.join(b'%d' % i for i in range(100 * 103))
    ciphertext3 = _encrypt(p3, plaintext3, b'aad3')
    with p1.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext3)), b'aad3') as ds:
      with self.assertRaises(tink.TinkError):
        ds.read()
    with p2.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext3)), b'aad3') as ds:
      self.assertEqual(ds.read(), plaintext3)
    with p3.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext3)), b'aad3') as ds:
      self.assertEqual(ds.read(), plaintext3)
    with p4.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext3)), b'aad3') as ds:
      self.assertEqual(ds.read(), plaintext3)

    # p4 encrypts with the newer key. So p2, p3 and p4 can decrypt it,
    # but not p1.
    plaintext4 = b' '.join(b'%d' % i for i in range(100 * 104))
    ciphertext4 = _encrypt(p4, plaintext4, b'aad4')
    with p1.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext4)), b'aad4') as ds:
      with self.assertRaises(tink.TinkError):
        ds.read()
    with p2.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext4)), b'aad4') as ds:
      self.assertEqual(ds.read(), plaintext4)
    with p3.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext4)), b'aad4') as ds:
      self.assertEqual(ds.read(), plaintext4)
    with p4.new_decrypting_stream(
        cast(BinaryIO, input_stream_factory(ciphertext4)), b'aad4') as ds:
      self.assertEqual(ds.read(), plaintext4)

  def test_decrypt_tink_output_prefix(self):
    key = aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingKey(
        version=0,
        params=aes_gcm_hkdf_streaming_pb2.AesGcmHkdfStreamingParams(
            ciphertext_segment_size=512,
            derived_key_size=16,
            hkdf_hash_type=common_pb2.HashType.SHA256,
        ),
        key_value=b'0123456789abcdef',
    )
    value1 = key.SerializeToString()
    key.key_value = b'ABCDEF0123456789'
    value2 = key.SerializeToString()

    # We use two keys so that we have at least 1 raw key in the keyset: Tink
    # has a check that creating a new StreamingAead fails when the keyset does
    # not contain any raw keys, and we only want this to fail on decryption.
    keyset = tink_pb2.Keyset(
        primary_key_id=1,
        key=[
            tink_pb2.Keyset.Key(
                key_data=tink_pb2.KeyData(
                    type_url=TYPE_URL,
                    value=value1,
                    key_material_type=tink_pb2.KeyData.SYMMETRIC,
                ),
                output_prefix_type=tink_pb2.OutputPrefixType.TINK,
                status=tink_pb2.KeyStatusType.ENABLED,
                key_id=1,
            ),
            tink_pb2.Keyset.Key(
                key_data=tink_pb2.KeyData(
                    type_url=TYPE_URL,
                    value=value2,
                    key_material_type=tink_pb2.KeyData.SYMMETRIC,
                ),
                output_prefix_type=tink_pb2.OutputPrefixType.RAW,
                status=tink_pb2.KeyStatusType.ENABLED,
                key_id=2,
            ),
        ],
    )

    keyset_handle = tink.proto_keyset_format.parse(
        keyset.SerializeToString(), secret_key_access.TOKEN
    )
    primitive = keyset_handle.primitive(streaming_aead.StreamingAead)

    plaintext = b'plaintext'
    associated_data = b'associated_data'

    ciphertext_dest = bytes_io.BytesIOWithValueAfterClose()
    with primitive.new_encrypting_stream(
        ciphertext_dest, associated_data
    ) as es:
      self.assertLen(plaintext, es.write(plaintext))
    self.assertTrue(ciphertext_dest.closed)

    ciphertext_src = io.BytesIO(ciphertext_dest.value_after_close())
    with primitive.new_decrypting_stream(ciphertext_src, associated_data) as ds:
      output = ds.read()
    self.assertTrue(ciphertext_src.closed)
    self.assertEqual(output, plaintext)


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