def Encrypt(stream: IO[bytes], key: bytes) -> IO[bytes]:
"""Encrypts given file-like object using AES algorithm with GCM mode.
The input stream is divided into small chunks of predefined size and then each
chunk is encrypted using AES GCM procedure. In the encoded stream, before each
proper chunk there is a nonce binary string prepended. As associated data for
each encrypted chunk, chunk index and information about whether it is the last
chunk is used.
Args:
stream: A file-like object to encrypt.
key: A secret key used for encrypting the data.
Returns:
A file-like object with encrypted data.
"""
aesgcm = aead.AESGCM(key)
def Generate() -> Iterator[bytes]:
chunks = ioutil.Chunk(stream, size=_AEAD_CHUNK_SIZE)
chunks = iterator.Lookahead(enumerate(chunks))
for idx, chunk in chunks:
nonce = os.urandom(_AEAD_NONCE_SIZE)
adata = _AEAD_ADATA_FORMAT.pack(idx, chunks.done)
yield nonce + aesgcm.encrypt(nonce, chunk, adata)
return ioutil.Unchunk(Generate())
def testMultiChunkReadByByte(self):
stream = ioutil.Unchunk(iter([b"foo", b"bar", b"baz"]))
self.assertEqual(stream.read(1), b"f")
self.assertEqual(stream.read(1), b"o")
self.assertEqual(stream.read(1), b"o")
self.assertEqual(stream.read(1), b"b")
self.assertEqual(stream.read(1), b"a")
self.assertEqual(stream.read(1), b"r")
self.assertEqual(stream.read(1), b"b")
self.assertEqual(stream.read(1), b"a")
self.assertEqual(stream.read(1), b"z")
self.assertEqual(stream.read(1), b"")
def Decrypt(stream: IO[bytes], key: bytes) -> IO[bytes]:
"""Decrypts given file-like object using AES algorithm in GCM mode.
Refer to the encryption documentation to learn about the details of the format
that this function allows to decode.
Args:
stream: A file-like object to decrypt.
key: A secret key used for decrypting the data.
Returns:
A file-like object with decrypted data.
"""
aesgcm = aead.AESGCM(key)
def Generate() -> Iterator[bytes]:
# Buffered reader should accept `IO[bytes]` but for now it accepts only
# `RawIOBase` (which is a concrete base class for all I/O implementations).
reader = io.BufferedReader(stream) # pytype: disable=wrong-arg-types
# We abort early if there is no data in the stream. Otherwise we would try
# to read nonce and fail.
if not reader.peek():
return
for idx in itertools.count():
nonce = reader.read(_AEAD_NONCE_SIZE)
# As long there is some data in the buffer (and there should be because of
# the initial check) there should be a fixed-size nonce prepended to each
# chunk.
if len(nonce) != _AEAD_NONCE_SIZE:
raise EOFError(f"Incorrect nonce length: {len(nonce)}")
chunk = reader.read(_AEAD_CHUNK_SIZE + 16)
# `BufferedReader#peek` will return non-empty byte string if there is more
# data available in the stream.
is_last = reader.peek() == b"" # pylint: disable=g-explicit-bool-comparison
adata = _AEAD_ADATA_FORMAT.pack(idx, is_last)
yield aesgcm.decrypt(nonce, chunk, adata)
if is_last:
break
return ioutil.Unchunk(Generate())
def Stream(url: str, encryption_key: bytes) -> Iterator[IO[bytes]]:
"""Streams a decrypted large file from the given URL."""
with requests.get(url, stream=True) as streamer:
stream = ioutil.Unchunk(streamer.iter_content(io.DEFAULT_BUFFER_SIZE))
yield aead.Decrypt(stream, encryption_key)
def testMultiChunkReadLarge(self):
stream = ioutil.Unchunk(itertools.cycle([b"\x42" * 1021]))
self.assertEqual(stream.read(16777259), b"\x42" * 16777259)
def testMultiChunkReadWhole(self):
stream = ioutil.Unchunk(iter([b"foo", b"bar", b"baz"]))
self.assertEqual(stream.read(), b"foobarbaz")
def testSingleChunkReadByByte(self):
stream = ioutil.Unchunk(iter([b"foo"]))
self.assertEqual(stream.read(1), b"f")
self.assertEqual(stream.read(1), b"o")
self.assertEqual(stream.read(1), b"o")
self.assertEqual(stream.read(1), b"")
def testSingleChunkReadWhole(self):
stream = ioutil.Unchunk(iter([b"foobar"]))
self.assertEqual(stream.read(), b"foobar")
def testEmpty(self):
stream = ioutil.Unchunk(iter([]))
self.assertEqual(stream.read(), b"")