Fix cross-signing and session persistence
This commit is contained in:
Astra Logical
parent
282363bfeb
commit
a28176de69
1 changed files with 73 additions and 32 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue