From fc516d78b74a5a3b6d0bd07366237513a09c3fbb Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 19 Nov 2021 23:58:17 +0200 Subject: [PATCH] Add encryption support to standalone mode --- maubot/client.py | 14 ++----- maubot/standalone/__main__.py | 56 +++++++++++++++++++++++++-- maubot/standalone/config.py | 1 + maubot/standalone/example-config.yaml | 12 +++++- requirements.txt | 2 +- 5 files changed, 69 insertions(+), 16 deletions(-) diff --git a/maubot/client.py b/maubot/client.py index 01eb0d9..b02ec24 100644 --- a/maubot/client.py +++ b/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] diff --git a/maubot/standalone/__main__.py b/maubot/standalone/__main__.py index ffd1a79..ff34f47 100644 --- a/maubot/standalone/__main__.py +++ b/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: diff --git a/maubot/standalone/config.py b/maubot/standalone/config.py index 8f0a2ee..558bc40 100644 --- a/maubot/standalone/config.py +++ b/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") diff --git a/maubot/standalone/example-config.yaml b/maubot/standalone/example-config.yaml index af671e0..ccabb74 100644 --- a/maubot/standalone/example-config.yaml +++ b/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. diff --git a/requirements.txt b/requirements.txt index 27126f8..78bd391 100644 --- a/requirements.txt +++ b/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