Switch to pickle for crypto persistence
This commit is contained in:
Astra Logical
parent
a28176de69
commit
247eeaa134
1 changed files with 82 additions and 59 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue