ai-video-fullstack/backend/services/openai_realtime/auth.py

"""Authentication for the public OpenAI-compatible Realtime API."""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

import settings
from db.models import RealtimeApiKey
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession


API_KEY_PREFIX = "sk-rt-"
CLIENT_SECRET_PREFIX = "ek-rt-"


class RealtimeAuthError(ValueError):
    """Raised when a public Realtime credential cannot be accepted."""


@dataclass(frozen=True)
class RealtimeCredential:
    api_key_id: str
    assistant_id: str | None = None
    session: dict[str, Any] | None = None
    safety_identifier_hash: str | None = None
    ephemeral: bool = False


def _b64encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64decode(value: str) -> bytes:
    padding = "=" * (-len(value) % 4)
    return base64.urlsafe_b64decode(f"{value}{padding}".encode("ascii"))


def _peppered_digest(value: str) -> str:
    return hmac.new(
        settings.REALTIME_TOKEN_SECRET.encode("utf-8"),
        value.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()


def _token_signature(encoded_payload: str) -> str:
    digest = hmac.new(
        settings.REALTIME_TOKEN_SECRET.encode("utf-8"),
        encoded_payload.encode("ascii"),
        hashlib.sha256,
    ).digest()
    return _b64encode(digest)


def hash_safety_identifier(value: str | None) -> str | None:
    normalized = (value or "").strip()
    if not normalized:
        return None
    return hashlib.sha256(normalized.encode("utf-8")).hexdigest()


def create_api_key_value() -> tuple[str, str, str, str]:
    key_id = f"rtkey_{uuid4().hex}"
    random_secret = secrets.token_urlsafe(32)
    value = f"{API_KEY_PREFIX}{key_id[6:18]}.{random_secret}"
    return key_id, value, value[:24], _peppered_digest(value)


def create_client_secret(
    *,
    api_key_id: str,
    assistant_id: str,
    session: dict[str, Any],
    safety_identifier_hash: str | None,
) -> tuple[str, int]:
    now = int(time.time())
    expires_at = now + settings.REALTIME_CLIENT_SECRET_TTL_SECONDS
    payload = {
        "sub": api_key_id,
        "assistant_id": assistant_id,
        "session": session,
        "safety_identifier_hash": safety_identifier_hash,
        "iat": now,
        "exp": expires_at,
        "jti": f"rtcs_{uuid4().hex}",
    }
    encoded = _b64encode(
        json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode(
            "utf-8"
        )
    )
    return f"{CLIENT_SECRET_PREFIX}{encoded}.{_token_signature(encoded)}", expires_at


async def _active_api_key(
    session: AsyncSession,
    *,
    key_id: str | None = None,
    key_prefix: str | None = None,
) -> RealtimeApiKey | None:
    if key_id:
        row = await session.get(RealtimeApiKey, key_id)
    elif key_prefix:
        row = (
            await session.execute(
                select(RealtimeApiKey).where(
                    RealtimeApiKey.key_prefix == key_prefix
                )
            )
        ).scalar_one_or_none()
    else:
        return None
    now = datetime.now(UTC)
    if row is None or row.revoked_at is not None:
        return None
    if row.expires_at is not None and row.expires_at <= now:
        return None
    return row


async def authenticate_bearer(
    session: AsyncSession,
    token: str,
) -> RealtimeCredential:
    if token.startswith(API_KEY_PREFIX):
        prefix = token[:24]
        row = await _active_api_key(session, key_prefix=prefix)
        if row is None or not hmac.compare_digest(
            row.key_hash,
            _peppered_digest(token),
        ):
            raise RealtimeAuthError("Invalid or expired API key")
        row.last_used_at = datetime.now(UTC)
        await session.commit()
        return RealtimeCredential(api_key_id=row.id)

    if not token.startswith(CLIENT_SECRET_PREFIX):
        raise RealtimeAuthError("Unsupported Realtime credential")
    compact = token[len(CLIENT_SECRET_PREFIX) :]
    try:
        encoded, signature = compact.rsplit(".", 1)
    except ValueError as exc:
        raise RealtimeAuthError("Invalid client secret") from exc
    if not hmac.compare_digest(_token_signature(encoded), signature):
        raise RealtimeAuthError("Invalid client secret")
    try:
        payload = json.loads(_b64decode(encoded))
    except (ValueError, json.JSONDecodeError) as exc:
        raise RealtimeAuthError("Invalid client secret") from exc
    if int(payload.get("exp", 0)) < int(time.time()):
        raise RealtimeAuthError("Client secret expired")
    api_key_id = str(payload.get("sub") or "")
    row = await _active_api_key(session, key_id=api_key_id)
    if row is None:
        raise RealtimeAuthError("Parent API key is no longer active")
    row.last_used_at = datetime.now(UTC)
    await session.commit()
    return RealtimeCredential(
        api_key_id=api_key_id,
        assistant_id=str(payload.get("assistant_id") or "") or None,
        session=(payload.get("session") if isinstance(payload.get("session"), dict) else None),
        safety_identifier_hash=str(payload.get("safety_identifier_hash") or "") or None,
        ephemeral=True,
    )


def bearer_from_authorization(value: str | None) -> str:
    scheme, _, token = (value or "").partition(" ")
    if scheme.lower() != "bearer" or not token.strip():
        raise RealtimeAuthError("Missing Bearer credential")
    return token.strip()