initial commit

Signed-off-by: Kamal Tufekcic <kamal@lo.sh>

soliton_py/tests/test_stream.py

@@ -0,0 +1,161 @@
+"""Tests for streaming AEAD."""
+
+import os
+import soliton
+
+
+CHUNK_SIZE = 1_048_576  # 1 MiB
+
+
+def _random_key():
+    return os.urandom(32)
+
+
+def test_stream_single_chunk():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        assert len(header) == 26
+        ct = enc.encrypt_chunk(b"hello stream", is_last=True)
+        assert enc.is_finalized()
+
+    with soliton.StreamDecryptor(key, header) as dec:
+        pt, is_last = dec.decrypt_chunk(ct)
+        assert pt == b"hello stream"
+        assert is_last
+
+
+def test_stream_multi_chunk():
+    key = _random_key()
+    chunk1_data = os.urandom(CHUNK_SIZE)  # non-final must be 1 MiB
+    chunk2_data = b"final chunk"
+
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        ct1 = enc.encrypt_chunk(chunk1_data, is_last=False)
+        ct2 = enc.encrypt_chunk(chunk2_data, is_last=True)
+
+    with soliton.StreamDecryptor(key, header) as dec:
+        pt1, last1 = dec.decrypt_chunk(ct1)
+        assert pt1 == chunk1_data
+        assert not last1
+
+        pt2, last2 = dec.decrypt_chunk(ct2)
+        assert pt2 == chunk2_data
+        assert last2
+
+
+def test_stream_with_aad():
+    key = _random_key()
+    aad = b"file-id:12345"
+
+    with soliton.StreamEncryptor(key, aad=aad) as enc:
+        header = enc.header()
+        ct = enc.encrypt_chunk(b"aad-bound data", is_last=True)
+
+    # Correct AAD decrypts.
+    with soliton.StreamDecryptor(key, header, aad=aad) as dec:
+        pt, _ = dec.decrypt_chunk(ct)
+        assert pt == b"aad-bound data"
+
+    # Wrong AAD fails.
+    try:
+        with soliton.StreamDecryptor(key, header, aad=b"wrong") as dec:
+            dec.decrypt_chunk(ct)
+            assert False, "should have raised"
+    except soliton.AeadError:
+        pass
+
+
+def test_stream_wrong_key_fails():
+    key1 = _random_key()
+    key2 = _random_key()
+    with soliton.StreamEncryptor(key1) as enc:
+        header = enc.header()
+        ct = enc.encrypt_chunk(b"data", is_last=True)
+
+    try:
+        with soliton.StreamDecryptor(key2, header) as dec:
+            dec.decrypt_chunk(ct)
+            assert False, "should have raised"
+    except soliton.AeadError:
+        pass
+
+
+def test_stream_tampered_ciphertext():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        ct = bytearray(enc.encrypt_chunk(b"data", is_last=True))
+    ct[0] ^= 0xFF
+
+    try:
+        with soliton.StreamDecryptor(key, header) as dec:
+            dec.decrypt_chunk(bytes(ct))
+            assert False, "should have raised"
+    except (soliton.AeadError, soliton.InvalidDataError):
+        pass
+
+
+def test_stream_empty_plaintext():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        ct = enc.encrypt_chunk(b"", is_last=True)
+
+    with soliton.StreamDecryptor(key, header) as dec:
+        pt, is_last = dec.decrypt_chunk(ct)
+        assert pt == b""
+        assert is_last
+
+
+def test_stream_encrypt_at_decrypt_at():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        # Short plaintext requires is_last=True (non-final must be exactly 1 MiB).
+        ct0 = enc.encrypt_chunk_at(0, b"at index zero", is_last=True)
+
+    with soliton.StreamDecryptor(key, header) as dec:
+        pt0, last0 = dec.decrypt_chunk_at(0, ct0)
+        assert pt0 == b"at index zero"
+        assert last0
+
+
+def test_stream_decrypt_at_different_indices():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        ct5 = enc.encrypt_chunk_at(5, b"at five", is_last=True)
+        ct99 = enc.encrypt_chunk_at(99, b"at ninety-nine", is_last=True)
+
+    # Decrypt in reverse order.
+    with soliton.StreamDecryptor(key, header) as dec:
+        pt99, _ = dec.decrypt_chunk_at(99, ct99)
+        assert pt99 == b"at ninety-nine"
+
+        pt5, _ = dec.decrypt_chunk_at(5, ct5)
+        assert pt5 == b"at five"
+
+
+def test_stream_expected_index():
+    key = _random_key()
+    with soliton.StreamEncryptor(key) as enc:
+        header = enc.header()
+        ct = enc.encrypt_chunk(b"chunk", is_last=True)
+
+    with soliton.StreamDecryptor(key, header) as dec:
+        assert dec.expected_index() == 0
+        dec.decrypt_chunk(ct)
+        assert dec.expected_index() == 1
+
+
+def test_stream_context_manager():
+    enc = soliton.StreamEncryptor(_random_key())
+    with enc:
+        enc.header()
+    try:
+        enc.header()
+        assert False, "should have raised"
+    except soliton.InvalidDataError:
+        pass