Add encryption support to standalone mode

maubot/client.py

@@ -31,21 +31,15 @@ from .db import DBClient
 from .matrix import MaubotMatrixClient
 
 try:
-    from mautrix.crypto import OlmMachine, StateStore as CryptoStateStore, CryptoStore
-
+    from mautrix.crypto import OlmMachine, StateStore as CryptoStateStore, PgCryptoStore
+    from mautrix.util.async_db import Database as AsyncDatabase
 
     class SQLStateStore(BaseSQLStateStore, CryptoStateStore):
         pass
 except ImportError as e:
-    OlmMachine = CryptoStateStore = CryptoStore = PickleCryptoStore = None
+    OlmMachine = CryptoStateStore = PgCryptoStore = AsyncDatabase = None
     SQLStateStore = BaseSQLStateStore
 
-try:
-    from mautrix.util.async_db import Database as AsyncDatabase
-    from mautrix.crypto import PgCryptoStore
-except ImportError:
-    AsyncDatabase = None
-    PgCryptoStore = None
 
 if TYPE_CHECKING:
     from .instance import PluginInstance
@@ -66,7 +60,7 @@ class Client:
     db_instance: DBClient
     client: MaubotMatrixClient
     crypto: Optional['OlmMachine']
-    crypto_store: Optional['CryptoStore']
+    crypto_store: Optional['PgCryptoStore']
     started: bool
 
     remote_displayname: Optional[str]

maubot/standalone/__main__.py

@@ -43,6 +43,15 @@ from .config import Config
 from .loader import FileSystemLoader
 from .database import NextBatch
 
+crypto_import_error = None
+
+try:
+    from mautrix.crypto import OlmMachine, PgCryptoStore, PgCryptoStateStore
+    from mautrix.util.async_db import Database as AsyncDatabase
+except ImportError as err:
+    crypto_import_error = err
+    OlmMachine = AsyncDatabase = PgCryptoStateStore = PgCryptoStore = None
+
 parser = argparse.ArgumentParser(
     description="A plugin-based Matrix bot system -- standalone mode.",
     prog="python -m maubot.standalone")
@@ -92,9 +101,19 @@ Base.metadata.create_all()
 NextBatch.bind(db)
 
 user_id = config["user.credentials.id"]
+device_id = config["user.credentials.device_id"]
 homeserver = config["user.credentials.homeserver"]
 access_token = config["user.credentials.access_token"]
 
+crypto_store = crypto_db = state_store = None
+if device_id and not OlmMachine:
+    log.warning("device_id set in config, but encryption dependencies not installed",
+                exc_info=crypto_import_error)
+elif device_id:
+    crypto_db = AsyncDatabase.create(config["database"], upgrade_table=PgCryptoStore.upgrade_table)
+    crypto_store = PgCryptoStore(account_id=user_id, pickle_key="mau.crypto", db=crypto_db)
+    state_store = PgCryptoStateStore(crypto_db)
+
 nb = NextBatch.get(user_id)
 if not nb:
     nb = NextBatch(user_id=user_id, next_batch=SyncToken(""), filter_id=FilterID(""))
@@ -140,7 +159,25 @@ async def main():
     client_log = logging.getLogger("maubot.client").getChild(user_id)
     client = MaubotMatrixClient(mxid=user_id, base_url=homeserver, token=access_token,
                                 client_session=http_client, loop=loop, log=client_log,
-                                sync_store=SyncStoreProxy(nb))
+                                sync_store=SyncStoreProxy(nb), state_store=state_store,
+                                device_id=device_id)
+    client.ignore_first_sync = config["user.ignore_first_sync"]
+    client.ignore_initial_sync = config["user.ignore_initial_sync"]
+    if crypto_store:
+        await crypto_db.start()
+        await state_store.upgrade_table.upgrade(crypto_db)
+        await crypto_store.open()
+
+        client.crypto = OlmMachine(client, crypto_store, state_store)
+        crypto_device_id = await crypto_store.get_device_id()
+        if crypto_device_id and crypto_device_id != device_id:
+            log.fatal("Mismatching device ID in crypto store and config "
+                      f"(store: {crypto_device_id}, config: {device_id})")
+            sys.exit(10)
+        await client.crypto.load()
+        if not crypto_device_id:
+            await crypto_store.put_device_id(device_id)
+        log.debug("Enabled encryption support")
 
     while True:
         try:
@@ -151,7 +188,12 @@ async def main():
             continue
         if whoami.user_id != user_id:
             log.fatal(f"User ID mismatch: configured {user_id}, but server said {whoami.user_id}")
-            sys.exit(1)
+            sys.exit(11)
+        elif whoami.device_id and device_id and whoami.device_id != device_id:
+            log.fatal(f"Device ID mismatch: configured {device_id}, "
+                      f"but server said {whoami.device_id}")
+            sys.exit(12)
+        log.debug(f"Confirmed connection as {whoami.user_id} / {whoami.device_id}")
         break
 
     if config["user.sync"]:
@@ -183,6 +225,13 @@ async def main():
     await bot.internal_start()
 
 
+async def stop() -> None:
+    client.stop()
+    await bot.internal_stop()
+    if crypto_db:
+        await crypto_db.stop()
+
+
 try:
     log.info("Starting plugin")
     loop.run_until_complete(main())
@@ -198,8 +247,7 @@ try:
     loop.run_forever()
 except KeyboardInterrupt:
     log.info("Interrupt received, stopping")
-    client.stop()
-    loop.run_until_complete(bot.internal_stop())
+    loop.run_until_complete(stop())
     loop.close()
     sys.exit(0)
 except Exception:

maubot/standalone/config.py

@@ -31,6 +31,7 @@ class Config(BaseFileConfig):
         copy("user.credentials.id")
         copy("user.credentials.homeserver")
         copy("user.credentials.access_token")
+        copy("user.credentials.device_id")
         copy("user.sync")
         copy("user.autojoin")
         copy("user.displayname")

maubot/standalone/example-config.yaml

@@ -4,15 +4,25 @@ user:
         id: "@bot:example.com"
         homeserver: https://example.com
         access_token: foo
+        # If you want to enable encryption, set the device ID corresponding to the access token here.
+        device_id: null
     # Enable /sync? This is not needed for purely unencrypted webhook-based bots, but is necessary in most other cases.
     sync: true
     # Automatically accept invites?
     autojoin: false
     # The displayname and avatar URL to set for the bot on startup.
+    # Set to "disable" to not change the the current displayname/avatar.
     displayname: Standalone Bot
     avatar_url: mxc://maunium.net/AKwRzQkTbggfVZGEqexbYLIO
 
-# The database for the plugin. Also used to store the sync token.
+    # Should events from the initial sync be ignored? This should usually always be true.
+    ignore_initial_sync: true
+    # Should events from the first sync after starting be ignored? This can be set to false
+    # if you want the bot to handle messages that were sent while the bot was down.
+    ignore_first_sync: true
+
+# The database for the plugin. Used for plugin data, the sync token and e2ee data (if enabled).
+# SQLite and Postgres are supported.
 database: sqlite:///bot.db
 
 # Config for the plugin. Refer to the plugin's base-config.yaml to find what (if anything) to put here.

requirements.txt

@@ -1,4 +1,4 @@
-mautrix>=0.11,<0.13
+mautrix>=0.12.1,<0.13
 aiohttp>=3,<4
 yarl>=1,<2
 SQLAlchemy>=1,<1.4