From a28176de69ccfd33d363f2f7c2b489f084740aa7 Mon Sep 17 00:00:00 2001 From: Astra Logical Date: Thu, 23 Jul 2026 13:28:43 -0400 Subject: [PATCH] Fix cross-signing and session persistence --- matrix_keycloak_bot.py | 105 ++++++++++++++++++++++++++++------------- 1 file changed, 73 insertions(+), 32 deletions(-) diff --git a/matrix_keycloak_bot.py b/matrix_keycloak_bot.py index 1277ea9..00fbaf9 100644 --- a/matrix_keycloak_bot.py +++ b/matrix_keycloak_bot.py @@ -9,8 +9,8 @@ from keycloak.exceptions import KeycloakError from mautrix.client import Client from mautrix.crypto import OlmMachine -from mautrix.crypto.store import MemoryCryptoStore -from mautrix.client.state_store import MemoryStateStore +from mautrix.crypto.store.sqlite import SqliteCryptoStore +from mautrix.client.state_store.sqlite import SqliteStateStore from mautrix.types import ( EventType, MessageType, TextMessageEventContent, Format, Event, RoomID, Membership, RelationType, UserID @@ -30,6 +30,10 @@ KEYCLOAK_REALM = os.getenv("KEYCLOAK_REALM", "master") KEYCLOAK_CLIENT_ID = os.getenv("KEYCLOAK_CLIENT_ID", "nest-matrix-bot") KEYCLOAK_CLIENT_SECRET = os.getenv("KEYCLOAK_CLIENT_SECRET", "your-client-secret") +SESSION_FILE = os.getenv("MATRIX_SESSION_FILE", "bot_session.json") +CRYPTO_DB = os.getenv("MATRIX_CRYPTO_DB", "bot_crypto.db") +STATE_DB = os.getenv("MATRIX_STATE_DB", "bot_state.db") + # In-memory tracking of active requests (Event ID -> Request Info) pending_requests = {} @@ -40,9 +44,9 @@ log = logging.getLogger("keycloak_bot") # ------------------------------------------------------------------------------ # Custom Stores to support find_shared_rooms and fix crypto cross-signing bug # ------------------------------------------------------------------------------ -class CustomMemoryStateStore(MemoryStateStore): - def __init__(self, bot_mxid: str): - super().__init__() +class CustomSqliteStateStore(SqliteStateStore): + def __init__(self, db_path: str, bot_mxid: str): + super().__init__(db_path) self.bot_mxid = bot_mxid async def find_shared_rooms(self, user_id: UserID) -> list[RoomID]: @@ -84,13 +88,9 @@ class CrossSigningKeyWrapper: return self.key == str(other) -class CustomMemoryCryptoStore(MemoryCryptoStore): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - if not hasattr(self, "_cross_signing_keys"): - self._cross_signing_keys = {} - +class CustomSqliteCryptoStore(SqliteCryptoStore): async def put_cross_signing_key(self, user_id: UserID, usage, key) -> None: + await super().put_cross_signing_key(user_id, usage, key) usage_str = usage.value if hasattr(usage, "value") else str(usage) if user_id not in self._cross_signing_keys: self._cross_signing_keys[user_id] = {} @@ -99,13 +99,13 @@ class CustomMemoryCryptoStore(MemoryCryptoStore): self._cross_signing_keys[user_id][usage] = wrapped async def get_cross_signing_key(self, user_id: UserID, usage): - usage_str = usage.value if hasattr(usage, "value") else str(usage) - user_dict = self._cross_signing_keys.get(user_id, {}) - val = user_dict.get(usage_str) or user_dict.get(usage) + val = await super().get_cross_signing_key(user_id, usage) if val is not None and not isinstance(val, CrossSigningKeyWrapper): val = CrossSigningKeyWrapper(val) - user_dict[usage_str] = val - user_dict[usage] = val + usage_str = usage.value if hasattr(usage, "value") else str(usage) + if user_id in self._cross_signing_keys: + self._cross_signing_keys[user_id][usage_str] = val + self._cross_signing_keys[user_id][usage] = val return val @@ -365,29 +365,70 @@ async def on_reaction(evt: Event): async def main(): global client - # 1. Initialize client - client = Client(base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER) - - # 2. Attach Custom Memory Crypto & State Stores first so keys align with login - state_store = CustomMemoryStateStore(client.mxid) + # Load session if it exists + session_data = {} + if os.path.exists(SESSION_FILE): + try: + with open(SESSION_FILE, "r") as f: + session_data = json.load(f) + except Exception as e: + log.warning(f"Could not load session file: {e}") + + # Initialize client with persisted token and device ID if available + client = Client( + base_url=MATRIX_HOMESERVER, + mxid=MATRIX_BOT_USER, + access_token=session_data.get("access_token"), + device_id=session_data.get("device_id", "KEYCLOAK_BOT_DEVICE") + ) + + # Initialize stores and crypto machine + state_store = CustomSqliteStateStore(STATE_DB, client.mxid) + crypto_store = CustomSqliteCryptoStore(CRYPTO_DB, client.mxid, client.device_id) client.state_store = state_store + client.crypto = OlmMachine(client, crypto_store, state_store) + + # Authenticate if access token is missing + if not client.access_token: + log.info("Logging in to Matrix homeserver...") + await client.login( + password=MATRIX_BOT_PASSWORD, + device_name="KEYCLOAK_ACCESS_BOT", + device_id="KEYCLOAK_BOT_DEVICE" + ) + # Save session data for future restarts + session_dir = os.path.dirname(SESSION_FILE) + if session_dir: + os.makedirs(session_dir, exist_ok=True) + with open(SESSION_FILE, "w") as f: + json.dump({ + "access_token": client.access_token, + "device_id": client.device_id + }, f) - crypto_store = CustomMemoryCryptoStore(client.mxid, client.device_id) - machine = OlmMachine(client, crypto_store, state_store) - await machine.load() - client.crypto = machine + # Start the crypto machine + await client.crypto.start() + + # Bootstrap Cross-Signing and Key Recovery using the environment variable + recovery_key = os.getenv("MATRIX_RECOVERY_KEY") + if recovery_key and client.crypto: + try: + log.info("Verifying recovery key and bootstrapping cross-signing...") + await client.crypto.verify_recovery_key(recovery_key) + if not await client.crypto.has_cross_signing_keys(): + await client.crypto.bootstrap_cross_signing( + auth_data={"password": MATRIX_BOT_PASSWORD} + ) + log.info("Cross-signing and key recovery initialized successfully.") + except Exception as e: + log.error(f"Failed to initialize cross-signing/recovery: {e}") - # 3. Perform Fresh Login on Startup (Syncs device ID and crypto keys correctly) - log.info(f"Logging in to {MATRIX_HOMESERVER}...") - await client.login(password=MATRIX_BOT_PASSWORD, device_name="KEYCLOAK_ACCESS_BOT") - log.info("Successfully logged in with fresh device session.") - - # 4. Register Event Handlers + # 5. Register Event Handlers client.add_event_handler(EventType.ROOM_MEMBER, on_invite) client.add_event_handler(EventType.ROOM_MESSAGE, on_message) client.add_event_handler(EventType.REACTION, on_reaction) - # 5. Start Sync Loop + # 6. Start Sync Loop log.info("Starting Mautrix sync loop...") await client.start(None)