Switch to pickle for crypto persistence

matrix_keycloak_bot.py

@@ -3,14 +3,15 @@ import os
 import sys
 import markdown
 import json
+import pickle
 import logging
 from keycloak import KeycloakAdmin
 from keycloak.exceptions import KeycloakError
 
 from mautrix.client import Client
 from mautrix.crypto import OlmMachine
-from mautrix.crypto.store.sqlite import SqliteCryptoStore
+from mautrix.crypto.store import MemoryCryptoStore
-from mautrix.client.state_store.sqlite import SqliteStateStore
+from mautrix.client.state_store import MemoryStateStore
 from mautrix.types import (
     EventType, MessageType, TextMessageEventContent, Format,
     Event, RoomID, Membership, RelationType, UserID
@@ -30,9 +31,8 @@ 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")
+SESSION_FILE = os.getenv("MATRIX_SESSION_FILE", "/app/store/bot_session.json")
-CRYPTO_DB = os.getenv("MATRIX_CRYPTO_DB", "bot_crypto.db")
+CRYPTO_STORE_FILE = os.getenv("MATRIX_CRYPTO_STORE", "/app/store/crypto_store.pkl")
-STATE_DB = os.getenv("MATRIX_STATE_DB", "bot_state.db")
 
 # In-memory tracking of active requests (Event ID -> Request Info)
 pending_requests = {}
@@ -42,11 +42,78 @@ log = logging.getLogger("keycloak_bot")
 
 
 # ------------------------------------------------------------------------------
-# Custom Stores to support find_shared_rooms and fix crypto cross-signing bug
+# Custom Pickle-Backed Crypto Store for Persistence
 # ------------------------------------------------------------------------------
-class CustomSqliteStateStore(SqliteStateStore):
+class PickleCryptoStore(MemoryCryptoStore):
-    def __init__(self, db_path: str, bot_mxid: str):
+    def __init__(self, filepath: str):
-        super().__init__(db_path)
+        super().__init__()
+        self.filepath = filepath
+        self.load()
+
+    def load(self):
+        if os.path.exists(self.filepath):
+            try:
+                with open(self.filepath, "rb") as f:
+                    data = pickle.load(f)
+                    self.accounts = data.get("accounts", {})
+                    self.sessions = data.get("sessions", {})
+                    self.inbound_group_sessions = data.get("inbound_group_sessions", {})
+                    self.outbound_group_sessions = data.get("outbound_group_sessions", {})
+                    self.devices = data.get("devices", {})
+                    self.cross_signing_keys = data.get("cross_signing_keys", {})
+                    log.info(f"Loaded crypto store from {self.filepath}")
+            except Exception as e:
+                log.warning(f"Failed to load pickle crypto store: {e}")
+
+    def save(self):
+        try:
+            data = {
+                "accounts": self.accounts,
+                "sessions": self.sessions,
+                "inbound_group_sessions": self.inbound_group_sessions,
+                "outbound_group_sessions": self.outbound_group_sessions,
+                "devices": self.devices,
+                "cross_signing_keys": self.cross_signing_keys,
+            }
+            directory = os.path.dirname(self.filepath)
+            if directory:
+                os.makedirs(directory, exist_ok=True)
+            with open(self.filepath, "wb") as f:
+                pickle.dump(data, f)
+        except Exception as e:
+            log.error(f"Failed to save pickle crypto store: {e}")
+
+    async def put_account(self, *args, **kwargs):
+        await super().put_account(*args, **kwargs)
+        self.save()
+
+    async def put_session(self, *args, **kwargs):
+        await super().put_session(*args, **kwargs)
+        self.save()
+
+    async def put_inbound_group_session(self, *args, **kwargs):
+        await super().put_inbound_group_session(*args, **kwargs)
+        self.save()
+
+    async def put_outbound_group_session(self, *args, **kwargs):
+        await super().put_outbound_group_session(*args, **kwargs)
+        self.save()
+
+    async def put_device(self, *args, **kwargs):
+        await super().put_device(*args, **kwargs)
+        self.save()
+
+    async def put_cross_signing_key(self, user_id, usage, key):
+        await super().put_cross_signing_key(user_id, usage, key)
+        self.save()
+
+
+# ------------------------------------------------------------------------------
+# Custom State Store supporting find_shared_rooms
+# ------------------------------------------------------------------------------
+class CustomMemoryStateStore(MemoryStateStore):
+    def __init__(self, bot_mxid: str):
+        super().__init__()
         self.bot_mxid = bot_mxid
 
     async def find_shared_rooms(self, user_id: UserID) -> list[RoomID]:
@@ -64,51 +131,6 @@ class CustomSqliteStateStore(SqliteStateStore):
         return shared
 
 
-class CrossSigningKeyWrapper:
-    def __init__(self, key_val):
-        if hasattr(key_val, "key"):
-            self.key = key_val.key
-            self.first = getattr(key_val, "first", self.key)
-        elif isinstance(key_val, str):
-            self.key = key_val
-            self.first = key_val
-        else:
-            self.key = str(key_val)
-            self.first = getattr(key_val, "first", self.key)
-
-    def __str__(self):
-        return self.key
-
-    def __repr__(self):
-        return f"CrossSigningKeyWrapper(key={self.key!r}, first={self.first!r})"
-
-    def __eq__(self, other):
-        if hasattr(other, "key"):
-            return self.key == other.key
-        return self.key == str(other)
-
-
-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] = {}
-        wrapped = CrossSigningKeyWrapper(key)
-        self._cross_signing_keys[user_id][usage_str] = wrapped
-        self._cross_signing_keys[user_id][usage] = wrapped
-
-    async def get_cross_signing_key(self, user_id: UserID, usage):
-        val = await super().get_cross_signing_key(user_id, usage)
-        if val is not None and not isinstance(val, CrossSigningKeyWrapper):
-            val = CrossSigningKeyWrapper(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
-
-
 # ------------------------------------------------------------------------------
 # Keycloak Helper Functions
 # ------------------------------------------------------------------------------
@@ -382,9 +404,10 @@ async def main():
         device_id=session_data.get("device_id", "KEYCLOAK_BOT_DEVICE")
     )
     
-    # Initialize stores and crypto machine
+    # Initialize stores using Pickle and Custom Memory State Stores
-    state_store = CustomSqliteStateStore(STATE_DB, client.mxid)
+    state_store = CustomMemoryStateStore(client.mxid)
-    crypto_store = CustomSqliteCryptoStore(CRYPTO_DB, client.mxid, client.device_id)
+    crypto_store = PickleCryptoStore(CRYPTO_STORE_FILE)
+    
     client.state_store = state_store
     client.crypto = OlmMachine(client, crypto_store, state_store)
     
@@ -423,12 +446,12 @@ async def main():
         except Exception as e:
             log.error(f"Failed to initialize cross-signing/recovery: {e}")
 
-    # 5. Register Event Handlers
+    # 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)
 
-    # 6. Start Sync Loop
+    # Start Sync Loop
     log.info("Starting Mautrix sync loop...")
     await client.start(None)