from typing import Annotated from fastapi import Depends, Header, HTTPException, Query, WebSocket from loguru import logger from api.constants import AUTH_PROVIDER from api.db import db_client from api.db.models import UserModel from api.enums import PostHogEvent from api.services.auth.stack_auth import stackauth from api.services.organization_bootstrap import ensure_organization_bootstrapped from api.services.posthog_client import ( POSTHOG_ORGANIZATION_GROUP_TYPE, capture_event, group_identify, set_person_properties, ) from api.utils.auth import decode_jwt_token async def require_local_auth() -> None: """Reject email/password auth requests outside OSS (local) deployments. The auth router stays mounted in every mode so the OpenAPI spec — and the clients generated from it — don't vary with AUTH_PROVIDER; the gate has to happen at request time. Without it, the SaaS deployment accepts unauthenticated signups that mint oss_* users bypassing Stack Auth. """ if AUTH_PROVIDER != "local": raise HTTPException(status_code=404, detail="Not found") async def get_user( authorization: Annotated[str | None, Header()] = None, x_api_key: Annotated[str | None, Header(alias="X-API-Key")] = None, ) -> UserModel: # ------------------------------------------------------------------ # Check if API key is provided (takes precedence) # ------------------------------------------------------------------ if x_api_key: return await _handle_api_key_auth(x_api_key) # ------------------------------------------------------------------ # Check if we're using local (email/password) auth # ------------------------------------------------------------------ if AUTH_PROVIDER == "local": return await _handle_oss_auth(authorization) # ------------------------------------------------------------------ # 1. Validate and fetch the authenticated Stack user # ------------------------------------------------------------------ stack_user = await stackauth.get_user(authorization) if stack_user is None: raise HTTPException(status_code=401, detail="Unauthorized") # ------------------------------------------------------------------ # 2. Ensure the user has a team (Stack "selected_team_id") # ------------------------------------------------------------------ selected_team_id: str | None = stack_user.get("selected_team_id") if not selected_team_id and stack_user.get("selected_team"): selected_team_id = stack_user["selected_team"].get("id") if not selected_team_id: raise HTTPException(status_code=400, detail="No team selected") # ------------------------------------------------------------------ # 3. Persist/Fetch the local User model # ------------------------------------------------------------------ try: ( user_model, user_was_created, ) = await db_client.get_or_create_user_by_provider_id(stack_user["id"]) # Sync email from Stack Auth if available and not already set stack_email = stack_user.get("primary_email_verified") and stack_user.get( "primary_email" ) if stack_email and user_model.email != stack_email: await db_client.update_user_email(user_model.id, stack_email) user_model.email = stack_email if user_was_created: capture_event( distinct_id=str(stack_user["id"]), event=PostHogEvent.SIGNED_UP, properties={ "auth_provider": "stack", }, ) except Exception as e: raise HTTPException( status_code=500, detail=f"Error while creating user from database {e}" ) # ------------------------------------------------------------------ # 4. Persist Organization (team) and mapping in local database # ------------------------------------------------------------------ try: ( organization, org_was_created, ) = await db_client.get_or_create_organization_by_provider_id( org_provider_id=selected_team_id, user_id=user_model.id ) if org_was_created: _sync_created_organization_to_posthog( organization=organization, stack_user=stack_user, ) # Check if user's selected organization differs from the current organization if user_model.selected_organization_id != organization.id: await db_client.add_user_to_organization(user_model.id, organization.id) # Update user's selected organization await db_client.update_user_selected_organization( user_model.id, organization.id ) # Update the user_model object to reflect the change user_model.selected_organization_id = organization.id _associate_user_with_posthog_organization( user=user_model, organization=organization, stack_user=stack_user, org_was_created=org_was_created, ) except Exception as exc: raise HTTPException( status_code=500, detail=f"Failed to map user to organization: {exc}", ) # Deliberately outside the org-mapping branch above: provisioning is keyed # on whether the organization is actually configured, not on whether this # request happened to create it or switch the user into it. That is what # lets a failed attempt be retried instead of stranding the organization # with no model configuration forever. It is safe to call every request — # it self-limits to one indexed read once the organization is provisioned. await ensure_organization_bootstrapped( organization.id, created_by=str(stack_user["id"]), ) return user_model def _sync_created_organization_to_posthog( *, organization, stack_user: dict | None = None, created_by_provider_id: str | None = None, ) -> None: """Create/update the PostHog organization group for a newly-created org.""" try: organization_id = int(organization.id) organization_provider_id = getattr(organization, "provider_id", None) created_by = created_by_provider_id if created_by is None and stack_user and stack_user.get("id"): created_by = str(stack_user["id"]) properties = { "organization_id": organization_id, "organization_provider_id": organization_provider_id, "auth_provider": "stack", } if created_by: properties["created_by_provider_id"] = created_by group_identify( POSTHOG_ORGANIZATION_GROUP_TYPE, str(organization_id), properties, distinct_id=created_by, ) if created_by: capture_event( distinct_id=created_by, event=PostHogEvent.ORGANIZATION_CREATED, properties=properties, groups={POSTHOG_ORGANIZATION_GROUP_TYPE: str(organization_id)}, ) except Exception: logger.exception("Failed to sync created organization to PostHog") def _associate_user_with_posthog_organization( *, user: UserModel, organization, stack_user: dict | None = None, user_distinct_id: str | None = None, org_was_created: bool, organization_ids: list[int] | None = None, selected_organization_id: int | None = None, selected_organization_provider_id: str | None = None, ) -> None: """Attach the Stack user to the PostHog organization group.""" try: organization_id = int(organization.id) organization_provider_id = getattr(organization, "provider_id", None) if user_distinct_id is None: if stack_user and stack_user.get("id"): user_distinct_id = str(stack_user["id"]) else: user_distinct_id = str(user.provider_id) selected_org_id = selected_organization_id or organization_id selected_org_provider_id = ( selected_organization_provider_id or organization_provider_id ) person_properties = { "user_id": user.id, "user_provider_id": user_distinct_id, "selected_organization_id": selected_org_id, "selected_organization_provider_id": selected_org_provider_id, } if organization_ids is not None: person_properties["organization_ids"] = organization_ids if user.email: person_properties["email"] = user.email set_person_properties(user_distinct_id, person_properties) event_properties = { "user_id": user.id, "organization_id": organization_id, "organization_provider_id": organization_provider_id, "auth_provider": "stack", "organization_was_created": org_was_created, } capture_event( distinct_id=user_distinct_id, event=PostHogEvent.ORGANIZATION_USER_ASSOCIATED, properties=event_properties, groups={POSTHOG_ORGANIZATION_GROUP_TYPE: str(organization_id)}, ) except Exception: logger.exception("Failed to associate user with PostHog organization") async def get_user_with_selected_organization( user: Annotated[UserModel, Depends(get_user)], ) -> UserModel: if not user.selected_organization_id: raise HTTPException(status_code=400, detail="No organization selected") return user async def _handle_oss_auth(authorization: str | None) -> UserModel: """ Handle authentication for OSS deployment mode. Validates JWT tokens issued by the email/password auth flow. """ if not authorization: raise HTTPException(status_code=401, detail="Authorization header required") # Remove "Bearer " prefix if present token = ( authorization.replace("Bearer ", "") if authorization.startswith("Bearer ") else authorization ) if not token: raise HTTPException(status_code=401, detail="Invalid authorization token") try: payload = decode_jwt_token(token) user = await db_client.get_user_by_id(int(payload["sub"])) if user is None: raise HTTPException(status_code=401, detail="User not found") except HTTPException: raise except Exception: raise HTTPException(status_code=401, detail="Invalid or expired token") # Deliberately outside the try above: a provisioning failure must not be # reported to the user as an expired token. if user.selected_organization_id: await ensure_organization_bootstrapped( user.selected_organization_id, created_by=user.provider_id, ) return user async def _handle_api_key_auth(api_key: str) -> UserModel: """ Handle authentication via X-API-Key header. Returns the user who created the API key with the correct organization context. """ # Validate the API key api_key_model = await db_client.validate_api_key(api_key) if not api_key_model: raise HTTPException(status_code=401, detail="Invalid or expired API key") # API key must have a created_by user if not api_key_model.created_by: raise HTTPException(status_code=401, detail="API key has no associated user") # Get the user who created this API key user = await db_client.get_user_by_id(api_key_model.created_by) if not user: raise HTTPException(status_code=401, detail="API key owner not found") # Set the organization context to the API key's organization user.selected_organization_id = api_key_model.organization_id logger.debug( f"Authenticated via API key: {api_key_model.key_prefix}... " f"(user_id={user.id}, org_id={api_key_model.organization_id})" ) return user async def get_superuser( authorization: Annotated[str | None, Header()] = None, x_api_key: Annotated[str | None, Header(alias="X-API-Key")] = None, ) -> UserModel: """ Dependency to check if the authenticated user is a superuser. Raises HTTPException if user is not authenticated or not a superuser. """ user = await get_user(authorization, x_api_key) if not user.is_superuser: raise HTTPException( status_code=403, detail="Access denied. Superuser privileges required." ) return user async def get_user_ws( websocket: WebSocket, token: str = Query(None), api_key: str = Query(None, alias="api_key"), ) -> UserModel: """ WebSocket authentication dependency. Uses token or api_key from query parameters for authentication. """ if not token and not api_key: await websocket.close(code=1008, reason="Missing authentication token") raise HTTPException(status_code=401, detail="Missing authentication token") try: # API key takes precedence if api_key: user = await get_user(None, api_key) else: # Use the same logic as get_user but with token from query authorization = f"Bearer {token}" user = await get_user(authorization, None) return user except HTTPException as e: await websocket.close(code=1008, reason=e.detail) raise