nacl_streams.py

from nacl.bindings.crypto_secretstream import (
    crypto_secretstream_xchacha20poly1305_ABYTES,
    crypto_secretstream_xchacha20poly1305_HEADERBYTES,
    crypto_secretstream_xchacha20poly1305_KEYBYTES,
    crypto_secretstream_xchacha20poly1305_TAG_MESSAGE,
    crypto_secretstream_xchacha20poly1305_TAG_FINAL,
    crypto_secretstream_xchacha20poly1305_init_pull,
    crypto_secretstream_xchacha20poly1305_init_push,
    crypto_secretstream_xchacha20poly1305_pull,
    crypto_secretstream_xchacha20poly1305_push,
    crypto_secretstream_xchacha20poly1305_state,
)
from nacl.utils import random
import struct

# version, chunk size
_header_struct = struct.Struct('<BQ')

def encrypt_stream(srcfp, destfp, symmetric_key, chunk_size=4096):
    if len(symmetric_key) != crypto_secretstream_xchacha20poly1305_KEYBYTES:
        raise ValueError('symmetric key is too short')
    state = crypto_secretstream_xchacha20poly1305_state()
    hdr = crypto_secretstream_xchacha20poly1305_init_push(state, symmetric_key)
    destfp.write(hdr)
    frame = crypto_secretstream_xchacha20poly1305_push(
        state,
        _header_struct.pack(1, chunk_size),
        None,
        crypto_secretstream_xchacha20poly1305_TAG_MESSAGE,
    )
    destfp.write(frame)

    eof = False
    while not eof:
        msg = srcfp.read(chunk_size)
        eof = len(msg) < chunk_size
        tag = crypto_secretstream_xchacha20poly1305_TAG_FINAL if eof else 0
        frame = crypto_secretstream_xchacha20poly1305_push(state, msg, tag=tag)
        destfp.write(frame)

def decrypt_stream(srcfp, destfp, symmetric_key):
    hdr = srcfp.read(crypto_secretstream_xchacha20poly1305_HEADERBYTES)
    if not hdr:
        raise ValueError('corrupted stream, missing header')
    state = crypto_secretstream_xchacha20poly1305_state()
    crypto_secretstream_xchacha20poly1305_init_pull(state, hdr, symmetric_key)

    # read the header version and chunk size
    hdr = srcfp.read(
        _header_struct.size + crypto_secretstream_xchacha20poly1305_ABYTES,
    )
    chunk, tag = crypto_secretstream_xchacha20poly1305_pull(state, hdr, None)
    version, chunk_size = _header_struct.unpack(chunk)
    if version != 1:
        raise ValueError('unsupported stream version={0}'.format(version))
    frame_size = chunk_size + crypto_secretstream_xchacha20poly1305_ABYTES

    while tag != crypto_secretstream_xchacha20poly1305_TAG_FINAL:
        frame = srcfp.read(frame_size)
        if not frame:
            raise ValueError('corrupted stream, missing chunks')
        chunk, tag = crypto_secretstream_xchacha20poly1305_pull(state, frame)
        destfp.write(chunk)

def main():
    import sys
    key = random(crypto_secretstream_xchacha20poly1305_KEYBYTES)
    with open(sys.argv[1], 'rb') as srcfp, open('foo.enc', 'wb') as dstfp:
        encrypt_stream(srcfp, dstfp, key)
    with open('foo.enc', 'rb') as srcfp, open('foo.dec', 'wb') as dstfp:
        decrypt_stream(srcfp, dstfp, key)

if __name__ == '__main__':
    main()

I know it's not supposed to be perfect, but as this example was linked in related PyNaCl PR, and since there's key length check already, maybe worth adding another check for chunk count in encrypt_stream?

As it is, anything larger than about "255 * chunk_size" bytes will be silently (!!!) encrypted into non-decryptable output.
Which can be a nasty surprise if someone uses this code for at-rest encryption after only basic testing on small streams.

As it is, anything larger than about "255 * chunk_size" bytes will be silently (!!!) encrypted into non-decryptable output.

Not sure if I was correct about that size, maybe example just breaks with anything above ~10K:

% rm -f foo && dd if=/dev/urandom of=foo bs=20K count=1 && py3 nacl_streams.py foo
Traceback (most recent call last):
  File "nacl_streams.py", line 74, in <module>
    main()
  File "nacl_streams.py", line 71, in main
    decrypt_stream(srcfp, dstfp, key)
  File "nacl_streams.py", line 62, in decrypt_stream
    chunk, tag = crypto_secretstream_xchacha20poly1305_pull(state, frame)
  File "/home/fraggod/.py3/nacl/bindings/crypto_secretstream.py", line 267, in crypto_secretstream_xchacha20poly1305_pull
    raising=exc.ValueError,
  File "/home/fraggod/.py3/nacl/exceptions.py", line 68, in ensure
    raise raising(*args)
nacl.exceptions.ValueError: Ciphertext is too short

EDIT:
Tracked the issue down to this line - frame = crypto_secretstream_xchacha20poly1305_push(state, msg, tag=tag) - error above happens with zero-length msg there, i.e. if file size divides into blocks without remainder.
To fix that, guess either different eof check should be used (e.g. C examples use eof = feof(fp_s);), or last block to always contain some dummy byte to be deliberately discarded.