setup workos

2026-05-02 20:57:09 -04:00
parent 45bdd7b629
commit a956c4a973
20 changed files with 1286 additions and 142 deletions
@@ -0,0 +1,202 @@
+"""WorkOS AuthKit helpers for the FastAPI web app."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from os import getenv
+from typing import Any
+
+from fastapi import Request
+from workos import WorkOSClient
+from workos.session import seal_session_from_auth_response
+
+
+@dataclass(frozen=True)
+class AuthConfig:
+    """Runtime configuration for WorkOS AuthKit."""
+
+    api_key: str
+    client_id: str
+    cookie_password: str
+    redirect_uri: str
+    logout_redirect_uri: str
+    session_cookie_name: str
+    organization_id: str
+
+    @property
+    def secure_cookies(self) -> bool:
+        return self.redirect_uri.startswith("https://")
+
+
+@dataclass(frozen=True)
+class AuthSession:
+    """Normalized auth session passed through the app."""
+
+    user_id: str
+    email: str
+    first_name: str | None
+    last_name: str | None
+    role_slugs: set[str]
+    organization_id: str | None
+    raw_user: Any
+    raw_session: Any
+
+    @property
+    def display_name(self) -> str:
+        parts = [part for part in (self.first_name, self.last_name) if part]
+        return " ".join(parts) if parts else self.email
+
+    @property
+    def is_admin(self) -> bool:
+        return "admin" in self.role_slugs
+
+
+@dataclass(frozen=True)
+class CallbackResult:
+    """Result of exchanging a WorkOS callback code."""
+
+    sealed_session: str
+    next_path: str
+
+
+def safe_next_path(value: str | None, default: str = "/dashboard") -> str:
+    """Allow only local relative redirect targets."""
+    if value and value.startswith("/") and not value.startswith("//"):
+        return value
+    return default
+
+
+def build_authorization_url(next_path: str) -> str:
+    """Build the WorkOS hosted login URL."""
+    config = get_auth_config()
+    return get_workos_client().user_management.get_authorization_url(
+        provider="authkit",
+        redirect_uri=config.redirect_uri,
+        state=safe_next_path(next_path),
+        organization_id=config.organization_id,
+    )
+
+
+def exchange_code(request: Request) -> CallbackResult:
+    """Exchange a WorkOS callback code for a sealed session cookie value."""
+    code = request.query_params.get("code")
+    if not code:
+        raise ValueError("Missing authentication code.")
+
+    config = get_auth_config()
+    auth_response = get_workos_client().user_management.authenticate_with_code(
+        code=code,
+        ip_address=_request_ip(request),
+        user_agent=request.headers.get("user-agent"),
+    )
+    sealed_session = seal_session_from_auth_response(
+        access_token=auth_response.access_token,
+        refresh_token=auth_response.refresh_token,
+        user=auth_response.user.to_dict(),
+        impersonator=auth_response.impersonator.to_dict()
+        if auth_response.impersonator is not None
+        else None,
+        cookie_password=config.cookie_password,
+    )
+
+    return CallbackResult(
+        sealed_session=sealed_session,
+        next_path=safe_next_path(request.query_params.get("state")),
+    )
+
+
+def get_current_session(request: Request) -> AuthSession | None:
+    """Load the current signed-in WorkOS session from the sealed cookie."""
+    cookie_name = getenv("WORKOS_SESSION_COOKIE_NAME", "workos_session")
+    sealed_session = request.cookies.get(cookie_name)
+    if not sealed_session:
+        return None
+
+    config = get_auth_config()
+    session = get_workos_client().user_management.load_sealed_session(
+        session_data=sealed_session,
+        cookie_password=config.cookie_password,
+    )
+    auth_response = session.authenticate()
+    if not getattr(auth_response, "authenticated", False):
+        return None
+
+    user = auth_response.user or {}
+    organization_id = getattr(auth_response, "organization_id", None)
+    if config.organization_id and organization_id != config.organization_id:
+        return None
+    role_slugs = set(getattr(auth_response, "roles", None) or [])
+    role = getattr(auth_response, "role", None)
+    if role:
+        role_slugs.add(role)
+
+    return AuthSession(
+        user_id=_user_field(user, "id") or "",
+        email=_user_field(user, "email") or "",
+        first_name=_user_field(user, "first_name"),
+        last_name=_user_field(user, "last_name"),
+        role_slugs=role_slugs,
+        organization_id=organization_id,
+        raw_user=user,
+        raw_session=auth_response,
+    )
+
+
+def get_logout_url(request: Request) -> str:
+    """Return the WorkOS logout URL for the current sealed session."""
+    config = get_auth_config()
+    sealed_session = request.cookies.get(config.session_cookie_name)
+    if not sealed_session:
+        return config.logout_redirect_uri
+
+    session = get_workos_client().user_management.load_sealed_session(
+        session_data=sealed_session,
+        cookie_password=config.cookie_password,
+    )
+    return session.get_logout_url(return_to=config.logout_redirect_uri)
+
+
+@lru_cache(maxsize=1)
+def get_auth_config() -> AuthConfig:
+    """Load and validate WorkOS environment configuration."""
+    values = {
+        "WORKOS_API_KEY": getenv("WORKOS_API_KEY"),
+        "WORKOS_CLIENT_ID": getenv("WORKOS_CLIENT_ID"),
+        "WORKOS_COOKIE_PASSWORD": getenv("WORKOS_COOKIE_PASSWORD"),
+        "WORKOS_ORGANIZATION_ID": getenv("WORKOS_ORGANIZATION_ID"),
+    }
+    missing = [name for name, value in values.items() if not value]
+    if missing:
+        raise RuntimeError(
+            "Missing WorkOS configuration: " + ", ".join(sorted(missing))
+        )
+
+    return AuthConfig(
+        api_key=values["WORKOS_API_KEY"] or "",
+        client_id=values["WORKOS_CLIENT_ID"] or "",
+        cookie_password=values["WORKOS_COOKIE_PASSWORD"] or "",
+        redirect_uri=getenv("WORKOS_REDIRECT_URI", "http://localhost:8000/callback"),
+        logout_redirect_uri=getenv("WORKOS_LOGOUT_REDIRECT_URI", "http://localhost:8000/"),
+        session_cookie_name=getenv("WORKOS_SESSION_COOKIE_NAME", "workos_session"),
+        organization_id=values["WORKOS_ORGANIZATION_ID"] or "",
+    )
+
+
+@lru_cache(maxsize=1)
+def get_workos_client():
+    """Create and cache the WorkOS SDK client."""
+    config = get_auth_config()
+    return WorkOSClient(api_key=config.api_key, client_id=config.client_id)
+
+
+def _request_ip(request: Request) -> str | None:
+    if request.client is None:
+        return None
+    return request.client.host
+
+
+def _user_field(user: Any, key: str) -> Any:
+    if isinstance(user, dict):
+        return user.get(key)
+    return getattr(user, key, None)