diff --git a/matrix_keycloak_bot.py b/matrix_keycloak_bot.py index 00fbaf9..8c50481 100644 --- a/matrix_keycloak_bot.py +++ b/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.client.state_store.sqlite import SqliteStateStore +from mautrix.crypto.store import MemoryCryptoStore +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") -CRYPTO_DB = os.getenv("MATRIX_CRYPTO_DB", "bot_crypto.db") -STATE_DB = os.getenv("MATRIX_STATE_DB", "bot_state.db") +SESSION_FILE = os.getenv("MATRIX_SESSION_FILE", "/app/store/bot_session.json") +CRYPTO_STORE_FILE = os.getenv("MATRIX_CRYPTO_STORE", "/app/store/crypto_store.pkl") # 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): - def __init__(self, db_path: str, bot_mxid: str): - super().__init__(db_path) +class PickleCryptoStore(MemoryCryptoStore): + def __init__(self, filepath: str): + 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 - state_store = CustomSqliteStateStore(STATE_DB, client.mxid) - crypto_store = CustomSqliteCryptoStore(CRYPTO_DB, client.mxid, client.device_id) + # Initialize stores using Pickle and Custom Memory State Stores + state_store = CustomMemoryStateStore(client.mxid) + 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)