Implement EncryptingStream.

This file object will be returned by streaming_aead.new_encrypting_stream(...).

PiperOrigin-RevId: 269839414

python/streaming_aead/encrypting_stream.py

+# 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.
+"""A file-like object that encrypts data written to it.
+
+It writes the ciphertext to a given other file-like object, which can later be
+decrypted and read using a DecryptingStream wrapper.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import google_type_annotations
+from __future__ import print_function
+
+import errno
+import io
+from typing import Iterable, BinaryIO
+
+from tink.python.cc.clif import cc_streaming_aead_wrappers
+from tink.python.core import tink_error
+from tink.python.util import file_object_adapter
+
+
+class EncryptingStream(io.BufferedIOBase):
+  """A file-like object which wraps writes to an underlying file-like object.
+
+  It encrypts any data written to it, and writes the ciphertext to the wrapped
+  object.
+
+  The additional method position() returns the number of written plaintext
+  bytes.
+
+  Writes to an EncryptingStream may be partial - it is important to check the
+  return value of write().
+
+  The close() method indicates that the message is complete, and will write a
+  final ciphertext block to signify end of message. The context manager will
+  only call this close() method on a normal exit - if an exception is raised
+  inside the context manager which causes it to exit early, the close() method
+  will not be called, and the ciphertext will not be complete.
+  """
+
+  def __init__(self, stream_aead, ciphertext_destination: BinaryIO,
+               associated_data: bytes):
+    """Create a new EncryptingStream.
+
+    Args:
+      stream_aead: C++ StreamingAead primitive from which a C++ EncryptingStream
+        will be obtained.
+      ciphertext_destination: A writable file-like object to which ciphertext
+        bytes will be written.
+      associated_data: The associated data to use for encryption. This must
+        match the associated_data used for decryption.
+    """
+    super(EncryptingStream, self).__init__()
+    self._closed = False
+    self._bytes_written = 0
+
+    # Create FileObjectAdapter
+    if not ciphertext_destination.writable():
+      raise ValueError('ciphertext_destination must be writable')
+    cc_ciphertext_destination = file_object_adapter.FileObjectAdapter(
+        ciphertext_destination)
+    # Get OutputStreamAdapter of C++ EncryptingStream
+    self._output_stream_adapter = self._get_output_stream_adapter(
+        stream_aead, associated_data, cc_ciphertext_destination)
+
+  @staticmethod
+  @tink_error.use_tink_errors
+  def _get_output_stream_adapter(cc_primitive, aad, destination):
+    """Implemented as a separate method to ensure correct error transform."""
+    return cc_streaming_aead_wrappers.new_cc_encrypting_stream(
+        cc_primitive, aad, destination)
+
+  @tink_error.use_tink_errors
+  def write(self, b: bytes) -> int:
+    """Write the given buffer to the stream.
+
+    May use multiple calls to the underlying file object's write() method.
+
+    Returns:
+      The number of bytes written, which will always be the length of b in
+      bytes.
+
+    Raises:
+      BlockingIOError: if the write could not be fully completed, with
+        characters_written set to the number of bytes successfully written.
+      TinkError: if there was a permanent error.
+
+    Args:
+      b: The buffer to write.
+    """
+    self._check_not_closed()
+
+    if not isinstance(b, (bytes, memoryview, bytearray)):
+      raise TypeError('a bytes-like object is required, not {}'.format(
+          type(b).__name__))
+
+    # One call to OutputStreamAdapter.write() may call next() multiple times
+    # on the C++ EncryptingStream, but will perform a partial write if there is
+    # a temporary write error. Permanent write errors will bubble up as
+    # exceptions.
+    written = self._output_stream_adapter.write(b)
+    if written < 0:
+      raise tink_error.TinkError('Number of written bytes was negative')
+
+    self._bytes_written += written
+
+    if written < len(b):
+      raise io.BlockingIOError(errno.EAGAIN,
+                               'Write could not complete without blocking.',
+                               written)
+    elif written > len(b):
+      raise tink_error.TinkError(
+          'Number of written bytes was greater than length of bytes given')
+
+    return written
+
+  def writelines(self, lines: Iterable[bytes]) -> None:
+    """Write a list of lines to the stream.
+
+    Line separators are not added, so it is usual for each of the lines
+    provided to have a line separator at the end.
+
+    Args:
+      lines: An iterable of buffers to write to the stream.
+    """
+    self._check_not_closed()
+    for line in lines:
+      self.write(line)
+
+  ### Internal ###
+
+  # TODO(tanujdhir) remove this function when b/141239409 b/141239470 are fixed.
+  def _check_not_closed(self, msg=None):
+    """Internal: raise a ValueError if file is closed."""
+    if self.closed:
+      raise ValueError('I/O operation on closed file.' if msg is None else msg)
+
+  ### Positioning ###
+
+  def position(self) -> int:
+    """Returns total number of written plaintext bytes."""
+    return self._bytes_written
+
+  ### Flush and close ###
+
+  def flush(self) -> None:
+    """Flush write buffers.
+
+    This method has no effect.
+    """
+    self._check_not_closed()
+    return
+
+  @tink_error.use_tink_errors
+  def close(self) -> None:
+    """Flush and close the stream.
+
+    This has no effect on a closed stream.
+    """
+    if self.closed:
+      return
+    self.flush()
+    self._output_stream_adapter.close()
+    self._closed = True
+
+  def __del__(self):
+    """Destructor.  Calls flush()."""
+    try:
+      # We deliberately don't close the file here, since we don't know if the
+      # user was really done writing or if there was an error.
+      self.flush()
+    except:  # pylint: disable=bare-except
+      pass
+
+  ### Inquiries ###
+
+  def writable(self) -> bool:
+    """Indicates whether object was opened for writing.
+
+    Returns:
+      Whether object was opened for writing.
+
+    If False, write() and truncate() will raise UnsupportedOperation.
+    """
+    return True
+
+  @property
+  def closed(self) -> bool:
+    """Indicates if the file has been closed.
+
+    Returns:
+      True if and only if the file has been closed.
+
+    For backwards compatibility, this is a property, not a predicate.
+    """
+    return self._closed
+
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    """Context management protocol.  Calls close() if there was no exception."""
+    # Calling close() signifies that the message is complete - we should not
+    # do this if there was an exception.
+    # Instead, we let the destructors be called, which should lead to sufficient
+    # cleanup on the other end, and if ciphertext_destination calls close() in
+    # __del__ (as IOBase does) then the underlying file descriptor should also
+    # be closed eventually.
+    if exc_type is None:
+      self.close()
+
+  ### Iterator ###
+  def __iter__(self):
+    """Iterator API."""
+    raise io.UnsupportedOperation('Cannot iterate an EncryptingStream')
+
+  def __next__(self):
+    """Iterator API."""
+    raise io.UnsupportedOperation('Cannot iterate an EncryptingStream')

python/streaming_aead/encrypting_stream_test.py

+# Lint as: python3
+# 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.streaming_aead.encrypting_stream."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import io
+import sys
+
+from absl.testing import absltest
+# TODO(b/141106504) Replace this with unittest.mock
+import mock
+
+from tink.python.streaming_aead import encrypting_stream
+
+
+def fake_get_output_stream_adapter(self, cc_primitive, aad, destination):
+  del cc_primitive, aad, self  # unused
+  return destination
+
+
+class TestBytesObject(io.BytesIO):
+  """A BytesIO object that does not close."""
+
+  def close(self):
+    pass
+
+
+def get_encrypting_stream(ciphertext_destination, aad):
+  return encrypting_stream.EncryptingStream(None, ciphertext_destination, aad)
+
+
+class EncryptingStreamTest(absltest.TestCase):
+
+  def setUp(self):
+    super(EncryptingStreamTest, self).setUp()
+    # Replace the EncryptingStream's staticmethod with a custom function to
+    # avoid the need for a Streaming AEAD primitive.
+    self.addCleanup(mock.patch.stopall)
+    mock.patch.object(
+        encrypting_stream.EncryptingStream,
+        '_get_output_stream_adapter',
+        new=fake_get_output_stream_adapter).start()
+
+  def test_non_writable_object(self):
+    f = mock.Mock()
+    f.writable = mock.Mock(return_value=False)
+    with self.assertRaisesRegex(ValueError, 'writable'):
+      get_encrypting_stream(f, b'aad')
+
+  def test_write(self):
+    f = TestBytesObject()
+    with get_encrypting_stream(f, b'aad') as es:
+      es.write(b'Hello world!')
+
+    self.assertEqual(b'Hello world!', f.getvalue())
+
+  @absltest.skipIf(sys.version_info[0] == 2, 'Python 2 strings are bytes')
+  def test_write_non_bytes(self):
+    with io.BytesIO() as f, get_encrypting_stream(f, b'aad') as es:
+      with self.assertRaisesRegex(TypeError, 'bytes-like object is required'):
+        es.write('This is a string, not a bytes object')
+
+  def test_textiowrapper_compatibility(self):
+    """A test that checks the TextIOWrapper works as expected.
+
+    It encrypts the same plaintext twice - once directly from bytes, and once
+    through TextIOWrapper's encoding. The two ciphertexts should have the same
+    length.
+    """
+    file_1 = TestBytesObject()
+    file_2 = TestBytesObject()
+
+    with get_encrypting_stream(file_1, b'aad') as es:
+      with io.TextIOWrapper(es) as wrapper:
+        # Need to specify this is a unicode string for Python 2.
+        wrapper.write(u'some data')
+
+    with get_encrypting_stream(file_2, b'aad') as es:
+      es.write(b'some data')
+
+    self.assertEqual(len(file_1.getvalue()), len(file_2.getvalue()))
+
+  def test_flush(self):
+    with io.BytesIO() as f, get_encrypting_stream(f, b'assoc') as es:
+      es.write(b'Hello world!')
+      es.flush()
+
+  def test_closed(self):
+    f = io.BytesIO()
+    es = get_encrypting_stream(f, b'assoc')
+    es.write(b'Hello world!')
+    es.close()
+
+    self.assertTrue(es.closed)
+    self.assertTrue(f.closed)
+
+  def test_closed_methods_raise(self):
+    f = io.BytesIO()
+    es = get_encrypting_stream(f, b'assoc')
+    es.write(b'Hello world!')
+    es.close()
+
+    with self.assertRaisesRegex(ValueError, 'closed'):
+      es.write(b'Goodbye world.')
+    with self.assertRaisesRegex(ValueError, 'closed'):
+      with es:
+        pass
+    with self.assertRaisesRegex(ValueError, 'closed'):
+      es.flush()
+
+  def test_unsupported_operation(self):
+    with io.BytesIO() as f, get_encrypting_stream(f, b'assoc') as es:
+      with self.assertRaisesRegex(io.UnsupportedOperation, 'seek'):
+        es.seek(0, 2)
+      with self.assertRaisesRegex(io.UnsupportedOperation, 'truncate'):
+        es.truncate(0)
+      with self.assertRaisesRegex(io.UnsupportedOperation, 'read'):
+        es.read(-1)
+
+  def test_inquiries(self):
+    with io.BytesIO() as f, get_encrypting_stream(f, b'assoc') as es:
+      self.assertTrue(es.writable())
+      self.assertFalse(es.readable())
+      self.assertFalse(es.seekable())
+
+  def test_position(self):
+    with io.BytesIO() as f:
+      with get_encrypting_stream(f, b'assoc') as es:
+        es.write(b'Hello world!')
+        self.assertEqual(es.position(), 12)
+
+  def test_position_works_closed(self):
+    with io.BytesIO() as f:
+      es = get_encrypting_stream(f, b'assoc')
+
+      es.write(b'Hello world!')
+      es.close()
+
+      self.assertTrue(es.closed)
+      self.assertEqual(es.position(), 12)
+
+  def test_blocking_io(self):
+
+    class OnlyWritesFirstFiveBytes(io.BytesIO):
+
+      def write(self, buffer):
+        buffer = buffer[:5]
+        n = super(OnlyWritesFirstFiveBytes, self).write(buffer)
+        return n
+
+    with OnlyWritesFirstFiveBytes() as f:
+      with get_encrypting_stream(f, b'assoc') as es:
+        with self.assertRaisesRegex(io.BlockingIOError, 'could not complete'):
+          es.write(b'Hello world!')
+
+  def test_context_manager_exception_close(self):
+    """Tests that exceptional exits do not trigger normal file closure.
+
+    Instead, the file will be closed without a proper final ciphertext block,
+    and will result in an invalid ciphertext. The ciphertext_destination file
+    object itself should in most cases still be closed when garbage collected.
+    """
+    f = io.BytesIO()
+    with self.assertRaisesRegex(ValueError, 'raised inside'):
+      with get_encrypting_stream(f, b'assoc') as es:
+        es.write(b'some message')
+        raise ValueError('Error raised inside context manager')
+
+    self.assertFalse(f.closed)
+
+
+if __name__ == '__main__':
+  absltest.main()

python/streaming_aead/streaming_aead.py

@@ -66,7 +66,6 @@ class StreamingAead(object):
         write()
         close()
         closed
-        tell()
         isatty()
         flush() (no-op)
         writable()
@@ -76,6 +75,8 @@ class StreamingAead(object):
       Other methods, including read() and seek(), raise io.UnsupportedOperation.
       Closing the wrapper also closes the ciphertext_destination.
 
+      The wrapper also supports the position() method, which returns the number
+      of written plaintext bytes.
     Raises:
       tink.TinkError if the creation fails.
     """