""" 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