feat: Matrix E2EE key management + multi-user isolation
- Add rag_key_manager.py: stores encryption key in private E2EE room - Bot loads key from Matrix on startup, injects into RAG via portal proxy - No plaintext key on disk (removed RAG_ENCRYPTION_KEY from .env) - Pass owner_id (matrix_user_id) to RAG search for user isolation - Stronger format_context instructions for source link rendering Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Christian Gick
186
rag_key_manager.py
Normal file
186
rag_key_manager.py
Normal file
@@ -0,0 +1,186 @@
|
||||
"""
|
||||
RAG Document Key Manager — stores per-user encryption keys in Matrix E2EE rooms.
|
||||
|
||||
The key is stored as an encrypted event in a private room that only the bot can access.
|
||||
On startup, the bot syncs the room and re-injects the key into the RAG service
|
||||
via the portal proxy (since RAG service is localhost-only on the customer VM).
|
||||
No plaintext keys are ever written to disk.
|
||||
"""
|
||||
|
||||
import secrets
|
||||
import logging
|
||||
import httpx
|
||||
from nio.api import RoomVisibility
|
||||
|
||||
logger = logging.getLogger("matrix-ai-bot")
|
||||
|
||||
KEY_EVENT_TYPE = "eu.matrixhost.rag_document_key"
|
||||
KEY_ROOM_TOPIC = "RAG Document Encryption Keys \u2014 DO NOT LEAVE"
|
||||
|
||||
|
||||
class RAGKeyManager:
|
||||
"""Manages per-user document encryption keys via Matrix E2EE."""
|
||||
|
||||
def __init__(self, client, portal_url: str, bot_api_key: str):
|
||||
self.client = client
|
||||
self.portal_url = portal_url.rstrip("/") if portal_url else ""
|
||||
self.bot_api_key = bot_api_key
|
||||
self._key_room_id: str | None = None
|
||||
|
||||
async def ensure_rag_key(self, seed_key_hex: str | None = None) -> bool:
|
||||
"""Ensure RAG service has encryption key loaded.
|
||||
|
||||
Args:
|
||||
seed_key_hex: Existing key to migrate into Matrix storage (one-time).
|
||||
"""
|
||||
if not self.portal_url:
|
||||
logger.warning("[rag-key] No portal URL configured")
|
||||
return False
|
||||
|
||||
# Check if RAG already has a key
|
||||
if await self._rag_has_key():
|
||||
logger.info("[rag-key] RAG service already has key loaded")
|
||||
room_id = await self._find_or_create_key_room()
|
||||
if room_id:
|
||||
existing = await self._load_key_from_room(room_id)
|
||||
if not existing and seed_key_hex:
|
||||
await self._store_key_in_room(room_id, seed_key_hex)
|
||||
logger.info("[rag-key] Migrated existing key into Matrix E2EE room")
|
||||
return True
|
||||
|
||||
# Find or create the key storage room
|
||||
room_id = await self._find_or_create_key_room()
|
||||
if not room_id:
|
||||
logger.error("[rag-key] Failed to find or create key room")
|
||||
return False
|
||||
|
||||
# Try to load existing key from room
|
||||
key_hex = await self._load_key_from_room(room_id)
|
||||
|
||||
if key_hex:
|
||||
logger.info("[rag-key] Loaded existing key from Matrix room")
|
||||
elif seed_key_hex:
|
||||
key_hex = seed_key_hex
|
||||
stored = await self._store_key_in_room(room_id, key_hex)
|
||||
if not stored:
|
||||
logger.error("[rag-key] Failed to store seed key in Matrix room")
|
||||
return False
|
||||
logger.info("[rag-key] Stored migration seed key in Matrix E2EE room")
|
||||
else:
|
||||
key_hex = secrets.token_hex(32)
|
||||
stored = await self._store_key_in_room(room_id, key_hex)
|
||||
if not stored:
|
||||
logger.error("[rag-key] Failed to store new key in Matrix room")
|
||||
return False
|
||||
logger.info("[rag-key] Generated and stored new encryption key")
|
||||
|
||||
return await self._inject_key(key_hex)
|
||||
|
||||
async def _rag_has_key(self) -> bool:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.get(
|
||||
f"{self.portal_url}/api/bot/rag-key",
|
||||
headers={"Authorization": f"Bearer {self.bot_api_key}"},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json().get("has_key", False)
|
||||
except Exception as e:
|
||||
logger.debug("[rag-key] Health check failed: %s", e)
|
||||
return False
|
||||
|
||||
async def _find_or_create_key_room(self) -> str | None:
|
||||
for room_id, room in self.client.rooms.items():
|
||||
if room.topic == KEY_ROOM_TOPIC:
|
||||
self._key_room_id = room_id
|
||||
logger.info("[rag-key] Found existing key room: %s", room_id)
|
||||
return room_id
|
||||
|
||||
try:
|
||||
from nio import EnableEncryptionBuilder
|
||||
initial_state = [EnableEncryptionBuilder().as_dict()]
|
||||
except ImportError:
|
||||
initial_state = [{
|
||||
"type": "m.room.encryption",
|
||||
"state_key": "",
|
||||
"content": {"algorithm": "m.megolm.v1.aes-sha2"},
|
||||
}]
|
||||
|
||||
resp = await self.client.room_create(
|
||||
name="RAG Key Storage",
|
||||
topic=KEY_ROOM_TOPIC,
|
||||
invite=[],
|
||||
initial_state=initial_state,
|
||||
visibility=RoomVisibility.private,
|
||||
)
|
||||
|
||||
if hasattr(resp, "room_id"):
|
||||
self._key_room_id = resp.room_id
|
||||
logger.info("[rag-key] Created new key room: %s", resp.room_id)
|
||||
return resp.room_id
|
||||
|
||||
logger.error("[rag-key] Failed to create key room: %s", resp)
|
||||
return None
|
||||
|
||||
async def _load_key_from_room(self, room_id: str) -> str | None:
|
||||
try:
|
||||
resp = await self.client.room_messages(
|
||||
room_id, start="", limit=50, direction="b",
|
||||
)
|
||||
if not hasattr(resp, "chunk"):
|
||||
return None
|
||||
|
||||
for event in resp.chunk:
|
||||
if hasattr(event, "source"):
|
||||
source = event.source
|
||||
if source.get("type") == KEY_EVENT_TYPE:
|
||||
key = source.get("content", {}).get("key_hex")
|
||||
if key:
|
||||
return key
|
||||
|
||||
if hasattr(event, "type") and event.type == KEY_EVENT_TYPE:
|
||||
if hasattr(event, "content"):
|
||||
key = event.content.get("key_hex")
|
||||
if key:
|
||||
return key
|
||||
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("[rag-key] Failed to load key from room: %s", e)
|
||||
return None
|
||||
|
||||
async def _store_key_in_room(self, room_id: str, key_hex: str) -> bool:
|
||||
try:
|
||||
content = {
|
||||
"key_hex": key_hex,
|
||||
"algorithm": "aes-256-gcm",
|
||||
"purpose": "rag-document-encryption",
|
||||
"msgtype": "eu.matrixhost.rag_key",
|
||||
}
|
||||
resp = await self.client.room_send(
|
||||
room_id, message_type=KEY_EVENT_TYPE,
|
||||
content=content, ignore_unverified_devices=True,
|
||||
)
|
||||
if hasattr(resp, "event_id"):
|
||||
logger.info("[rag-key] Key stored as event %s", resp.event_id)
|
||||
return True
|
||||
logger.error("[rag-key] Failed to send key event: %s", resp)
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error("[rag-key] Failed to store key: %s", e)
|
||||
return False
|
||||
|
||||
async def _inject_key(self, key_hex: str) -> bool:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
resp = await client.post(
|
||||
f"{self.portal_url}/api/bot/rag-key",
|
||||
json={"key_hex": key_hex},
|
||||
headers={"Authorization": f"Bearer {self.bot_api_key}"},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
logger.info("[rag-key] Key injected into RAG service via portal proxy")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("[rag-key] Failed to inject key: %s", e)
|
||||
return False
|
||||
Reference in New Issue
Block a user