Add API calls for SSSS

Author: tulir

mautrix/crypto/base.py

@@ -31,6 +31,7 @@
 from mautrix.util.logging import TraceLogger
 
 from .. import client as cli, crypto
+from .ssss import Machine as SSSSMachine
 
 
 class SignedObject(TypedDict):
@@ -40,6 +41,7 @@ class SignedObject(TypedDict):
 
 class BaseOlmMachine:
     client: cli.Client
+    ssss: SSSSMachine
     log: TraceLogger
     crypto_store: crypto.CryptoStore
     state_store: crypto.StateStore

mautrix/crypto/machine.py

@@ -36,6 +36,7 @@
 from .encrypt_megolm import MegolmEncryptionMachine
 from .key_request import KeyRequestingMachine
 from .key_share import KeySharingMachine
+from .ssss import Machine as SSSSMachine
 from .store import CryptoStore, StateStore
 from .unwedge import OlmUnwedgingMachine
 
@@ -58,6 +59,7 @@ class OlmMachine(
     log: TraceLogger
     crypto_store: CryptoStore
     state_store: StateStore
+    ssss: SSSSMachine
 
     account: Optional[OlmAccount]
 
@@ -70,6 +72,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.client = client
+        self.ssss = SSSSMachine(client)
         self.log = log or logging.getLogger("mau.crypto")
         self.crypto_store = crypto_store
         self.state_store = state_store

mautrix/crypto/ssss/__init__.py

@@ -1,4 +1,5 @@
 from .key import Key, KeyMetadata, PassphraseMetadata
+from .machine import Machine
 from .types import (
     Algorithm,
     EncryptedAccountDataEventContent,

mautrix/crypto/ssss/machine.py

@@ -0,0 +1,65 @@
+# 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 mautrix import client as cli
+from mautrix.errors import MNotFound
+from mautrix.types import EventType, SecretStorageDefaultKeyEventContent
+
+from .key import Key, KeyMetadata
+from .types import EncryptedAccountDataEventContent
+
+
+class Machine:
+    client: cli.Client
+
+    def __init__(self, client: cli.Client) -> None:
+        self.client = client
+
+    async def get_default_key_id(self) -> str | None:
+        try:
+            data = await self.client.get_account_data(EventType.SECRET_STORAGE_DEFAULT_KEY)
+            return SecretStorageDefaultKeyEventContent.deserialize(data).key
+        except (MNotFound, ValueError):
+            return None
+
+    async def set_default_key_id(self, key_id: str) -> None:
+        await self.client.set_account_data(
+            EventType.SECRET_STORAGE_DEFAULT_KEY,
+            SecretStorageDefaultKeyEventContent(key=key_id),
+        )
+
+    async def get_key_data(self, key_id: str) -> KeyMetadata:
+        data = await self.client.get_account_data(f"m.secret_storage.key.{key_id}")
+        return KeyMetadata.deserialize(data)
+
+    async def set_key_data(self, key_id: str, data: KeyMetadata) -> None:
+        await self.client.set_account_data(f"m.secret_storage.key.{key_id}", data)
+
+    async def get_default_key_data(self) -> tuple[str, KeyMetadata]:
+        key_id = await self.get_default_key_id()
+        if not key_id:
+            raise ValueError("No default key ID set")
+        return key_id, await self.get_key_data(key_id)
+
+    async def get_decrypted_account_data(self, event_type: EventType | str, key: Key) -> bytes:
+        data = await self.client.get_account_data(event_type)
+        parsed = EncryptedAccountDataEventContent.deserialize(data)
+        return parsed.decrypt(event_type, key)
+
+    async def set_encrypted_account_data(
+        self, event_type: EventType | str, data: bytes, *keys: Key
+    ) -> None:
+        encrypted_data = {}
+        for key in keys:
+            encrypted_data[key.id] = key.encrypt(event_type, data)
+        await self.client.set_account_data(
+            event_type,
+            EncryptedAccountDataEventContent(encrypted=encrypted_data),
+        )
+
+    async def generate_and_upload_key(self, passphrase: str | None = None) -> Key:
+        key = Key.generate(passphrase)
+        await self.set_key_data(key.id, key.metadata)
+        return key