Add and update tests

Author: tulir

mautrix/client/state_store/tests/store_test.py

@@ -3,7 +3,9 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
-from typing import AsyncContextManager, AsyncIterator, Callable, Dict, List
+from __future__ import annotations
+
+from typing import AsyncContextManager, AsyncIterator, Callable
 from contextlib import asynccontextmanager
 import json
 import os
@@ -83,7 +85,7 @@ async def store(request) -> AsyncIterator[StateStore]:
         yield state_store
 
 
-def read_state_file(request, file) -> Dict[RoomID, List[StateEvent]]:
+def read_state_file(request, file) -> dict[RoomID, list[StateEvent]]:
     path = pathlib.Path(request.node.fspath).with_name(file)
     with path.open() as fp:
         content = json.load(fp)
@@ -122,7 +124,6 @@ async def get_joined_members(request, store: StateStore) -> None:
         await store.set_members(room_id, parsed_members, only_membership=Membership.JOIN)
 
 
-@pytest.mark.asyncio
 async def test_basic(store: StateStore) -> None:
     room_id = RoomID("!foo:example.com")
     user_id = UserID("@tulir:example.com")
@@ -136,7 +137,6 @@ async def test_basic(store: StateStore) -> None:
     assert await store.is_encrypted(RoomID("!unknown-room:example.com")) is None
 
 
-@pytest.mark.asyncio
 async def test_basic_updated(request, store: StateStore) -> None:
     await store_room_state(request, store)
     test_group = RoomID("!telegram-group:example.com")
@@ -145,7 +145,6 @@ async def test_basic_updated(request, store: StateStore) -> None:
     assert not await store.is_encrypted(RoomID("!unencrypted-room:example.com"))
 
 
-@pytest.mark.asyncio
 async def test_updates(request, store: StateStore) -> None:
     await store_room_state(request, store)
     room_id = RoomID("!telegram-group:example.com")

mautrix/crypto/attachments/async_attachments_test.py

@@ -0,0 +1,44 @@
+# Copyright © 2019 Damir Jelić <[email protected]> (under the Apache 2.0 license)
+# Copyright © 2019 miruka <[email protected]> (under the Apache 2.0 license)
+# Copyright (c) 2022 Tulir Asokan
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+from mautrix.types import EncryptedFile
+
+from .async_attachments import async_encrypt_attachment, async_inplace_encrypt_attachment
+from .attachments import decrypt_attachment
+
+try:
+    from Crypto import Random
+except ImportError:
+    from Cryptodome import Random
+
+
+async def _get_data_cypher_keys(data: bytes) -> tuple[bytes, EncryptedFile]:
+    *chunks, keys = [i async for i in async_encrypt_attachment(data)]
+    return b"".join(chunks), keys
+
+
+async def test_async_encrypt():
+    data = b"Test bytes"
+
+    cyphertext, keys = await _get_data_cypher_keys(data)
+
+    plaintext = decrypt_attachment(cyphertext, keys.key.key, keys.hashes["sha256"], keys.iv)
+
+    assert data == plaintext
+
+
+async def test_async_inplace_encrypt():
+    orig_data = b"Test bytes"
+    data = bytearray(orig_data)
+
+    keys = await async_inplace_encrypt_attachment(data)
+
+    assert data != orig_data
+
+    decrypt_attachment(data, keys.key.key, keys.hashes["sha256"], keys.iv, inplace=True)
+
+    assert data == orig_data

mautrix/crypto/attachments/attachments.py

@@ -29,42 +29,46 @@
     from Cryptodome.Util import Counter
 
 
-def decrypt_attachment(ciphertext: bytes, key: str, hash: str, iv: str) -> bytes:
+def decrypt_attachment(
+    ciphertext: bytes | bytearray | memoryview, key: str, hash: str, iv: str, inplace: bool = False
+) -> bytes:
     """Decrypt an encrypted attachment.
 
     Args:
         ciphertext: The data to decrypt.
         key: AES_CTR JWK key object.
         hash: Base64 encoded SHA-256 hash of the ciphertext.
         iv: Base64 encoded 16 byte AES-CTR IV.
+        inplace: Should the decryption be performed in-place?
+                 The input must be a bytearray or writable memoryview to use this.
     Returns:
         The plaintext bytes.
     Raises:
         EncryptionError: if the integrity check fails.
-
-
     """
     expected_hash = unpaddedbase64.decode_base64(hash)
 
     h = SHA256.new()
     h.update(ciphertext)
 
     if h.digest() != expected_hash:
-        raise DecryptionError("Mismatched SHA-256 digest.")
+        raise DecryptionError("Mismatched SHA-256 digest")
 
     try:
         byte_key: bytes = unpaddedbase64.decode_base64(key)
     except (binascii.Error, TypeError):
-        raise DecryptionError("Error decoding key.")
+        raise DecryptionError("Error decoding key")
 
     try:
         byte_iv: bytes = unpaddedbase64.decode_base64(iv)
+        if len(byte_iv) != 16:
+            raise DecryptionError("Invalid IV length")
         prefix = byte_iv[:8]
         # A non-zero IV counter is not spec-compliant, but some clients still do it,
         # so decode the counter part too.
         initial_value = struct.unpack(">Q", byte_iv[8:])[0]
     except (binascii.Error, TypeError, IndexError, struct.error):
-        raise DecryptionError("Error decoding initial values.")
+        raise DecryptionError("Error decoding IV")
 
     ctr = Counter.new(64, prefix=prefix, initial_value=initial_value)
 
@@ -73,7 +77,11 @@ def decrypt_attachment(ciphertext: bytes, key: str, hash: str, iv: str) -> bytes
     except ValueError as e:
         raise DecryptionError("Failed to create AES cipher") from e
 
-    return cipher.decrypt(ciphertext)
+    if inplace:
+        cipher.decrypt(ciphertext, ciphertext)
+        return ciphertext
+    else:
+        return cipher.decrypt(ciphertext)
 
 
 def encrypt_attachment(plaintext: bytes) -> tuple[bytes, EncryptedFile]:
@@ -103,7 +111,7 @@ def _prepare_encryption() -> tuple[bytes, bytes, AES, SHA256.SHA256Hash]:
     return key, iv, cipher, sha256
 
 
-def inplace_encrypt_attachment(data: bytearray) -> EncryptedFile:
+def inplace_encrypt_attachment(data: bytearray | memoryview) -> EncryptedFile:
     key, iv, cipher, sha256 = _prepare_encryption()
 
     cipher.encrypt(plaintext=data, output=data)

mautrix/crypto/attachments/attachments_test.py

@@ -0,0 +1,111 @@
+# Copyright © 2019 Damir Jelić <[email protected]> (under the Apache 2.0 license)
+# Copyright (c) 2022 Tulir Asokan
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+import pytest
+import unpaddedbase64
+
+from mautrix.errors import DecryptionError
+
+from .attachments import decrypt_attachment, encrypt_attachment, inplace_encrypt_attachment
+
+try:
+    from Crypto import Random
+except ImportError:
+    from Cryptodome import Random
+
+
+def test_encrypt():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    plaintext = decrypt_attachment(cyphertext, keys.key.key, keys.hashes["sha256"], keys.iv)
+
+    assert data == plaintext
+
+
+def test_inplace_encrypt():
+    orig_data = b"Test bytes"
+    data = bytearray(orig_data)
+
+    keys = inplace_encrypt_attachment(data)
+
+    assert data != orig_data
+
+    decrypt_attachment(data, keys.key.key, keys.hashes["sha256"], keys.iv, inplace=True)
+
+    assert data == orig_data
+
+
+def test_hash_verification():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    with pytest.raises(DecryptionError):
+        decrypt_attachment(cyphertext, keys.key.key, "Fake hash", keys.iv)
+
+
+def test_invalid_key():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    with pytest.raises(DecryptionError):
+        decrypt_attachment(cyphertext, "Fake key", keys.hashes["sha256"], keys.iv)
+
+
+def test_invalid_iv():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    with pytest.raises(DecryptionError):
+        decrypt_attachment(cyphertext, keys.key.key, keys.hashes["sha256"], "Fake iv")
+
+
+def test_short_key():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    with pytest.raises(DecryptionError):
+        decrypt_attachment(
+            cyphertext,
+            unpaddedbase64.encode_base64(b"Fake key", urlsafe=True),
+            keys["hashes"]["sha256"],
+            keys["iv"],
+        )
+
+
+def test_short_iv():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    with pytest.raises(DecryptionError):
+        decrypt_attachment(
+            cyphertext,
+            keys.key.key,
+            keys.hashes["sha256"],
+            unpaddedbase64.encode_base64(b"F" + b"\x00" * 8),
+        )
+
+
+def test_fake_key():
+    data = b"Test bytes"
+
+    cyphertext, keys = encrypt_attachment(data)
+
+    fake_key = Random.new().read(32)
+
+    plaintext = decrypt_attachment(
+        cyphertext,
+        unpaddedbase64.encode_base64(fake_key, urlsafe=True),
+        keys["hashes"]["sha256"],
+        keys["iv"],
+    )
+    assert plaintext != data