""" Unified authentication provider - orchestrates WorkOS SSO and ORY Kratos. This module provides a single interface for authentication that transparently routes between WorkOS (enterprise SSO) and ORY Kratos (self-service auth). """ import os import logging from typing import Optional, Dict, Any, Tuple from datetime import datetime from .workos_provider import WorkOSAuthProvider from .ory_provider import ORYAuthProvider from ..models.user_context import ( AuthResult, UserContext, TokenClaims, UserRole ) from ..identity_sync import IdentitySync logger = logging.getLogger(__name__) class UnifiedAuthProvider: """ Unified authentication provider. Responsibilities: 1. Route authentication requests to appropriate provider 2. Sync WorkOS identities to ORY Kratos 3. Create unified ORY sessions for all auth methods 4. Provide consistent session validation interface """ def __init__( self, workos_provider: Optional[WorkOSAuthProvider] = None, ory_provider: Optional[ORYAuthProvider] = None, identity_sync: Optional[IdentitySync] = None ): """ Initialize unified auth provider. Args: workos_provider: WorkOS SSO provider ory_provider: ORY Kratos provider identity_sync: Identity sync service """ self.workos = workos_provider or WorkOSAuthProvider() self.ory = ory_provider or ORYAuthProvider() self.identity_sync = identity_sync or IdentitySync( workos_provider=self.workos, ory_provider=self.ory ) async def authenticate_sso( self, authorization_code: str, organization_id: str, redirect_uri: Optional[str] = None ) -> AuthResult: """ Authenticate user via WorkOS SSO and create ORY session. Flow: 1. Exchange WorkOS auth code for ID token 2. Validate WorkOS ID token 3. Sync identity to ORY Kratos (create or update) 4. Create ORY session for the identity 5. Return unified AuthResult with ORY session token Args: authorization_code: OAuth authorization code from WorkOS organization_id: WorkOS organization ID redirect_uri: OAuth redirect URI Returns: AuthResult with ORY session token and user context """ try: # Step 1: Authenticate with WorkOS logger.info(f"SSO authentication for org: {organization_id}") workos_result = await self.workos.authenticate( authorization_code=authorization_code, redirect_uri=redirect_uri ) if not workos_result.success: logger.error(f"WorkOS authentication failed: {workos_result.error}") return workos_result # Step 2: Validate WorkOS token workos_claims = await self.workos.validate_token( workos_result.access_token ) if not workos_claims: return AuthResult( success=False, error="invalid_token", error_description="WorkOS token validation failed" ) # Step 3: Sync to ORY Kratos logger.info(f"Syncing SSO user to ORY: {workos_claims.email}") identity_id = await self.identity_sync.sync_from_workos( workos_profile={ 'id': workos_claims.sub, 'email': workos_claims.email, 'organization_id': organization_id, 'metadata': workos_claims.metadata or {} }, firm_id=workos_claims.firm_id ) if not identity_id: return AuthResult( success=False, error="sync_failed", error_description="Failed to sync identity to ORY" ) # Step 4: Create ORY session logger.info(f"Creating ORY session for identity: {identity_id}") ory_session = await self._create_ory_session_for_identity( identity_id=identity_id, auth_method="workos_sso", metadata={ 'workos_org_id': organization_id, 'sso_login_at': datetime.utcnow().isoformat() } ) if not ory_session: return AuthResult( success=False, error="session_creation_failed", error_description="Failed to create ORY session" ) # Step 5: Return unified result logger.info(f"SSO authentication successful for: {identity_id}") return AuthResult( success=True, user_context=ory_session['user_context'], access_token=ory_session['session_token'], expires_in=86400 # 24 hours ) except Exception as e: logger.exception(f"SSO authentication error: {e}") return AuthResult( success=False, error="internal_error", error_description=str(e) ) async def authenticate_password( self, email: str, password: str ) -> AuthResult: """ Authenticate user via email/password (ORY Kratos). Flow: 1. Create ORY login flow 2. Submit credentials 3. Handle MFA if required 4. Return session token Args: email: User email password: User password Returns: AuthResult with ORY session token """ try: logger.info(f"Password authentication for: {email}") # Create login flow flow = await self.ory.create_login_flow() if not flow: return AuthResult( success=False, error="flow_creation_failed" ) # Submit credentials result = await self.ory.submit_login( flow_id=flow['id'], identifier=email, password=password ) if not result.success: logger.warning(f"Password login failed for {email}: {result.error}") return result except Exception as e: logger.exception(f"Password authentication error: {e}") return AuthResult( success=False, error="internal_error", error_description=str(e) ) async def validate_session( self, session_token: str ) -> Optional[UserContext]: """ Validate session token and return user context. Works for both SSO and self-service authenticated users since all sessions are managed by ORY Kratos. Args: session_token: ORY session token Returns: UserContext if session is valid, None otherwise """ try: return await self.ory.get_session(session_token) except Exception as e: logger.warning(f"Session validation error: {e}") return None async def refresh_session( self, session_token: str ) -> Optional[str]: """ Refresh session token. ORY sessions auto-extend on activity, so this is mainly for explicit refresh requests. Args: session_token: Current session token Returns: New session token if successful """ try: # Validate current session user_context = await self.validate_session(session_token) if not user_context: return None # Create new session new_session = await self._create_ory_session_for_identity( identity_id=user_context.user_id, auth_method="session_refresh" ) return new_session.get('session_token') if new_session else None except Exception as e: logger.exception(f"Session refresh error: {e}") return None async def logout( self, session_token: str ): """ Logout user by revoking session. Args: session_token: Session token to revoke """ try: await self.ory.revoke_session(session_token) logger.info("Session revoked successfully") except Exception as e: logger.exception(f"Logout error: {e}") async def register_user( self, email: str, password: str, firm_id: str, role: str = "client", **additional_traits ) -> AuthResult: """ Register new user via ORY Kratos. Args: email: User email password: User password firm_id: Firm ID role: User role **additional_traits: Additional identity traits Returns: AuthResult with session token """ try: logger.info(f"Registering new user: {email} for firm: {firm_id}") # Create identity via Admin API identity_id = await self.ory.create_identity( email=email, firm_id=firm_id, role=role, auth_provider="password", **additional_traits ) if not identity_id: return AuthResult( success=False, error="identity_creation_failed" ) # Set password via Admin API # Note: In production, use registration flow instead # This is a simplified version for admin-created accounts # Create session for new user session = await self._create_ory_session_for_identity( identity_id=identity_id, auth_method="registration" ) if not session: return AuthResult( success=False, error="session_creation_failed" ) logger.info(f"User registered successfully: {identity_id}") return AuthResult( success=True, user_context=session['user_context'], access_token=session['session_token'], expires_in=86400 ) except Exception as e: logger.exception(f"Registration error: {e}") return AuthResult( success=False, error="internal_error", error_description=str(e) ) async def get_authentication_method( self, email: str, firm_id: str ) -> Dict[str, Any]: """ Determine authentication method for user. Args: email: User email firm_id: Firm ID Returns: Dict with authentication method info: { 'method': 'sso' | 'password', 'sso_enabled': bool, 'organization_id': str (if SSO), 'requires_mfa': bool } """ try: # Look up identity in ORY identities = await self.ory.list_identities(firm_id=firm_id) identity = next( (i for i in identities if i['traits']['email'] == email), None ) if not identity: # New user - check if firm has SSO configured # TODO: Check firm SSO settings return { 'method': 'password', 'sso_enabled': False, 'requires_mfa': False } traits = identity['traits'] auth_provider = traits.get('auth_provider', 'password') if auth_provider == 'workos_sso': return { 'method': 'sso', 'sso_enabled': True, 'organization_id': traits.get('workos_org_id'), 'requires_mfa': False # MFA handled by IdP } else: return { 'method': 'password', 'sso_enabled': False, 'requires_mfa': traits.get('mfa_enabled', False) } except Exception as e: logger.exception(f"Error determining auth method: {e}") return { 'method': 'password', 'sso_enabled': False, 'requires_mfa': False } async def _create_ory_session_for_identity( self, identity_id: str, auth_method: str, metadata: Optional[Dict] = None ) -> Optional[Dict[str, Any]]: """ Create ORY session for identity via Admin API. Args: identity_id: ORY identity ID auth_method: Authentication method used metadata: Additional session metadata Returns: Dict with session_token and user_context """ try: # Get identity details identity = await self.ory.get_identity(identity_id) if not identity: logger.error(f"Identity not found: {identity_id}") return None traits = identity['traits'] # Create user context user_context = UserContext( user_id=identity_id, firm_id=traits['firm_id'], role=UserRole(traits['role']), email=traits['email'], metadata={ 'auth_method': auth_method, 'auth_provider': traits.get('auth_provider', 'password'), **(metadata or {}) } ) # Create session via Admin API # POST /admin/identities/{id}/sessions session_data = await self.ory.admin_client.post( f"/admin/identities/{identity_id}/sessions", json={ "expires_in": "24h", "metadata": user_context.metadata } ) session_json = session_data.json() return { 'session_token': session_json['session_token'], 'session_id': session_json['id'], 'user_context': user_context, 'expires_at': session_json['expires_at'] } except Exception as e: logger.exception(f"Session creation error: {e}") return None async def close(self): """Close HTTP clients.""" await self.workos.close() await self.ory.close()