sweatstack 0.60.0__py3-none-any.whl → 0.62.0__py3-none-any.whl

sweatstack/fastapi/models.py

@@ -0,0 +1,175 @@
+"""Data models for FastAPI session management."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Protocol, runtime_checkable
+
+from pydantic import SecretStr
+
+from ..utils import decode_jwt_body
+
+
+@dataclass(frozen=True, slots=True)
+class TokenSet:
+    """Immutable token pair with user ID.
+
+    This represents either principal or delegated tokens stored in the session.
+    The frozen=True ensures tokens can't be accidentally modified.
+    """
+
+    access_token: str
+    refresh_token: str
+    user_id: str
+
+    def to_dict(self) -> dict[str, str]:
+        """Serialize to dictionary for session storage."""
+        return {
+            "access_token": self.access_token,
+            "refresh_token": self.refresh_token,
+            "user_id": self.user_id,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> TokenSet:
+        """Deserialize from dictionary."""
+        return cls(
+            access_token=data["access_token"],
+            refresh_token=data["refresh_token"],
+            user_id=data["user_id"],
+        )
+
+
+@dataclass(slots=True)
+class SessionData:
+    """Type-safe wrapper for session data.
+
+    Handles both the new format (with principal/delegated) and legacy format
+    (flat access_token/refresh_token/user_id) for backwards compatibility.
+    """
+
+    principal: TokenSet
+    delegated: TokenSet | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary for cookie storage."""
+        data: dict[str, Any] = {"principal": self.principal.to_dict()}
+        if self.delegated:
+            data["delegated"] = self.delegated.to_dict()
+        return data
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SessionData:
+        """Deserialize from dictionary.
+
+        Handles both new format and legacy format for backwards compatibility.
+        """
+        # New format: has "principal" key
+        if "principal" in data:
+            return cls(
+                principal=TokenSet.from_dict(data["principal"]),
+                delegated=TokenSet.from_dict(data["delegated"]) if data.get("delegated") else None,
+            )
+
+        # Legacy format: flat structure with access_token, refresh_token, user_id
+        # Migrate to new format by treating as principal
+        return cls(
+            principal=TokenSet(
+                access_token=data["access_token"],
+                refresh_token=data["refresh_token"],
+                user_id=data["user_id"],
+            ),
+            delegated=None,
+        )
+
+
+def extract_user_id(jwt_token: str | SecretStr) -> str:
+    """Extract user ID ('sub' claim) from a JWT token.
+
+    This does not validate the signature - the token was already validated
+    by the API when it was issued.
+
+    Args:
+        jwt_token: The JWT access token (str or SecretStr).
+
+    Returns:
+        The user ID from the token's 'sub' claim.
+
+    Raises:
+        ValueError: If the token is malformed or missing the 'sub' claim.
+    """
+    try:
+        token_str = jwt_token.get_secret_value() if isinstance(jwt_token, SecretStr) else jwt_token
+        payload = decode_jwt_body(token_str)
+        user_id = payload.get("sub")
+        if not user_id:
+            raise ValueError("Token missing 'sub' claim")
+        return user_id
+    except (IndexError, KeyError) as e:
+        raise ValueError(f"Malformed JWT token: {e}") from e
+
+
+# ---------------------------------------------------------------------------
+# Token storage for webhook support
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class StoredTokens:
+    """Token data for persistent storage.
+
+    Used by TokenStore implementations to persist tokens for webhook handling.
+    The library does NOT encrypt tokens - see Security section in documentation.
+
+    Attributes:
+        user_id: The SweatStack user ID.
+        access_token: The OAuth access token.
+        refresh_token: The OAuth refresh token.
+        expires_at: When the access token expires.
+    """
+
+    user_id: str
+    access_token: str
+    refresh_token: str
+    expires_at: datetime
+
+    def __repr__(self) -> str:
+        """Hide sensitive data in logs."""
+        return (
+            f"StoredTokens(user_id={self.user_id!r}, "
+            f"access_token='***', refresh_token='***', "
+            f"expires_at={self.expires_at!r})"
+        )
+
+
+@runtime_checkable
+class TokenStore(Protocol):
+    """Protocol for persisting OAuth tokens.
+
+    Implement this interface to enable AuthenticatedUser in webhook handlers.
+    The library calls these methods automatically:
+    - save(): After OAuth callback and token refresh
+    - load(): When handling webhooks
+    - delete(): On logout
+
+    Thread Safety:
+        All methods may be called concurrently. Your implementation must be thread-safe.
+
+    Error Handling:
+        - save(): Use upsert semantics (don't raise on duplicate user_id)
+        - load(): Return None if user not found (don't raise)
+        - delete(): Be idempotent (don't raise if user doesn't exist)
+    """
+
+    def save(self, tokens: StoredTokens) -> None:
+        """Save or update tokens for a user."""
+        ...
+
+    def load(self, user_id: str) -> StoredTokens | None:
+        """Load tokens for a user. Returns None if not found."""
+        ...
+
+    def delete(self, user_id: str) -> None:
+        """Delete tokens for a user. Idempotent."""
+        ...

sweatstack/fastapi/routes.py

@@ -1,25 +1,34 @@
 """OAuth routes for the FastAPI plugin."""
 
+from __future__ import annotations
+
 import base64
 import json
+import logging
 import secrets
-from urllib.parse import urlencode
+from urllib.parse import urlencode, urlparse
 
 import httpx
-from fastapi import APIRouter, FastAPI, Request, Response
+from fastapi import APIRouter, FastAPI, HTTPException, Request, Response
 from fastapi.responses import RedirectResponse
 
 from ..constants import DEFAULT_URL
 from ..utils import decode_jwt_body
 from .config import get_config
+from .dependencies import _extract_expiry
+from .models import SessionData, StoredTokens, TokenSet
 from .session import (
+    SESSION_COOKIE_NAME,
     STATE_COOKIE_NAME,
     clear_session_cookie,
     clear_state_cookie,
+    decrypt_session,
     set_session_cookie,
     set_state_cookie,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def validate_redirect(url: str | None) -> str | None:
     """Validate that a redirect URL is a safe relative path.
@@ -31,6 +40,82 @@ def validate_redirect(url: str | None) -> str | None:
     return None
 
 
+def _is_same_origin(referer: str | None, app_url: str) -> bool:
+    """Check if a referer URL is from the same origin as the app."""
+    if not referer:
+        return False
+    try:
+        ref_parsed = urlparse(referer)
+        app_parsed = urlparse(app_url)
+        return (
+            ref_parsed.scheme == app_parsed.scheme
+            and ref_parsed.netloc == app_parsed.netloc
+        )
+    except Exception:
+        return False
+
+
+def _get_redirect_url(request: Request, next_param: str | None) -> str:
+    """Determine the redirect URL after a user selection change.
+
+    Priority: ?next= parameter > Referer header (if same-origin) > /
+    """
+    # First try the explicit next parameter
+    if validated := validate_redirect(next_param):
+        return validated
+
+    # Then try the Referer header if same-origin
+    config = get_config()
+    referer = request.headers.get("referer")
+    if _is_same_origin(referer, config.app_url):
+        # Extract just the path from referer
+        parsed = urlparse(referer)
+        path = parsed.path
+        if parsed.query:
+            path += f"?{parsed.query}"
+        if validated := validate_redirect(path):
+            return validated
+
+    # Default to root
+    return "/"
+
+
+def _get_session_data(request: Request) -> SessionData | None:
+    """Get session data from request cookie."""
+    raw_session = decrypt_session(request.cookies.get(SESSION_COOKIE_NAME))
+    if not raw_session:
+        return None
+    try:
+        return SessionData.from_dict(raw_session)
+    except (KeyError, TypeError):
+        return None
+
+
+def _fetch_delegated_token(principal_tokens: TokenSet, target_user_id: str) -> TokenSet:
+    """Fetch a delegated token for the target user using principal credentials."""
+    config = get_config()
+
+    response = httpx.post(
+        f"{DEFAULT_URL}/api/v1/oauth/delegated-token",
+        headers={"Authorization": f"Bearer {principal_tokens.access_token}"},
+        json={"sub": target_user_id},
+    )
+
+    if response.status_code == 403:
+        raise HTTPException(status_code=403, detail="You don't have access to this user")
+    if response.status_code == 404:
+        raise HTTPException(status_code=404, detail="User not found")
+
+    response.raise_for_status()
+    tokens = response.json()
+
+    return TokenSet(
+        access_token=tokens["access_token"],
+        refresh_token=tokens["refresh_token"],
+        user_id=target_user_id,
+    )
+
+
 def create_state(next_url: str | None) -> str:
     """Create an OAuth state value with nonce and optional redirect."""
     nonce = secrets.token_urlsafe(32)
@@ -82,28 +167,34 @@ def create_router() -> APIRouter:
         code: str | None = None,
         state: str | None = None,
         error: str | None = None,
+        error_description: str | None = None,
     ) -> Response:
         """Handle OAuth callback from SweatStack."""
         config = get_config()
 
+        def error_redirect(error_code: str) -> Response:
+            """Redirect to / with error code in query params."""
+            response = RedirectResponse(url=f"/?auth_error={error_code}", status_code=302)
+            clear_state_cookie(response)
+            return response
+
         # Get state cookie
         state_cookie = request.cookies.get(STATE_COOKIE_NAME)
 
-        # Clear state cookie regardless of outcome
-        response = RedirectResponse(url="/", status_code=302)
-        clear_state_cookie(response)
-
-        # Handle OAuth errors
+        # Handle OAuth errors from provider
         if error:
-            return response
+            logger.warning("OAuth error from provider: %s - %s", error, error_description)
+            return error_redirect(error)
 
         # Verify state (CSRF protection)
         if not state or not state_cookie or state != state_cookie:
-            return Response(content="Invalid state", status_code=400)
+            logger.warning("OAuth state mismatch (possible CSRF)")
+            return error_redirect("invalid_state")
 
         # Exchange code for tokens
         if not code:
-            return Response(content="Missing authorization code", status_code=400)
+            logger.warning("OAuth callback missing authorization code")
+            return error_redirect("missing_code")
 
         try:
             token_response = httpx.post(
@@ -118,24 +209,31 @@ def create_router() -> APIRouter:
             )
             token_response.raise_for_status()
             tokens = token_response.json()
-        except Exception:
-            return response  # Redirect to / on token exchange failure
+        except httpx.HTTPStatusError as e:
+            logger.error("Token exchange failed: %s - %s", e.response.status_code, e.response.text)
+            return error_redirect("token_exchange_failed")
+        except Exception as e:
+            logger.error("Token exchange error: %s", e)
+            return error_redirect("token_exchange_failed")
 
         access_token = tokens.get("access_token")
         refresh_token = tokens.get("refresh_token")
 
         if not access_token:
-            return response
+            logger.error("Token response missing access_token")
+            return error_redirect("invalid_token_response")
 
         # Extract user_id from JWT
         try:
             token_body = decode_jwt_body(access_token)
             user_id = token_body.get("sub")
-        except Exception:
-            return response
+        except Exception as e:
+            logger.error("Failed to decode access token: %s", e)
+            return error_redirect("invalid_token")
 
         if not user_id:
-            return response
+            logger.error("Access token missing 'sub' claim")
+            return error_redirect("invalid_token")
 
         # Create session
         session_data = {
@@ -144,6 +242,17 @@ def create_router() -> APIRouter:
             "user_id": user_id,
         }
 
+        # Persist tokens to store if configured
+        if config.token_store:
+            config.token_store.save(
+                StoredTokens(
+                    user_id=user_id,
+                    access_token=access_token,
+                    refresh_token=refresh_token,
+                    expires_at=_extract_expiry(access_token),
+                )
+            )
+
         # Determine redirect URL from state
         state_data = parse_state(state)
         redirect_url = state_data.get("next", "/")
@@ -154,15 +263,107 @@ def create_router() -> APIRouter:
         return response
 
     @router.post("/logout")
-    def logout() -> Response:
+    def logout(request: Request) -> Response:
         """Clear session and redirect to /."""
+        config = get_config()
+
+        # Delete tokens from store if configured
+        if config.token_store:
+            session = _get_session_data(request)
+            if session:
+                config.token_store.delete(session.principal.user_id)
+
         response = RedirectResponse(url="/", status_code=302)
         clear_session_cookie(response)
         return response
 
+    @router.post("/select-user/{user_id}")
+    def select_user(request: Request, user_id: str, next: str | None = None) -> Response:
+        """Switch to viewing as another user.
+
+        Fetches a delegated token for the target user and stores it in the session.
+        Redirects to Referer (if same-origin), ?next= parameter, or /.
+        """
+        session = _get_session_data(request)
+        if not session:
+            raise HTTPException(status_code=401, detail="Not authenticated")
+
+        # Fetch delegated token for the target user
+        try:
+            delegated_tokens = _fetch_delegated_token(session.principal, user_id)
+        except httpx.HTTPStatusError as e:
+            logger.warning("Failed to fetch delegated token for user %s: %s", user_id, e)
+            raise HTTPException(status_code=403, detail="You don't have access to this user")
+
+        # Update session with delegated tokens
+        updated_session = SessionData(
+            principal=session.principal,
+            delegated=delegated_tokens,
+        )
+
+        redirect_url = _get_redirect_url(request, next)
+        response = RedirectResponse(url=redirect_url, status_code=303)
+        set_session_cookie(response, updated_session.to_dict())
+        return response
+
+    @router.post("/select-self")
+    def select_self(request: Request, next: str | None = None) -> Response:
+        """Switch back to viewing as yourself (clear delegation).
+
+        Removes the delegated tokens from the session.
+        Redirects to Referer (if same-origin), ?next= parameter, or /.
+        """
+        session = _get_session_data(request)
+        if not session:
+            raise HTTPException(status_code=401, detail="Not authenticated")
+
+        # Clear delegation
+        updated_session = SessionData(
+            principal=session.principal,
+            delegated=None,
+        )
+
+        redirect_url = _get_redirect_url(request, next)
+        response = RedirectResponse(url=redirect_url, status_code=303)
+        set_session_cookie(response, updated_session.to_dict())
+        return response
+
     return router
 
 
+def _warn_if_webhook_misconfigured(app: FastAPI) -> None:
+    """Log error if WebhookPayload is used but webhook_secret not configured."""
+    config = get_config()
+
+    if config.webhook_secret:
+        return  # Properly configured
+
+    # Import here to avoid circular imports
+    from .webhooks import _require_webhook_payload
+
+    # Check if any route uses WebhookPayload dependency
+    for route in app.routes:
+        if not hasattr(route, "dependant"):
+            continue
+
+        if _uses_dependency(route.dependant, _require_webhook_payload):
+            raise RuntimeError(
+                f"Route '{route.path}' uses WebhookPayload but webhook_secret is not configured. "
+                "Webhook signature verification will fail at runtime. "
+                "Configure with the SWEATSTACK_WEBHOOK_SECRET env variable or configure(webhook_secret='whsec_...')"
+            )
+
+
+def _uses_dependency(dependant, target_callable) -> bool:
+    """Check if a dependency tree includes the target callable."""
+    for dep in dependant.dependencies:
+        if dep.call is target_callable:
+            return True
+        if hasattr(dep, "dependant") and _uses_dependency(dep.dependant, target_callable):
+            return True
+    return False
+
+
 def instrument(app: FastAPI) -> None:
     """Add SweatStack auth routes to a FastAPI application.
 
@@ -175,3 +376,8 @@ def instrument(app: FastAPI) -> None:
     config = get_config()  # This will raise if not configured
     router = create_router()
     app.include_router(router, prefix=config.auth_route_prefix)
+
+    # Validate webhook configuration at startup (after all routes are registered)
+    @app.on_event("startup")
+    def _check_webhook_config():
+        _warn_if_webhook_misconfigured(app)