feat: Add E2E encryption support to Matrix bot
- matrix-nio[e2e] with libolm for Megolm encryption - Persistent crypto store volume for key persistence - Auto-accept key verification (SAS) - Upload device keys on first login CF-1147 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
57
bot.py
57
bot.py
@@ -2,7 +2,17 @@ import os
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from nio import AsyncClient, LoginResponse, InviteMemberEvent
|
||||
from nio import (
|
||||
AsyncClient,
|
||||
AsyncClientConfig,
|
||||
LoginResponse,
|
||||
InviteMemberEvent,
|
||||
KeyVerificationStart,
|
||||
KeyVerificationCancel,
|
||||
KeyVerificationKey,
|
||||
KeyVerificationMac,
|
||||
ToDeviceError,
|
||||
)
|
||||
from livekit import api
|
||||
|
||||
logger = logging.getLogger("matrix-ai-bot")
|
||||
@@ -14,11 +24,23 @@ LK_URL = os.environ["LIVEKIT_URL"]
|
||||
LK_KEY = os.environ["LIVEKIT_API_KEY"]
|
||||
LK_SECRET = os.environ["LIVEKIT_API_SECRET"]
|
||||
AGENT_NAME = os.environ.get("AGENT_NAME", "matrix-ai")
|
||||
STORE_PATH = os.environ.get("CRYPTO_STORE_PATH", "/data/crypto_store")
|
||||
|
||||
|
||||
class Bot:
|
||||
def __init__(self):
|
||||
self.client = AsyncClient(HOMESERVER, BOT_USER)
|
||||
config = AsyncClientConfig(
|
||||
max_limit_exceeded=0,
|
||||
max_timeouts=0,
|
||||
store_sync_tokens=True,
|
||||
encryption_enabled=True,
|
||||
)
|
||||
self.client = AsyncClient(
|
||||
HOMESERVER,
|
||||
BOT_USER,
|
||||
store_path=STORE_PATH,
|
||||
config=config,
|
||||
)
|
||||
self.lkapi = None
|
||||
self.dispatched_rooms = set()
|
||||
|
||||
@@ -27,12 +49,20 @@ class Bot:
|
||||
if not isinstance(resp, LoginResponse):
|
||||
logger.error("Login failed: %s", resp)
|
||||
return
|
||||
logger.info("Logged in as %s", resp.user_id)
|
||||
logger.info("Logged in as %s (device %s)", resp.user_id, resp.device_id)
|
||||
|
||||
# Trust our own device keys
|
||||
if self.client.should_upload_keys:
|
||||
await self.client.keys_upload()
|
||||
|
||||
self.lkapi = api.LiveKitAPI(LK_URL, LK_KEY, LK_SECRET)
|
||||
self.client.add_event_callback(self.on_invite, InviteMemberEvent)
|
||||
self.client.add_to_device_callback(self.on_key_verification, KeyVerificationStart)
|
||||
self.client.add_to_device_callback(self.on_key_verification, KeyVerificationKey)
|
||||
self.client.add_to_device_callback(self.on_key_verification, KeyVerificationMac)
|
||||
self.client.add_to_device_callback(self.on_key_verification, KeyVerificationCancel)
|
||||
|
||||
await self.client.sync_forever(timeout=30000)
|
||||
await self.client.sync_forever(timeout=30000, full_state=True)
|
||||
|
||||
async def on_invite(self, room, event: InviteMemberEvent):
|
||||
if event.state_key != BOT_USER:
|
||||
@@ -54,6 +84,24 @@ class Bot:
|
||||
except Exception:
|
||||
logger.exception("Dispatch failed for %s", room.room_id)
|
||||
|
||||
async def on_key_verification(self, event):
|
||||
"""Auto-accept key verification requests."""
|
||||
if isinstance(event, KeyVerificationStart):
|
||||
sas = self.client.key_verifications.get(event.transaction_id)
|
||||
if sas:
|
||||
await self.client.accept_key_verification(event.transaction_id)
|
||||
await self.client.to_device(sas.share_key())
|
||||
elif isinstance(event, KeyVerificationKey):
|
||||
sas = self.client.key_verifications.get(event.transaction_id)
|
||||
if sas:
|
||||
await self.client.confirm_short_auth_string(event.transaction_id)
|
||||
elif isinstance(event, KeyVerificationMac):
|
||||
sas = self.client.key_verifications.get(event.transaction_id)
|
||||
if sas:
|
||||
mac = sas.get_mac()
|
||||
if not isinstance(mac, ToDeviceError):
|
||||
await self.client.to_device(mac)
|
||||
|
||||
async def cleanup(self):
|
||||
await self.client.close()
|
||||
if self.lkapi:
|
||||
@@ -61,6 +109,7 @@ class Bot:
|
||||
|
||||
|
||||
async def main():
|
||||
os.makedirs(STORE_PATH, exist_ok=True)
|
||||
bot = Bot()
|
||||
try:
|
||||
await bot.start()
|
||||
|
||||
Reference in New Issue
Block a user