konstruct/packages/shared/shared/api/llm_keys.py

"""
LLM key CRUD API endpoints for the Konstruct portal.

Endpoints:
  GET    /api/portal/tenants/{tenant_id}/llm-keys           → list BYO keys (redacted)
  POST   /api/portal/tenants/{tenant_id}/llm-keys           → create encrypted key
  DELETE /api/portal/tenants/{tenant_id}/llm-keys/{key_id} → remove key

Security design:
  - The raw API key is NEVER stored — only the Fernet-encrypted ciphertext.
  - The `key_hint` column stores the last 4 characters of the plaintext key
    so the portal can display "...ABCD" without requiring decryption.
  - GET responses include key_hint but never encrypted_key or the plaintext.
  - DELETE verifies the key belongs to the requested tenant_id before removing.

UNIQUE constraint: (tenant_id, provider) — one BYO key per provider per tenant.
If a tenant needs to rotate a key, DELETE the existing row and POST a new one.
"""

from __future__ import annotations

import uuid
from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from shared.api.rbac import PortalCaller, require_tenant_admin
from shared.config import settings
from shared.crypto import KeyEncryptionService
from shared.db import get_session
from shared.models.billing import TenantLlmKey
from shared.models.tenant import Tenant

llm_keys_router = APIRouter(
    prefix="/api/portal/tenants/{tenant_id}/llm-keys",
    tags=["llm-keys"],
)


# ---------------------------------------------------------------------------
# Pydantic schemas
# ---------------------------------------------------------------------------

class LlmKeyCreate(BaseModel):
    """Request body for creating a new BYO API key."""
    provider: str = Field(
        ...,
        min_length=1,
        max_length=100,
        description="LLM provider name (e.g. 'openai', 'anthropic', 'cohere')",
        examples=["openai"],
    )
    label: str = Field(
        ...,
        min_length=1,
        max_length=255,
        description="Human-readable label for the key (e.g. 'Production OpenAI Key')",
        examples=["Production OpenAI Key"],
    )
    api_key: str = Field(
        ...,
        min_length=1,
        description="The raw API key — will be encrypted before storage, never logged",
        examples=["sk-..."],
    )


class LlmKeyResponse(BaseModel):
    """Response body for LLM key operations — never contains raw or encrypted key."""
    id: str
    provider: str
    label: str
    key_hint: str | None
    created_at: datetime

    model_config = {"from_attributes": True}

    @classmethod
    def from_orm_key(cls, key: TenantLlmKey) -> "LlmKeyResponse":
        return cls(
            id=str(key.id),
            provider=key.provider,
            label=key.label,
            key_hint=key.key_hint,
            created_at=key.created_at,
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _get_encryption_service() -> KeyEncryptionService:
    """Return the platform-level KeyEncryptionService from settings."""
    if not settings.platform_encryption_key:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="PLATFORM_ENCRYPTION_KEY not configured",
        )
    return KeyEncryptionService(
        primary_key=settings.platform_encryption_key,
        previous_key=settings.platform_encryption_key_previous,
    )


async def _get_tenant_or_404(tenant_id: uuid.UUID, session: AsyncSession) -> Tenant:
    result = await session.execute(select(Tenant).where(Tenant.id == tenant_id))
    tenant = result.scalar_one_or_none()
    if tenant is None:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Tenant not found")
    return tenant


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@llm_keys_router.get("", response_model=list[LlmKeyResponse])
async def list_llm_keys(
    tenant_id: uuid.UUID,
    caller: PortalCaller = Depends(require_tenant_admin),
    session: AsyncSession = Depends(get_session),
) -> list[LlmKeyResponse]:
    """
    List all BYO API keys for a tenant.

    Returns: provider, label, key_hint (last 4 chars), created_at.
    The encrypted_key and any plaintext key are NEVER returned.
    """
    await _get_tenant_or_404(tenant_id, session)

    result = await session.execute(
        select(TenantLlmKey)
        .where(TenantLlmKey.tenant_id == tenant_id)
        .order_by(TenantLlmKey.created_at.desc())
    )
    keys = result.scalars().all()
    return [LlmKeyResponse.from_orm_key(k) for k in keys]


@llm_keys_router.post("", response_model=LlmKeyResponse, status_code=status.HTTP_201_CREATED)
async def create_llm_key(
    tenant_id: uuid.UUID,
    body: LlmKeyCreate,
    caller: PortalCaller = Depends(require_tenant_admin),
    session: AsyncSession = Depends(get_session),
) -> LlmKeyResponse:
    """
    Store a new encrypted BYO API key for a tenant.

    The raw api_key is encrypted with Fernet before storage. The response
    contains only the key_hint (last 4 chars) — the raw key is discarded
    after encryption.

    Returns 409 if a key for the same provider already exists for this tenant.
    Returns 201 with the created key record on success.
    """
    await _get_tenant_or_404(tenant_id, session)

    # Check for existing key for this tenant+provider (UNIQUE constraint guard)
    existing = await session.execute(
        select(TenantLlmKey).where(
            TenantLlmKey.tenant_id == tenant_id,
            TenantLlmKey.provider == body.provider,
        )
    )
    if existing.scalar_one_or_none() is not None:
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail=f"A BYO key for provider '{body.provider}' already exists for this tenant. "
                   "Delete the existing key before adding a new one.",
        )

    # Encrypt the API key — raw key must not be stored or logged
    enc_svc = _get_encryption_service()
    encrypted_key = enc_svc.encrypt(body.api_key)
    key_hint = body.api_key[-4:] if len(body.api_key) >= 4 else body.api_key

    key = TenantLlmKey(
        tenant_id=tenant_id,
        provider=body.provider,
        label=body.label,
        encrypted_key=encrypted_key,
        key_hint=key_hint,
        key_version=1,
    )
    session.add(key)
    await session.commit()
    await session.refresh(key)

    return LlmKeyResponse.from_orm_key(key)


@llm_keys_router.delete("/{key_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_llm_key(
    tenant_id: uuid.UUID,
    key_id: uuid.UUID,
    caller: PortalCaller = Depends(require_tenant_admin),
    session: AsyncSession = Depends(get_session),
) -> None:
    """
    Delete a BYO API key.

    Verifies that the key belongs to the specified tenant_id before deletion
    to prevent cross-tenant key removal. Returns 404 if not found.
    """
    await _get_tenant_or_404(tenant_id, session)

    result = await session.execute(
        select(TenantLlmKey).where(
            TenantLlmKey.id == key_id,
            TenantLlmKey.tenant_id == tenant_id,  # Cross-tenant protection
        )
    )
    key = result.scalar_one_or_none()
    if key is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="LLM key not found or does not belong to this tenant",
        )

    await session.delete(key)
    await session.commit()