secrets_bootstrap.py

"""
First-boot secrets bootstrap.

Auto-generates JWT_SECRET and ENCRYPTION_KEY if missing and persists them to a
file in the data volume (/app/data/.secrets.env by default). On subsequent
boots the same values are read back, so secrets survive image pulls, stack
restarts, and Portainer redeploys as long as the bind-mounted data volume
survives.

Runs at import time from app_fastapi.py and worker.py so os.environ is
populated before anything else reads these values.
"""

import logging
import os
import secrets as py_secrets
from pathlib import Path
from typing import Dict

from config import CACHE_ROOT

logger = logging.getLogger(__name__)

_PLACEHOLDERS = {
    "",
    "changeme",
    "changeme_generate_fernet_key",
    "changeme_generate_64_random_chars",
    "changeme_strong_password",
    "change_me_in_production_with_random_string",
}

_SECRETS_FILE = Path(CACHE_ROOT) / ".secrets.env"


def _parse_env_file(path: Path) -> Dict[str, str]:
    out: Dict[str, str] = {}
    if not path.exists():
        return out
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        out[k.strip()] = v.strip()
    return out


def _write_env_file(path: Path, values: Dict[str, str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    body = "# Auto-generated by rag on first boot. Keep this file safe.\n"
    body += "\n".join(f"{k}={v}" for k, v in sorted(values.items())) + "\n"
    path.write_text(body, encoding="utf-8")
    try:
        path.chmod(0o600)
    except OSError:
        pass


def bootstrap_secrets() -> None:
    """
    Ensure JWT_SECRET and ENCRYPTION_KEY are set in os.environ.

    Priority: explicit env var (non-placeholder) > persisted file > freshly
    generated value (which is then persisted).
    """
    try:
        from cryptography.fernet import Fernet
    except ImportError:
        logger.warning("cryptography not installed; cannot bootstrap ENCRYPTION_KEY")
        return

    saved = _parse_env_file(_SECRETS_FILE)
    final = dict(saved)
    generated = []

    generators = {
        "JWT_SECRET": lambda: py_secrets.token_hex(32),
        "ENCRYPTION_KEY": lambda: Fernet.generate_key().decode(),
    }

    for name, gen in generators.items():
        current = os.environ.get(name, "").strip()

        # 1. Explicit real value wins
        if current and current not in _PLACEHOLDERS:
            final[name] = current
            continue

        # 2. Persisted real value from the data volume
        if saved.get(name) and saved[name] not in _PLACEHOLDERS:
            os.environ[name] = saved[name]
            continue

        # 3. Generate a fresh value and persist it
        new_val = gen()
        os.environ[name] = new_val
        final[name] = new_val
        generated.append(name)

    if generated:
        _write_env_file(_SECRETS_FILE, final)
        logger.warning(
            "Auto-generated and persisted secrets %s to %s. "
            "Back up this file to survive data-volume loss.",
            generated, _SECRETS_FILE,
        )