replypier/backend/app/core/encryption.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Encryption of user content at rest.

Conversation content (chat messages, AI-written session summaries, the agno agent
memory blob) is stored as ciphertext so that a database dump, a stolen backup or a
stray `psql` session does not expose it. The application decrypts transparently.

Stored format is ``enc:<key_id>:<fernet-token>``. The key id is what makes key
rotation possible without a schema change: a future key is added to KEYS as ``v2``,
CURRENT_KEY_ID moves to it, and ``v1`` stays in the registry so old rows keep
reading while a backfill re-encrypts them.

This module deliberately depends on nothing but the standard library, cryptography,
SQLAlchemy and app.core.{logger,config} (both light), so the backfill script can load
it by file path on the host without importing the `app` package — whose __init__
pulls in the FastAPI app and the whole ML stack.
"""

import base64
import json
import os
from functools import lru_cache
from typing import Any, Dict, Optional

from cryptography.fernet import Fernet
from sqlalchemy import Text
from sqlalchemy.types import TypeDecorator

from app.core.config import ENCRYPTION_KEY_HINT, is_throwaway_environment
from app.core.logger import get_logger

logger = get_logger(__name__)

# Wire format. Nothing outside this module should restate these.
ENCRYPTED_PREFIX = "enc"
SEPARATOR = ":"
CURRENT_KEY_ID = "v1"
# Key under which an encrypted JSON blob is stored, so an encrypted dict is
# distinguishable from a plaintext one written before this change.
JSON_MARKER = "__enc__"

ENCRYPTION_KEY_ENV = "ENCRYPTION_KEY"


# Indented for the multi-line error messages below; the command itself is defined
# once in app.core.config so this and the startup audit cannot drift apart.
GENERATE_KEY_HINT = f"  {ENCRYPTION_KEY_HINT}"


def load_key() -> bytes:
    """Return the raw Fernet key from ``ENCRYPTION_KEY``.

    The env var holds the Fernet key base64-encoded a second time, which is the
    format already in use for API-key encryption — decoding must stay identical or
    every existing encrypted API key becomes unreadable.

    Outside development and tests a missing or malformed key is fatal: chat
    messages, session summaries and agent memory are all encrypted with it.

    Usability of the key is checked here rather than left to ``Fernet()`` in
    get_keys(), so that the .env.example placeholder and the common near-miss of
    pasting a raw ``Fernet.generate_key()`` (base64-encoded once, not twice) both
    surface the message below instead of a bare "must be 32 url-safe base64-encoded
    bytes" from deep inside the library.
    """
    env_key = os.getenv(ENCRYPTION_KEY_ENV)
    if env_key:
        try:
            key = base64.b64decode(env_key)
            Fernet(key)
            return key
        except Exception as e:
            logger.error(f"Invalid encryption key format: {str(e)}")
            if not is_throwaway_environment():
                raise RuntimeError(
                    f"{ENCRYPTION_KEY_ENV} is set but is not a usable Fernet key. "
                    "Encrypted data cannot be read; refusing to start with a "
                    f"different key. Generate one with:\n{GENERATE_KEY_HINT}"
                ) from e

    if not is_throwaway_environment():
        raise RuntimeError(
            f"{ENCRYPTION_KEY_ENV} is not set. It encrypts chat messages, session "
            "summaries and agent memory at rest, so starting without it would make "
            f"existing conversations unreadable. Generate one with:\n{GENERATE_KEY_HINT}"
        )

    key = Fernet.generate_key()
    logger.info(
        f"Generated a throwaway encryption key. Set "
        f"{ENCRYPTION_KEY_ENV}={base64.b64encode(key).decode()} to keep data readable"
    )
    return key


@lru_cache(maxsize=1)
def get_keys() -> Dict[str, Fernet]:
    """Every key we can still decrypt with, by key id. Add an id here to rotate.

    Loaded on first use rather than at import: a module import must not depend on
    environment setup, and the explicit check belongs in verify_encryption_key()
    at startup where the failure is legible.
    """
    return {CURRENT_KEY_ID: Fernet(load_key())}


def verify_encryption_key() -> None:
    """Fail fast at startup if the key is missing or unusable.

    Without this the first customer message would be the thing that discovers a
    misconfigured deployment — after the app has already accepted traffic.
    """
    get_keys()  # raises with an actionable message if the key is missing or invalid
    if decrypt_value(encrypt_value("verify")) != "verify":
        raise RuntimeError(f"{ENCRYPTION_KEY_ENV} does not round-trip; refusing to start")


def is_encrypted(value: Any) -> bool:
    """True when the value is already in the stored ciphertext format."""
    return isinstance(value, str) and value.startswith(ENCRYPTED_PREFIX + SEPARATOR)


def encrypt_value(value: Optional[str]) -> Optional[str]:
    """Encrypt a string for storage. Idempotent: already-encrypted input is returned
    untouched, so it is safe to call on a value of unknown provenance."""
    if value is None:
        return None
    if is_encrypted(value):
        return value
    # Outside the try: a missing key raises a RuntimeError explaining itself, and
    # wrapping that in "failed to encrypt" would hide the actual problem.
    fernet = get_keys()[CURRENT_KEY_ID]
    try:
        token = fernet.encrypt(str(value).encode()).decode()
    except Exception as e:
        logger.error(f"Encryption error: {str(e)}")
        raise ValueError("Failed to encrypt value")
    return SEPARATOR.join((ENCRYPTED_PREFIX, CURRENT_KEY_ID, token))


def decrypt_value(value: Optional[str]) -> Optional[str]:
    """Decrypt a stored string.

    Values without the ``enc:`` prefix are returned unchanged — that is a row the
    backfill has not reached yet, and reads must keep working throughout the
    migration. A key we no longer hold, or a corrupt token, logs and returns None
    rather than raising: one damaged row must not fail the whole request.
    """
    if not is_encrypted(value):
        return value

    parts = value.split(SEPARATOR, 2)
    if len(parts) != 3:
        logger.error("Malformed ciphertext: missing key id or token")
        return None
    _, key_id, token = parts

    fernet = get_keys().get(key_id)
    if fernet is None:
        logger.error(f"No encryption key registered for key id '{key_id}'")
        return None

    try:
        return fernet.decrypt(token.encode()).decode()
    except Exception as e:
        logger.error(f"Decryption error: {str(e)}")
        return None


def encrypt_json(value: Optional[Any]) -> Optional[Dict[str, str]]:
    """Encrypt a JSON-serialisable object into a single-key marker dict."""
    if value is None:
        return None
    if isinstance(value, dict) and JSON_MARKER in value:
        return value
    return {JSON_MARKER: encrypt_value(json.dumps(value))}


def decrypt_json(value: Optional[Any]) -> Optional[Any]:
    """Inverse of encrypt_json. A plaintext object written before this change is
    returned as-is; an undecryptable blob degrades to None."""
    if not isinstance(value, dict) or JSON_MARKER not in value:
        return value

    plaintext = decrypt_value(value[JSON_MARKER])
    if plaintext is None:
        return None
    try:
        return json.loads(plaintext)
    except ValueError as e:
        logger.error(f"Decrypted blob is not valid JSON: {str(e)}")
        return None


class EncryptedText(TypeDecorator):
    """Text column whose value is encrypted at rest and decrypted on read.

    Underlying storage stays unbounded text, so applying this to an existing
    String/Text column needs no schema migration.
    """

    impl = Text
    cache_ok = True

    def process_bind_param(self, value, dialect):
        return encrypt_value(value)

    def process_result_value(self, value, dialect):
        return decrypt_value(value)