Add utilities for SSSS

Author: tulir

mautrix/crypto/ssss/__init__.py

@@ -0,0 +1,7 @@
+from .key import Key, KeyMetadata, PassphraseMetadata
+from .types import (
+    Algorithm,
+    EncryptedAccountDataEventContent,
+    EncryptedKeyData,
+    PassphraseAlgorithm,
+)

mautrix/crypto/ssss/key.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2025 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 typing import Optional
+import base64
+import hashlib
+import hmac
+
+from attr import dataclass
+import unpaddedbase64
+
+from mautrix.types import EventType, SerializableAttrs
+
+from .types import Algorithm, EncryptedKeyData, PassphraseAlgorithm
+from .util import (
+    calculate_hash,
+    cryptorand,
+    decode_base58_recovery_key,
+    derive_keys,
+    encode_base58_recovery_key,
+    prepare_aes,
+)
+
+try:
+    from Crypto.Cipher import AES
+    from Crypto.Util import Counter
+except ImportError:
+    from Cryptodome.Cipher import AES
+    from Cryptodome.Util import Counter
+
+
+@dataclass
+class PassphraseMetadata(SerializableAttrs):
+    algorithm: PassphraseAlgorithm
+    iterations: int
+    salt: str
+    bits: int = 256
+
+    def get_key(self, passphrase: str) -> bytes:
+        if self.algorithm != PassphraseAlgorithm.PBKDF2:
+            raise ValueError(f"Unsupported passphrase algorithm {self.algorithm}")
+        return hashlib.pbkdf2_hmac(
+            "sha512",
+            passphrase.encode("utf-8"),
+            self.salt.encode("utf-8"),
+            self.iterations,
+            self.bits // 8,
+        )
+
+
+@dataclass
+class KeyMetadata(SerializableAttrs):
+    algorithm: Algorithm
+
+    iv: str | None = None
+    mac: str | None = None
+
+    name: str | None = None
+    passphrase: Optional[PassphraseMetadata] = None
+
+    def verify_passphrase(self, key_id: str, phrase: str) -> "Key":
+        if not self.passphrase:
+            raise ValueError("Passphrase not set on this key")
+        return self.verify_raw_key(key_id, self.passphrase.get_key(phrase))
+
+    def verify_recovery_key(self, key_id: str, recovery_key: str) -> "Key":
+        return self.verify_raw_key(key_id, decode_base58_recovery_key(recovery_key))
+
+    def verify_raw_key(self, key_id: str, key: bytes) -> "Key":
+        if self.mac.rstrip("=") != calculate_hash(key, self.iv):
+            raise ValueError("Key MAC does not match")
+        return Key(id=key_id, key=key, metadata=self)
+
+
+@dataclass
+class Key:
+    id: str
+    key: bytes
+    metadata: KeyMetadata
+
+    @classmethod
+    def generate(cls, passphrase: str | None = None) -> "Key":
+        passphrase_meta = (
+            PassphraseMetadata(
+                algorithm=PassphraseAlgorithm.PBKDF2,
+                iterations=500_000,
+                salt=base64.b64encode(cryptorand.read(24)).decode("utf-8"),
+                bits=256,
+            )
+            if passphrase
+            else None
+        )
+        key = passphrase_meta.get_key(passphrase) if passphrase else cryptorand.read(32)
+        iv = unpaddedbase64.encode_base64(cryptorand.read(16))
+        metadata = KeyMetadata(
+            algorithm=Algorithm.AES_HMAC_SHA2,
+            passphrase=passphrase_meta,
+            mac=calculate_hash(key, iv),
+            iv=iv,
+        )
+        key_id = unpaddedbase64.encode_base64(cryptorand.read(24))
+        return cls(key=key, id=key_id, metadata=metadata)
+
+    @property
+    def recovery_key(self) -> str:
+        return encode_base58_recovery_key(self.key)
+
+    def encrypt(self, event_type: str | EventType, data: str | bytes) -> EncryptedKeyData:
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        data = base64.b64encode(data).rstrip(b"=")
+
+        aes_key, hmac_key = derive_keys(self.key, event_type)
+        iv = bytearray(cryptorand.read(16))
+        iv[8] &= 0x7F
+        ciphertext = prepare_aes(aes_key, iv).encrypt(data)
+        digest = hmac.digest(hmac_key, ciphertext, hashlib.sha256)
+        return EncryptedKeyData(
+            ciphertext=unpaddedbase64.encode_base64(ciphertext),
+            iv=unpaddedbase64.encode_base64(iv),
+            mac=unpaddedbase64.encode_base64(digest),
+        )
+
+    def decrypt(self, event_type: str | EventType, data: EncryptedKeyData) -> bytes:
+        aes_key, hmac_key = derive_keys(self.key, event_type)
+        ciphertext = unpaddedbase64.decode_base64(data.ciphertext)
+        mac = unpaddedbase64.decode_base64(data.mac)
+
+        expected_mac = hmac.digest(hmac_key, ciphertext, hashlib.sha256)
+        if not hmac.compare_digest(mac, expected_mac):
+            raise ValueError("Invalid MAC")
+
+        plaintext = prepare_aes(aes_key, data.iv).decrypt(ciphertext)
+        return unpaddedbase64.decode_base64(plaintext.decode("utf-8"))

mautrix/crypto/ssss/key_test.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2025 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
+
+from ...types.event.type import EventType
+from .key import Key, KeyMetadata
+from .types import EncryptedAccountDataEventContent
+
+KEY1_CROSS_SIGNING_MASTER_KEY = """{
+  "encrypted": {
+    "gEJqbfSEMnP5JXXcukpXEX1l0aI3MDs0": {
+      "iv": "BpKP9nQJTE9jrsAssoxPqQ==",
+      "ciphertext": "fNRiiiidezjerTgV+G6pUtmeF3izzj5re/mVvY0hO2kM6kYGrxLuIu2ej80=",
+      "mac": "/gWGDGMyOLmbJp+aoSLh5JxCs0AdS6nAhjzpe+9G2Q0="
+    }
+  }
+}"""
+
+KEY1_CROSS_SIGNING_MASTER_KEY_DECRYPTED = bytes(
+    [
+        0x68,
+        0xF9,
+        0x7F,
+        0xD1,
+        0x92,
+        0x2E,
+        0xEC,
+        0xF6,
+        0xB8,
+        0x2B,
+        0xB8,
+        0x90,
+        0xD2,
+        0x4D,
+        0x06,
+        0x52,
+        0x98,
+        0x4E,
+        0x7A,
+        0x1D,
+        0x70,
+        0x3B,
+        0x9E,
+        0x86,
+        0x7B,
+        0x7E,
+        0xBA,
+        0xF7,
+        0xFE,
+        0xB9,
+        0x5B,
+        0x6F,
+    ]
+)
+
+KEY1_META = """{
+  "algorithm": "m.secret_storage.v1.aes-hmac-sha2",
+  "passphrase": {
+    "algorithm": "m.pbkdf2",
+    "iterations": 500000,
+    "salt": "y863BOoqOadgDp8S3FtHXikDJEalsQ7d"
+  },
+  "iv": "xxkTK0L4UzxgAFkQ6XPwsw",
+  "mac": "MEhooO0ZhFJNxUhvRMSxBnJfL20wkLgle3ocY0ee/eA"
+}"""
+KEY1_ID = "gEJqbfSEMnP5JXXcukpXEX1l0aI3MDs0"
+KEY1_RECOVERY_KEY = "EsTE s92N EtaX s2h6 VQYF 9Kao tHYL mkyL GKMh isZb KJ4E tvoC"
+KEY1_PASSPHRASE = "correct horse battery staple"
+
+KEY2_META = """{
+  "algorithm": "m.secret_storage.v1.aes-hmac-sha2",
+  "iv": "O0BOvTqiIAYjC+RMcyHfWw==",
+  "mac": "7k6OruQlWg0UmQjxGZ0ad4Q6DdwkgnoI7G6X3IjBYtI="
+}"""
+KEY2_ID = "NVe5vK6lZS9gEMQLJw0yqkzmE5Mr7dLv"
+KEY2_RECOVERY_KEY = "EsUC xSxt XJgQ dz19 8WBZ rHdE GZo7 ybsn EFmG Y5HY MDAG GNWe"
+
+KEY2_META_BROKEN_IV = """{
+  "algorithm": "m.secret_storage.v1.aes-hmac-sha2",
+  "iv": "O0BOvTqiIAYjC+RMcyHfWwMeowMeowMeow",
+  "mac": "7k6OruQlWg0UmQjxGZ0ad4Q6DdwkgnoI7G6X3IjBYtI="
+}"""
+
+KEY2_META_BROKEN_MAC = """{
+  "algorithm": "m.secret_storage.v1.aes-hmac-sha2",
+  "iv": "O0BOvTqiIAYjC+RMcyHfWw==",
+  "mac": "7k6OruQlWg0UmQjxGZ0ad4Q6DdwkgnoI7G6X3IjBYtIMeowMeowMeow"
+}"""
+
+
+def get_key_meta(meta: str) -> KeyMetadata:
+    return KeyMetadata.parse_json(meta)
+
+
+def get_key1() -> Key:
+    return get_key_meta(KEY1_META).verify_recovery_key(KEY1_ID, KEY1_RECOVERY_KEY)
+
+
+def get_key2() -> Key:
+    return get_key_meta(KEY2_META).verify_recovery_key(KEY2_ID, KEY2_RECOVERY_KEY)
+
+
+def get_encrypted_master_key() -> EncryptedAccountDataEventContent:
+    return EncryptedAccountDataEventContent.parse_json(KEY1_CROSS_SIGNING_MASTER_KEY)
+
+
+def test_decrypt_success() -> None:
+    key = get_key1()
+    emk = get_encrypted_master_key()
+    assert (
+        emk.decrypt(EventType.CROSS_SIGNING_MASTER, key) == KEY1_CROSS_SIGNING_MASTER_KEY_DECRYPTED
+    )
+
+
+def test_decrypt_fail_wrong_key() -> None:
+    key = get_key2()
+    emk = get_encrypted_master_key()
+    with pytest.raises(ValueError):
+        emk.decrypt(EventType.CROSS_SIGNING_MASTER, key)
+
+
+def test_decrypt_fail_fake_key() -> None:
+    key = get_key2()
+    key.id = KEY1_ID
+    emk = get_encrypted_master_key()
+    with pytest.raises(ValueError):
+        emk.decrypt(EventType.CROSS_SIGNING_MASTER, key)
+
+
+def test_decrypt_fail_wrong_type() -> None:
+    key = get_key1()
+    emk = get_encrypted_master_key()
+    with pytest.raises(ValueError):
+        emk.decrypt(EventType.CROSS_SIGNING_SELF_SIGNING, key)
+
+
+def test_encrypt_roundtrip() -> None:
+    key = get_key1()
+    data = bytes([0xDE, 0xAD, 0xBE, 0xEF])
+    ciphertext = key.encrypt("net.maunium.data", data)
+    plaintext = key.decrypt("net.maunium.data", ciphertext)
+    assert plaintext == data
+
+
+def test_verify_recovery_key_correct() -> None:
+    meta = get_key_meta(KEY1_META)
+    key = meta.verify_recovery_key(KEY1_ID, KEY1_RECOVERY_KEY)
+    assert key.recovery_key == KEY1_RECOVERY_KEY
+
+
+def test_verify_recovery_key_correct2() -> None:
+    meta = get_key_meta(KEY2_META)
+    key = meta.verify_recovery_key(KEY2_ID, KEY2_RECOVERY_KEY)
+    assert key.recovery_key == KEY2_RECOVERY_KEY
+
+
+def test_verify_recovery_key_invalid() -> None:
+    meta = get_key_meta(KEY1_META)
+    with pytest.raises(ValueError):
+        meta.verify_recovery_key(KEY1_ID, "foo")
+
+
+def test_verify_recovery_key_incorrect() -> None:
+    meta = get_key_meta(KEY1_META)
+    with pytest.raises(ValueError):
+        meta.verify_recovery_key(KEY2_ID, KEY2_RECOVERY_KEY)
+
+
+def test_verify_recovery_key_broken_iv() -> None:
+    meta = get_key_meta(KEY2_META_BROKEN_IV)
+    with pytest.raises(ValueError):
+        meta.verify_recovery_key(KEY2_ID, KEY2_RECOVERY_KEY)
+
+
+def test_verify_recovery_key_broken_mac() -> None:
+    meta = get_key_meta(KEY2_META_BROKEN_MAC)
+    with pytest.raises(ValueError):
+        meta.verify_recovery_key(KEY2_ID, KEY2_RECOVERY_KEY)
+
+
+def test_verify_passphrase_correct() -> None:
+    meta = get_key_meta(KEY1_META)
+    key = meta.verify_passphrase(KEY1_ID, KEY1_PASSPHRASE)
+    assert key.recovery_key == KEY1_RECOVERY_KEY
+
+
+def test_verify_passphrase_incorrect() -> None:
+    meta = get_key_meta(KEY1_META)
+    with pytest.raises(ValueError):
+        meta.verify_passphrase(KEY1_ID, "incorrect horse battery staple")
+
+
+def test_verify_passphrase_notset() -> None:
+    meta = get_key_meta(KEY2_META)
+    with pytest.raises(ValueError):
+        meta.verify_passphrase(KEY2_ID, "hmm")

mautrix/crypto/ssss/types.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2025 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 typing import TYPE_CHECKING
+
+from attr import dataclass
+
+from mautrix.types import EventType, SerializableAttrs, SerializableEnum
+from mautrix.types.event.account_data import account_data_event_content_map
+
+if TYPE_CHECKING:
+    from .key import Key
+
+
+class Algorithm(SerializableEnum):
+    AES_HMAC_SHA2 = "m.secret_storage.v1.aes-hmac-sha2"
+    CURVE25519_AES_SHA2 = "m.secret_storage.v1.curve25519-aes-sha2"
+
+
+class PassphraseAlgorithm(SerializableEnum):
+    PBKDF2 = "m.pbkdf2"
+
+
+@dataclass
+class EncryptedKeyData(SerializableAttrs):
+    ciphertext: str
+    iv: str
+    mac: str
+
+
+@dataclass
+class EncryptedAccountDataEventContent(SerializableAttrs):
+    encrypted: dict[str, EncryptedKeyData]
+
+    def decrypt(self, event_type: str | EventType, key: "Key") -> bytes:
+        try:
+            encrypted_data = self.encrypted[key.id]
+        except KeyError as e:
+            raise ValueError(f"Event not encrypted for provided key") from e
+        return key.decrypt(event_type, encrypted_data)
+
+
+for encrypted_account_data_type in (
+    EventType.CROSS_SIGNING_MASTER,
+    EventType.CROSS_SIGNING_USER_SIGNING,
+    EventType.CROSS_SIGNING_SELF_SIGNING,
+    EventType.MEGOLM_BACKUP_V1,
+):
+    account_data_event_content_map[encrypted_account_data_type] = EncryptedAccountDataEventContent