108144696b
- Per-user Fernet encryption for fact/chunk_text/summary fields - Postgres RLS with memory_app restricted role - SSL for memory-db connections - Data migration script (migrate_encrypt.py) - DB migration (migrate_rls.sql) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""MAT-107: One-time migration to encrypt existing plaintext memory data.
|
|
|
|
Run INSIDE the memory-service container after deploying new code:
|
|
docker exec -it matrix-ai-agent-memory-service-1 python migrate_encrypt.py
|
|
|
|
Connects as owner (memory) to bypass RLS.
|
|
"""
|
|
import os
|
|
import sys
|
|
import hashlib
|
|
import base64
|
|
import asyncio
|
|
|
|
import asyncpg
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
|
|
OWNER_DSN = os.environ.get(
|
|
"OWNER_DATABASE_URL",
|
|
"postgresql://memory:{password}@memory-db:5432/memories".format(
|
|
password=os.environ.get("MEMORY_DB_OWNER_PASSWORD", "memory")
|
|
),
|
|
)
|
|
ENCRYPTION_KEY = os.environ.get("MEMORY_ENCRYPTION_KEY", "")
|
|
|
|
|
|
def _derive_user_key(user_id: str) -> bytes:
|
|
derived = hashlib.pbkdf2_hmac("sha256", ENCRYPTION_KEY.encode(), user_id.encode(), iterations=1)
|
|
return base64.urlsafe_b64encode(derived)
|
|
|
|
|
|
def _encrypt(text: str, user_id: str) -> str:
|
|
f = Fernet(_derive_user_key(user_id))
|
|
return f.encrypt(text.encode()).decode()
|
|
|
|
|
|
def _is_encrypted(text: str, user_id: str) -> bool:
|
|
"""Check if text is already Fernet-encrypted."""
|
|
try:
|
|
f = Fernet(_derive_user_key(user_id))
|
|
f.decrypt(text.encode())
|
|
return True
|
|
except (InvalidToken, Exception):
|
|
return False
|
|
|
|
|
|
async def migrate():
|
|
if not ENCRYPTION_KEY:
|
|
print("ERROR: MEMORY_ENCRYPTION_KEY not set")
|
|
sys.exit(1)
|
|
|
|
conn = await asyncpg.connect(OWNER_DSN)
|
|
|
|
# Migrate memories
|
|
rows = await conn.fetch("SELECT id, user_id, fact FROM memories ORDER BY id")
|
|
print(f"Migrating {len(rows)} memories...")
|
|
encrypted = 0
|
|
skipped = 0
|
|
for row in rows:
|
|
if _is_encrypted(row["fact"], row["user_id"]):
|
|
skipped += 1
|
|
continue
|
|
enc_fact = _encrypt(row["fact"], row["user_id"])
|
|
await conn.execute("UPDATE memories SET fact = $1 WHERE id = $2", enc_fact, row["id"])
|
|
encrypted += 1
|
|
if encrypted % 100 == 0:
|
|
print(f" memories: {encrypted}/{len(rows)} encrypted")
|
|
print(f"Memories done: {encrypted} encrypted, {skipped} already encrypted")
|
|
|
|
# Migrate conversation_chunks
|
|
rows = await conn.fetch("SELECT id, user_id, chunk_text, summary FROM conversation_chunks ORDER BY id")
|
|
print(f"Migrating {len(rows)} chunks...")
|
|
encrypted = 0
|
|
skipped = 0
|
|
for row in rows:
|
|
if _is_encrypted(row["chunk_text"], row["user_id"]):
|
|
skipped += 1
|
|
continue
|
|
enc_text = _encrypt(row["chunk_text"], row["user_id"])
|
|
enc_summary = _encrypt(row["summary"], row["user_id"])
|
|
await conn.execute(
|
|
"UPDATE conversation_chunks SET chunk_text = $1, summary = $2 WHERE id = $3",
|
|
enc_text, enc_summary, row["id"],
|
|
)
|
|
encrypted += 1
|
|
if encrypted % 500 == 0:
|
|
print(f" chunks: {encrypted}/{len(rows)} encrypted")
|
|
print(f"Chunks done: {encrypted} encrypted, {skipped} already encrypted")
|
|
|
|
await conn.close()
|
|
print("Migration complete.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(migrate())
|