Fix cross-signing and session persistence

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.crypto.store.sqlite import SqliteCryptoStore
-from mautrix.client.state_store import MemoryStateStore
+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):
+class CustomSqliteStateStore(SqliteStateStore):
-    def __init__(self, bot_mxid: str):
+    def __init__(self, db_path: str, bot_mxid: str):
-        super().__init__()
+        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):
+class CustomSqliteCryptoStore(SqliteCryptoStore):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        if not hasattr(self, "_cross_signing_keys"):
-            self._cross_signing_keys = {}
-
     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)
+        val = await super().get_cross_signing_key(user_id, usage)
-        user_dict = self._cross_signing_keys.get(user_id, {})
-        val = user_dict.get(usage_str) or user_dict.get(usage)
         if val is not None and not isinstance(val, CrossSigningKeyWrapper):
             val = CrossSigningKeyWrapper(val)
-            user_dict[usage_str] = val
+            usage_str = usage.value if hasattr(usage, "value") else str(usage)
-            user_dict[usage] = val
+            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
+    # Load session if it exists
-    client = Client(base_url=MATRIX_HOMESERVER, mxid=MATRIX_BOT_USER)
+    session_data = {}
-
+    if os.path.exists(SESSION_FILE):
-    # 2. Attach Custom Memory Crypto & State Stores first so keys align with login
+        try:
-    state_store = CustomMemoryStateStore(client.mxid)
+            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)
+    # Start the crypto machine
-    machine = OlmMachine(client, crypto_store, state_store)
+    await client.crypto.start()
-    await machine.load()
+    
-    client.crypto = machine
+    # 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)
+    # 5. Register Event Handlers
-    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
     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)