deepcell-cli 0.6.1__py3-none-any.whl

deepcell_cli/oauth_provider.py

@@ -0,0 +1,580 @@
+"""OAuth 2.1 provider for MCP remote authentication.
+
+Implements the ``OAuthAuthorizationServerProvider`` protocol from the MCP SDK
+so that AI platforms (Claude.ai, ChatGPT, Manus) can authenticate users via
+OAuth 2.1 + PKCE before accessing DeepCell tools.
+
+Storage:
+- Client registrations: ``~/.deepcell/oauth_clients.json`` (persistent; the
+  ``deepcell-mcp`` container mounts a volume here so DCR survives redeploys).
+- Authorization codes: in-memory dict with 10-min TTL.
+- Refresh tokens: delegated to the Jingwei API. An MCP refresh token *is* a
+  Jingwei refresh token (opaque, hashed in Postgres), so rotation, token-family
+  reuse-detection and revocation live server-side and survive container
+  recreates — there is no in-memory refresh-token store.
+- Access tokens: JWT issued by Jingwei and verified locally using the same
+  shared secret.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import secrets
+import time
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote, urlparse
+from uuid import UUID
+
+import httpx
+import jwt as pyjwt
+
+from mcp.server.auth.provider import (
+    AccessToken,
+    AuthorizationCode,
+    AuthorizationParams,
+    AuthorizeError,
+    OAuthAuthorizationServerProvider,
+    RefreshToken,
+    TokenError,
+    construct_redirect_uri,
+)
+from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
+
+from deepcell_cli.config import _DIR, _ensure_dir, _write_private_file, client_headers
+
+logger = logging.getLogger(__name__)
+
+_CLIENTS_FILE = _DIR / "oauth_clients.json"
+
+# Authorization codes expire after 10 minutes (RFC 6749 §4.1.2).
+_AUTH_CODE_TTL = 600
+
+# A pending authorization is a browser mid-consent, so it lives about as long as
+# a person takes to sign in. Without a TTL a `session_id` stayed valid until the
+# process restarted, which both kept phishing links alive indefinitely and let
+# `/authorize` — reachable without any credential — grow the map without bound
+# in a 512 MB container. The cap is the backstop for a burst inside one window.
+_PENDING_AUTH_TTL = 900
+_MAX_PENDING_AUTH = 2000
+
+
+class TransientAuthError(RuntimeError):
+    """The Jingwei auth backend was *transiently* unavailable (network error,
+    timeout, or 5xx) during a token exchange.
+
+    Distinct from ``TokenError(invalid_grant)`` on purpose. The MCP SDK's token
+    handler only catches ``TokenError`` (serialized as HTTP 400), so this
+    propagates as an HTTP 500 instead. That difference matters to the client:
+
+    - ``invalid_grant`` means "your grant is permanently bad" → a well-behaved
+      client discards the refresh token and forces a full re-authorization.
+    - a 5xx means "server problem, try again" → the client retries and KEEPS its
+      refresh token.
+
+    Mapping a transient outage to ``invalid_grant`` would log everyone out during
+    exactly the redeploy window #833 is meant to make seamless, so transient
+    failures must never be reported as a bad grant.
+    """
+
+
+# ── Persistent client store ──────────────────────────────────
+
+
+def _load_clients() -> dict[str, dict[str, Any]]:
+    if _CLIENTS_FILE.exists():
+        try:
+            return json.loads(_CLIENTS_FILE.read_text())
+        except (json.JSONDecodeError, OSError):
+            return {}
+    return {}
+
+
+def _save_clients(clients: dict[str, dict[str, Any]]) -> None:
+    _ensure_dir()
+    _write_private_file(_CLIENTS_FILE, json.dumps(clients, indent=2) + "\n")
+
+
+# ── Extended models with user identity ───────────────────────
+
+
+class DeepCellAuthorizationCode(AuthorizationCode):
+    """Authorization code carrying the authenticated user_id and the Jingwei
+    refresh token issued at login (reused at code-exchange time)."""
+
+    user_id: str
+    jingwei_refresh_token: str
+
+
+class DeepCellAccessToken(AccessToken):
+    """Access token enriched with user identity from JWT claims."""
+
+    user_id: str
+
+
+class DeepCellRefreshToken(RefreshToken):
+    """Refresh token wrapper. ``user_id`` is optional: when delegating to
+    Jingwei the user is resolved server-side from the token, so we don't always
+    know it locally (e.g. in ``load_refresh_token``)."""
+
+    user_id: str = ""
+
+
+# ── Provider ─────────────────────────────────────────────────
+
+
+class DeepCellOAuthProvider(
+    OAuthAuthorizationServerProvider[
+        DeepCellAuthorizationCode,
+        DeepCellRefreshToken,
+        DeepCellAccessToken,
+    ]
+):
+    """MCP OAuth 2.1 provider backed by Jingwei API for user auth."""
+
+    def __init__(
+        self,
+        *,
+        api_url: str,
+        server_url: str,
+        jwt_secret: str,
+        frontend_url: str | None = None,
+        jwt_algorithm: str = "HS256",
+        jwt_issuer: str = "deepcell-api",
+        jwt_audience: str = "deepcell-user",
+    ) -> None:
+        self.api_url = api_url.rstrip("/")
+        self.server_url = server_url.rstrip("/")
+        # Where the consent UI lives. Defaults to this server's own origin
+        # because nginx serves the Next.js frontend at `/` on the same host
+        # that proxies `/mcp` and `/oauth/*` here.
+        self.frontend_url = (frontend_url or server_url).rstrip("/")
+        self.jwt_secret = jwt_secret
+        self.jwt_algorithm = jwt_algorithm
+        self.jwt_issuer = jwt_issuer
+        self.jwt_audience = jwt_audience
+
+        # In-memory stores. Refresh tokens are NOT stored here — they are
+        # delegated to Jingwei (see ``_jingwei_refresh``).
+        self._clients: dict[str, dict[str, Any]] = _load_clients()
+        self._auth_codes: dict[str, DeepCellAuthorizationCode] = {}
+        # session_id -> (client, params, created_at). Pruned by _prune_pending.
+        self._pending_auth: dict[
+            str, tuple[OAuthClientInformationFull, AuthorizationParams, float]
+        ] = {}
+        self._api_key_cache: dict[str, tuple[DeepCellAccessToken, float]] = {}
+
+    # ── Client registration (DCR) ────────────────────────────
+
+    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
+        data = self._clients.get(client_id)
+        if data is None:
+            return None
+        return OAuthClientInformationFull.model_validate(data)
+
+    async def register_client(
+        self, client_info: OAuthClientInformationFull
+    ) -> None:
+        if client_info.client_id is None:
+            client_info.client_id = f"dc_{secrets.token_urlsafe(16)}"
+        if client_info.client_secret is None:
+            client_info.client_secret = secrets.token_urlsafe(32)
+        client_info.client_id_issued_at = int(time.time())
+
+        self._clients[client_info.client_id] = client_info.model_dump(
+            mode="json", exclude_none=True
+        )
+        _save_clients(self._clients)
+
+    # ── Authorization ────────────────────────────────────────
+
+    async def authorize(
+        self,
+        client: OAuthClientInformationFull,
+        params: AuthorizationParams,
+    ) -> str:
+        session_id = secrets.token_urlsafe(20)
+        self._pending_auth[session_id] = (client, params, time.time())
+        # Prune AFTER inserting so the cap holds including this entry, and so
+        # oldest-first eviction can never drop the session we just handed out.
+        self._prune_pending()
+        # Hand the browser to the shared consent page on the frontend, which
+        # renders the same AuthPanel as every other sign-in surface (Google
+        # included — this server can only check passwords, and OAuth-only
+        # accounts have none). The OAuth state machine stays here: the page
+        # comes back to POST /oauth/complete with an authorized device code.
+        return f"{self.frontend_url}/mcp-auth?session_id={quote(session_id, safe='')}"
+
+    def _prune_pending(self) -> None:
+        """Drop expired pending authorizations, then enforce the hard cap.
+
+        Called on every ``authorize()`` because that is the only unauthenticated
+        entry point that grows the map.
+        """
+        cutoff = time.time() - _PENDING_AUTH_TTL
+        for sid in [s for s, (_, _, ts) in self._pending_auth.items() if ts < cutoff]:
+            self._pending_auth.pop(sid, None)
+
+        # Still over cap after pruning: a burst inside one TTL window. Evict
+        # oldest-first (dicts preserve insertion order, and insertion order is
+        # creation order here) so a flood can't pin memory until restart.
+        overflow = len(self._pending_auth) - _MAX_PENDING_AUTH
+        if overflow > 0:
+            logger.warning(
+                "pending-authorization cap reached; evicting %d oldest", overflow
+            )
+            for sid in list(self._pending_auth)[:overflow]:
+                self._pending_auth.pop(sid, None)
+
+    def describe_pending(self, session_id: str) -> dict[str, Any] | None:
+        """Consent-screen metadata for a pending authorization.
+
+        Returns ``None`` when the session is unknown, expired, or already
+        consumed, so the page can show "expired" instead of a form that cannot
+        succeed.
+
+        ``redirect_host`` is the load-bearing field. Dynamic client registration
+        is open, so ``client_name`` and ``client_uri`` are free text chosen by
+        whoever registered the client — anyone can register one calling itself
+        "Claude.ai". The registered redirect target is the only value in here
+        the client cannot lie about, because the SDK sends the code there and
+        nowhere else, so it is what actually tells a user where their grant is
+        going.
+        """
+        pending = self._pending_auth.get(session_id)
+        if pending is None:
+            return None
+
+        client, params, created_at = pending
+        if time.time() - created_at > _PENDING_AUTH_TTL:
+            self._pending_auth.pop(session_id, None)
+            return None
+        redirect_host = ""
+        if params.redirect_uri:
+            parsed = urlparse(str(params.redirect_uri))
+            redirect_host = parsed.netloc or str(params.redirect_uri)
+
+        return {
+            "client_name": client.client_name or client.client_id or "",
+            "client_uri": str(client.client_uri) if client.client_uri else None,
+            "redirect_host": redirect_host,
+            "scopes": list(params.scopes or ["deepcell"]),
+        }
+
+    async def redeem_device_code(self, device_code: str) -> tuple[str, str] | None:
+        """Exchange an *authorized* device code for ``(user_id, refresh_token)``.
+
+        The consent page mints the code (``/auth/device/code``) and approves it
+        as the signed-in user (``/auth/device/authorize``); this redeems it. All
+        token minting therefore stays server-side in Jingwei's existing device
+        flow — no refresh token is ever handed to browser JavaScript, and the
+        user's own web session is left untouched rather than rotated out from
+        under them.
+
+        Returns ``None`` when the code is unknown, expired, or still pending.
+        """
+        try:
+            async with httpx.AsyncClient(timeout=10, headers=client_headers("mcp")) as client:
+                resp = await client.post(
+                    f"{self.api_url}/auth/device/token",
+                    json={"device_code": device_code},
+                )
+        except httpx.HTTPError as exc:
+            logger.error("Jingwei device-token request failed: %s", exc)
+            return None
+
+        if resp.status_code != 200:
+            logger.warning(
+                "Device-code redemption rejected (HTTP %s)", resp.status_code
+            )
+            return None
+
+        data = resp.json()
+        # A code that was never approved comes back 200 with {"status":
+        # "pending"} rather than an error, so this is a real branch, not
+        # defensive padding.
+        if data.get("status") == "pending":
+            return None
+
+        user_id = data.get("user", {}).get("id")
+        refresh_token = data.get("refresh_token")
+        if not user_id or not refresh_token:
+            logger.error("Device-token response missing user id or refresh token")
+            return None
+        return user_id, refresh_token
+
+    # ── Authorization code lifecycle ─────────────────────────
+
+    def create_authorization_code(
+        self,
+        session_id: str,
+        user_id: str,
+        jingwei_refresh_token: str,
+    ) -> tuple[str, str, str | None]:
+        """Called by the login callback after successful user authentication.
+
+        ``jingwei_refresh_token`` is the refresh token Jingwei issued during
+        login; it is carried on the (single-use, short-lived) auth code so that
+        ``exchange_authorization_code`` can rotate it into the tokens handed to
+        the MCP client.
+
+        Returns ``(code, redirect_uri, state)`` so the caller can redirect.
+        """
+        pending = self._pending_auth.pop(session_id, None)
+        if pending is None:
+            raise ValueError("Invalid or expired session")
+
+        client, params, created_at = pending
+        if time.time() - created_at > _PENDING_AUTH_TTL:
+            raise ValueError("Invalid or expired session")
+        code = secrets.token_urlsafe(20)  # ≥160 bits entropy
+
+        self._auth_codes[code] = DeepCellAuthorizationCode(
+            code=code,
+            scopes=params.scopes or [],
+            expires_at=time.time() + _AUTH_CODE_TTL,
+            client_id=client.client_id or "",
+            code_challenge=params.code_challenge,
+            redirect_uri=params.redirect_uri,
+            redirect_uri_provided_explicitly=params.redirect_uri_provided_explicitly,
+            resource=params.resource,
+            user_id=user_id,
+            jingwei_refresh_token=jingwei_refresh_token,
+        )
+
+        return code, str(params.redirect_uri), params.state
+
+    async def load_authorization_code(
+        self,
+        client: OAuthClientInformationFull,
+        authorization_code: str,
+    ) -> DeepCellAuthorizationCode | None:
+        ac = self._auth_codes.get(authorization_code)
+        if ac is None:
+            return None
+        if time.time() > ac.expires_at:
+            self._auth_codes.pop(authorization_code, None)
+            return None
+        if ac.client_id != client.client_id:
+            return None
+        return ac
+
+    async def exchange_authorization_code(
+        self,
+        client: OAuthClientInformationFull,
+        authorization_code: DeepCellAuthorizationCode,
+    ) -> OAuthToken:
+        # Remove the code (single use)
+        self._auth_codes.pop(authorization_code.code, None)
+
+        # Rotate the login-issued Jingwei refresh token into the tokens we hand
+        # the MCP client. This immediately consumes the login token (so the copy
+        # stashed on the auth code can't be replayed) and returns a fresh,
+        # Postgres-backed access/refresh pair.
+        access_token, refresh_token, expires_in = await self._jingwei_refresh(
+            authorization_code.jingwei_refresh_token
+        )
+
+        scopes = authorization_code.scopes
+        return OAuthToken(
+            access_token=access_token,
+            token_type="Bearer",
+            expires_in=expires_in,
+            scope=" ".join(scopes) if scopes else None,
+            refresh_token=refresh_token,
+        )
+
+    # ── Token verification ───────────────────────────────────
+
+    async def load_access_token(self, token: str) -> DeepCellAccessToken | None:
+        # API key path — verify via Jingwei API
+        if token.startswith("dck_"):
+            return await self._verify_api_key(token)
+
+        # JWT path
+        if not self.jwt_secret:
+            return None
+        try:
+            payload = pyjwt.decode(
+                token,
+                self.jwt_secret,
+                algorithms=[self.jwt_algorithm],
+                issuer=self.jwt_issuer,
+                audience=self.jwt_audience,
+            )
+            if payload.get("type") != "access":
+                return None
+            return DeepCellAccessToken(
+                token=token,
+                client_id="",  # Not tracked per-token
+                scopes=["deepcell"],
+                expires_at=payload.get("exp"),
+                user_id=payload["sub"],
+            )
+        except (pyjwt.ExpiredSignatureError, pyjwt.InvalidTokenError):
+            return None
+
+    async def _verify_api_key(self, token: str) -> DeepCellAccessToken | None:
+        """Verify a ``dck_`` API key via the Jingwei ``/auth/api-keys/verify``
+        endpoint.  Results are cached in-memory for 60 seconds."""
+        import hashlib
+        import os
+
+        cache_key = hashlib.sha256(token.encode()).hexdigest()
+        now = time.time()
+
+        # Check cache
+        cached = self._api_key_cache.get(cache_key)
+        if cached and now - cached[1] < 60:
+            return cached[0]
+
+        # Call Jingwei API
+        jingwei_api_key = os.environ.get("JINGWEI_API_KEY", "")
+        headers: dict[str, str] = {}
+        if jingwei_api_key:
+            headers["Authorization"] = f"Bearer {jingwei_api_key}"
+
+        try:
+            async with httpx.AsyncClient(timeout=10, headers=client_headers("mcp")) as client:
+                resp = await client.post(
+                    f"{self.api_url}/auth/api-keys/verify",
+                    json={"key": token},
+                    headers=headers,
+                )
+                if resp.status_code == 200:
+                    data = resp.json()
+                    result = DeepCellAccessToken(
+                        token=token,
+                        client_id="api-key",
+                        scopes=data.get("scopes", ["deepcell"]),
+                        expires_at=None,
+                        user_id=data["user_id"],
+                    )
+                    self._api_key_cache[cache_key] = (result, now)
+                    return result
+        except httpx.HTTPError as exc:
+            logger.error("API key verification request failed: %s", exc)
+
+        return None
+
+    # ── Refresh token lifecycle ──────────────────────────────
+
+    async def load_refresh_token(
+        self,
+        client: OAuthClientInformationFull,
+        refresh_token: str,
+    ) -> DeepCellRefreshToken | None:
+        # There is no local refresh-token store — the token is a Jingwei refresh
+        # token, validated server-side at rotation time (see _jingwei_refresh).
+        # Return an optimistic wrapper bound to the requesting client so the SDK
+        # token handler's pre-exchange checks (client_id match, expiry, scope
+        # subset) pass; Jingwei is the real authority on validity.
+        return DeepCellRefreshToken(
+            token=refresh_token,
+            client_id=client.client_id or "",
+            scopes=["deepcell"],
+            expires_at=None,
+        )
+
+    async def exchange_refresh_token(
+        self,
+        client: OAuthClientInformationFull,
+        refresh_token: DeepCellRefreshToken,
+        scopes: list[str],
+    ) -> OAuthToken:
+        # Delegate rotation to Jingwei (Postgres-backed, with token-family
+        # reuse-detection). Raises TokenError(invalid_grant) if the token is
+        # rejected, or TransientAuthError if the backend is unreachable.
+        access_token, new_refresh, expires_in = await self._jingwei_refresh(
+            refresh_token.token
+        )
+
+        return OAuthToken(
+            access_token=access_token,
+            token_type="Bearer",
+            expires_in=expires_in,
+            scope=" ".join(scopes) if scopes else None,
+            refresh_token=new_refresh,
+        )
+
+    # ── Revocation ───────────────────────────────────────────
+
+    async def revoke_token(
+        self,
+        token: DeepCellAccessToken | DeepCellRefreshToken,
+    ) -> None:
+        # A refresh token *is* a Jingwei refresh token, so an explicit OAuth
+        # revocation (e.g. a connector being disconnected) is delegated to
+        # Jingwei's ``/auth/revoke-token`` endpoint, which invalidates the
+        # underlying session in Postgres. Best-effort: per RFC 7009 the client's
+        # revoke request should succeed regardless, so transport/HTTP failures
+        # are logged and swallowed rather than surfaced. Access tokens are
+        # stateless short-lived JWTs that cannot be revoked without a blocklist,
+        # so they simply expire (15 min).
+        if isinstance(token, DeepCellRefreshToken):
+            try:
+                async with httpx.AsyncClient(timeout=10, headers=client_headers("mcp")) as client:
+                    await client.post(
+                        f"{self.api_url}/auth/revoke-token",
+                        json={"refresh_token": token.token},
+                    )
+            except httpx.HTTPError as exc:
+                logger.warning("Jingwei revoke request failed (best-effort): %s", exc)
+        return None
+
+    # ── Internal helpers ─────────────────────────────────────
+
+    async def _jingwei_refresh(self, raw_refresh_token: str) -> tuple[str, str, int]:
+        """Rotate a Jingwei refresh token via the API's ``/auth/refresh`` endpoint.
+
+        An MCP refresh token *is* a Jingwei refresh token, so rotation,
+        token-family reuse-detection and revocation all live in Postgres and
+        survive container recreates. ``/auth/refresh`` accepts the token in the
+        request body (``RefreshRequest.refresh_token``).
+
+        Returns ``(access_token, refresh_token, expires_in)``.
+
+        Error mapping distinguishes a *bad grant* from a *transient outage* so a
+        backend blip doesn't log clients out (see ``TransientAuthError``):
+
+        Raises:
+            TokenError(``invalid_grant``): Jingwei rejected the token with a 4xx
+                (invalid / expired / reused) — the grant is genuinely bad, so the
+                client should re-authorize.
+            TransientAuthError: the backend was unreachable, timed out, returned
+                a 5xx, or returned 200 without tokens (a contract violation). The
+                grant may still be valid, so the client should retry and keep it.
+        """
+        try:
+            async with httpx.AsyncClient(timeout=10, headers=client_headers("mcp")) as client:
+                resp = await client.post(
+                    f"{self.api_url}/auth/refresh",
+                    json={"refresh_token": raw_refresh_token},
+                )
+        except httpx.HTTPError as exc:
+            logger.error("Jingwei refresh request failed (transient): %s", exc)
+            raise TransientAuthError("Auth service unreachable") from exc
+
+        if resp.status_code == 200:
+            data = resp.json()
+            access_token = data.get("access_token")
+            new_refresh = data.get("refresh_token")
+            if not access_token or not new_refresh:
+                # 200 but no tokens is a backend bug, not a rejected grant —
+                # don't force the client to re-auth over it.
+                logger.error("Jingwei refresh response missing tokens")
+                raise TransientAuthError("Auth service returned no tokens")
+            return access_token, new_refresh, int(data.get("expires_in", 900))
+
+        if 400 <= resp.status_code < 500:
+            logger.info("Jingwei rejected refresh token (status %s)", resp.status_code)
+            raise TokenError(
+                error="invalid_grant",
+                error_description="Refresh token is invalid or expired",
+            )
+
+        # 5xx (and anything else non-2xx): transient backend failure.
+        logger.error("Jingwei refresh failed transiently (status %s)", resp.status_code)
+        raise TransientAuthError(f"Auth service error (status {resp.status_code})")
+