kite/backend/services/ai_config.py

"""
AI provider configuration service for DocuVault.

Provides HKDF/Fernet encryption helpers, a provider config loader that reads
from the system_settings DB table, and a startup seed function that populates
the default provider row from env vars on first boot.

Security design (D-05, T-07-02):
  HKDF domain separation — the info bytes b"ai-provider-settings" differ from
  b"cloud-credentials" used by storage/cloud_utils.py. Both use the same master
  key (settings.cloud_creds_key) but produce DIFFERENT derived Fernet keys, so
  a leaked cloud credential cannot decrypt an AI API key and vice versa.

AlreadyFinalized warning (RESEARCH.md Pitfall 2 / .continue-here.md anti-pattern):
  The cryptography library raises AlreadyFinalized if .derive() is called twice
  on the same HKDF instance. _derive_ai_settings_key() creates a FRESH HKDF(...)
  object on every call — never cache or reuse the HKDF object between calls.

Pattern reference: storage/cloud_utils.py:_derive_fernet_key().
"""
from __future__ import annotations

import base64
import logging
from typing import Optional

import structlog
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ai.provider_config import ProviderConfig, PROVIDER_DEFAULTS
from config import settings

logger = structlog.get_logger(__name__)


# ── HKDF key derivation ───────────────────────────────────────────────────────

def _derive_ai_settings_key(master_key: bytes, provider_id: str) -> Fernet:
    """Derive a per-provider Fernet encryption key using HKDF-SHA256.

    Security notes:
      - A FRESH HKDF instance is created on every call. The cryptography library
        raises AlreadyFinalized if .derive() is called twice on the same instance.
        Never cache or reuse the HKDF object (RESEARCH.md Pitfall 2).
      - salt = provider_id.encode("utf-8") provides per-provider derivation
        (deterministic: same provider → same key for encrypt/decrypt consistency).
      - info = b"ai-provider-settings" provides domain separation from
        b"cloud-credentials" — same master key, different derived keys.
        A leaked cloud credential cannot decrypt an AI API key (T-07-02 mitigated).

    Args:
        master_key: The CLOUD_CREDS_KEY env var as bytes.
        provider_id: The provider slug, e.g. "openai", "anthropic" (used as HKDF salt).

    Returns:
        A Fernet instance ready for encrypt/decrypt operations.
    """
    # Create a FRESH HKDF instance — never cache (AlreadyFinalized guard)
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=provider_id.encode("utf-8"),
        info=b"ai-provider-settings",  # domain-separated from b"cloud-credentials"
    )
    raw_key: bytes = hkdf.derive(master_key)
    fernet_key = base64.urlsafe_b64encode(raw_key)
    return Fernet(fernet_key)


# ── Encryption helpers ────────────────────────────────────────────────────────

def encrypt_api_key(master_key: bytes, provider_id: str, api_key: str) -> str:
    """Encrypt a plaintext API key string to a Fernet token.

    The returned string is safe to store in system_settings.api_key_enc.
    No JSON wrapping — the raw API key string is encrypted directly.

    Args:
        master_key: The CLOUD_CREDS_KEY env var as bytes.
        provider_id: The provider slug (used as HKDF salt for key derivation).
        api_key: The plaintext API key, e.g. "sk-proj-...".

    Returns:
        A URL-safe base64 Fernet token (str).
    """
    f = _derive_ai_settings_key(master_key, provider_id)
    return f.encrypt(api_key.encode("utf-8")).decode("utf-8")


def decrypt_api_key(master_key: bytes, provider_id: str, api_key_enc: str) -> str:
    """Decrypt a Fernet token back to the original plaintext API key.

    Args:
        master_key: The CLOUD_CREDS_KEY env var as bytes.
        provider_id: The provider slug (used as HKDF salt for key derivation).
        api_key_enc: The Fernet token string from the database.

    Returns:
        The original plaintext API key string.
    """
    f = _derive_ai_settings_key(master_key, provider_id)
    return f.decrypt(api_key_enc.encode("utf-8")).decode("utf-8")


# ── Provider config loader ────────────────────────────────────────────────────

async def load_provider_config(session: AsyncSession) -> Optional[ProviderConfig]:
    """Load the active AI provider config from the system_settings table.

    Returns a ProviderConfig built from the row where is_active=True,
    decrypting api_key_enc when present.  Returns None when no active row exists.

    Args:
        session: An open AsyncSession.

    Returns:
        A ProviderConfig if an active row exists; None if the table is empty or
        no row is marked active.
    """
    from db.models import SystemSettings  # local import to avoid circular deps

    stmt = select(SystemSettings).where(SystemSettings.is_active.is_(True))
    result = await session.execute(stmt)
    row = result.scalar_one_or_none()

    if row is None:
        return None

    # Decrypt API key if present
    api_key = ""
    if row.api_key_enc:
        master_key = settings.cloud_creds_key.encode("utf-8")
        try:
            api_key = decrypt_api_key(master_key, row.provider_id, row.api_key_enc)
        except Exception:
            logger.warning(
                "ai_config.load_provider_config: failed to decrypt api_key_enc",
                provider_id=row.provider_id,
            )

    return ProviderConfig(
        provider_id=row.provider_id,
        api_key=api_key,
        base_url=row.base_url,
        model=row.model_name,
        context_chars=row.context_chars,
    )


# ── Startup seed ──────────────────────────────────────────────────────────────

async def seed_system_settings_from_env(session: AsyncSession) -> None:
    """Populate system_settings with a default provider row on first boot.

    Reads settings.default_ai_provider and settings.default_ai_model from config.
    If no row exists for that provider_id, inserts one with is_active=True.
    Never overwrites an existing row — idempotent across restarts (D-04).

    This function is called from the FastAPI lifespan in main.py after the
    session factory is available.  Caller is responsible for committing.

    Args:
        session: An open AsyncSession.
    """
    from db.models import SystemSettings  # local import to avoid circular deps

    provider_id = settings.default_ai_provider
    model_name = settings.default_ai_model

    stmt = select(SystemSettings).where(SystemSettings.provider_id == provider_id)
    result = await session.execute(stmt)
    existing = result.scalar_one_or_none()

    if existing is not None:
        # Row already exists — never overwrite (idempotent)
        return

    # Use PROVIDER_DEFAULTS context_chars for the provider, fallback to 8000
    context_chars = PROVIDER_DEFAULTS.get(provider_id, {}).get("context_chars", 8000)

    # Insert default row with no API key (local providers like Ollama don't need one)
    row = SystemSettings(
        provider_id=provider_id,
        model_name=model_name,
        context_chars=context_chars,
        is_active=True,
        api_key_enc=None,
        base_url=None,
    )
    session.add(row)
    logger.info(
        "ai_config.seed_system_settings_from_env: seeded default provider",
        provider_id=provider_id,
        model_name=model_name,
    )