fh-saas 0.9.5__py3-none-any.whl

fh_saas/utils_auth.py

@@ -0,0 +1,647 @@
+"""Multi-user tenant authentication with Google OAuth, CSRF protection, and automatic tenant provisioning."""
+
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_utils_auth.ipynb.
+
+# %% auto 0
+__all__ = ['logger', 'ROLE_HIERARCHY', 'DEFAULT_SKIP_AUTH', 'has_min_role', 'get_user_role', 'require_role',
+           'invalidate_auth_cache', 'create_auth_beforeware', 'get_google_oauth_client', 'generate_oauth_state',
+           'verify_oauth_state', 'create_or_get_global_user', 'get_user_membership', 'verify_membership',
+           'provision_new_user', 'create_user_session', 'get_current_user', 'clear_session', 'route_user_after_login',
+           'require_tenant_access', 'handle_login_request', 'handle_oauth_callback', 'handle_logout']
+
+# %% ../nbs/04_utils_auth.ipynb 2
+from fastsql import *
+from fastcore.utils import *
+from fh_saas.db_host import (
+    timestamp, gen_id, 
+    GlobalUser, TenantCatalog, Membership, HostAuditLog,
+    HostDatabase
+)
+from fh_saas.db_tenant import (
+    get_or_create_tenant_db, 
+    init_tenant_core_schema, 
+    TenantUser
+)
+from fasthtml.oauth import GoogleAppClient, redir_url
+from starlette.responses import RedirectResponse
+import os
+import uuid
+import json
+import logging
+import time
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Module-level logger - configured by app via configure_logging()
+logger = logging.getLogger(__name__)
+
+# %% ../nbs/04_utils_auth.ipynb 6
+# Role hierarchy: higher number = more permissions
+ROLE_HIERARCHY = {
+    'admin': 3,
+    'editor': 2,
+    'viewer': 1
+}
+
+def has_min_role(user: dict, required_role: str) -> bool:
+    """Check if user meets the minimum role requirement.
+    
+    Args:
+        user: User dict with 'role' field (from request.state.user)
+        required_role: Minimum role needed ('admin', 'editor', 'viewer')
+    
+    Returns:
+        True if user's role >= required_role in hierarchy
+        
+    Example:
+        >>> user = {'role': 'editor'}
+        >>> has_min_role(user, 'viewer')  # True - editor > viewer
+        >>> has_min_role(user, 'admin')   # False - editor < admin
+    """
+    user_role = user.get('role', 'viewer')
+    user_level = ROLE_HIERARCHY.get(user_role, 0)
+    required_level = ROLE_HIERARCHY.get(required_role, 0)
+    return user_level >= required_level
+
+
+def get_user_role(session: dict, tenant_db: 'Database' = None) -> str:
+    """Derive effective role from session and tenant database.
+    
+    Rules:
+    1. Tenant owner (session['tenant_role'] == 'owner') → 'admin'
+    2. System admin → 'admin'
+    3. Otherwise → lookup TenantUser.local_role from tenant DB
+    4. Fallback → None (user must be explicitly assigned a role)
+    
+    Args:
+        session: User session dict
+        tenant_db: Tenant database connection (optional)
+    
+    Returns:
+        Effective role string: 'admin', 'editor', 'viewer', or None
+    """
+    # Owner is always admin
+    if session.get('tenant_role') == 'owner':
+        return 'admin'
+    
+    # System admin is always admin
+    if session.get('is_sys_admin'):
+        return 'admin'
+    
+    # Look up local_role from TenantUser
+    if tenant_db:
+        user_id = session.get('user_id')
+        if user_id:
+            try:
+                tenant_db.conn.rollback()
+                tenant_users = tenant_db.t.core_tenant_users
+                all_users = tenant_users()
+                matching = [u for u in all_users if u.id == user_id]
+                if matching:
+                    return matching[0].local_role
+            except Exception as e:
+                logger.warning(f"Failed to lookup TenantUser role: {e}")
+    
+    # No role found - user must be assigned explicitly
+    return None
+
+# %% ../nbs/04_utils_auth.ipynb 9
+from functools import wraps
+from starlette.responses import Response
+
+def require_role(min_role: str):
+    """Decorator to protect routes with minimum role requirement.
+    
+    Args:
+        min_role: Minimum role required ('admin', 'editor', 'viewer')
+    
+    Returns:
+        Decorator that checks request.state.user['role']
+        Returns 403 Forbidden if user lacks required role
+        
+    Example:
+        >>> @app.get('/admin/settings')
+        >>> @require_role('admin')
+        >>> def admin_settings(request):
+        ...     return "Admin only content"
+        
+        >>> @app.get('/reports')  
+        >>> @require_role('viewer')  # All authenticated users
+        >>> def view_reports(request):
+        ...     return "Reports"
+    """
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(request, *args, **kwargs):
+            user = getattr(request.state, 'user', None)
+            
+            if not user:
+                return Response(
+                    content="Authentication required",
+                    status_code=401
+                )
+            
+            if not user.get('role'):
+                return Response(
+                    content="No role assigned. Contact your administrator.",
+                    status_code=403
+                )
+            
+            if not has_min_role(user, min_role):
+                return Response(
+                    content=f"Access denied. Required role: {min_role}",
+                    status_code=403
+                )
+            
+            # Handle both sync and async route functions
+            result = func(request, *args, **kwargs)
+            if hasattr(result, '__await__'):
+                return await result
+            return result
+        
+        return wrapper
+    return decorator
+
+# %% ../nbs/04_utils_auth.ipynb 12
+# Session cache key
+_AUTH_CACHE_KEY = '_auth_cache'
+
+def invalidate_auth_cache(session: dict):
+    """Clear the auth cache from session.
+    
+    Call this when:
+    - User role or permissions change
+    - User is added/removed from tenant
+    - Admin changes user's local_role
+    
+    Args:
+        session: User session dict
+        
+    Example:
+        >>> # After admin changes user role
+        >>> tenant_user.local_role = 'editor'
+        >>> tables['tenant_users'].update(tenant_user)
+        >>> invalidate_auth_cache(session)  # Force fresh lookup
+    """
+    session.pop(_AUTH_CACHE_KEY, None)
+    logger.debug("Auth cache invalidated")
+
+
+def _get_cached_auth(session: dict, cache_ttl: int) -> dict | None:
+    """Get cached auth data if valid and not expired.
+    
+    Args:
+        session: User session dict
+        cache_ttl: Cache time-to-live in seconds
+        
+    Returns:
+        Cached auth dict with 'user' and 'tenant_id', or None if expired/missing
+    """
+    cache = session.get(_AUTH_CACHE_KEY)
+    if not cache:
+        return None
+    
+    cached_at = cache.get('cached_at', 0)
+    if time.time() - cached_at > cache_ttl:
+        logger.debug("Auth cache expired")
+        return None
+    
+    return cache
+
+
+def _set_auth_cache(session: dict, user: dict, tenant_id: str):
+    """Store auth data in session cache.
+    
+    Args:
+        session: User session dict
+        user: User dict to cache
+        tenant_id: Tenant ID to cache
+    """
+    session[_AUTH_CACHE_KEY] = {
+        'user': user.copy(),
+        'tenant_id': tenant_id,
+        'cached_at': time.time()
+    }
+    logger.debug(f"Auth cache set for user {user.get('user_id')}")
+
+# %% ../nbs/04_utils_auth.ipynb 17
+from fasthtml.common import Beforeware
+from typing import Callable, Any
+
+# Default URL patterns to skip authentication
+DEFAULT_SKIP_AUTH = [
+    r'/login.*',
+    r'/logout',
+    r'/oauth/callback',
+    r'/auth/callback',
+    r'/health',
+    r'/favicon\.ico',
+    r'/static/.*',
+]
+
+def create_auth_beforeware(
+    redirect_path: str = '/login',
+    session_key: str = 'user_id',
+    skip: list[str] = None,
+    include_defaults: bool = True,
+    setup_tenant_db: bool = True,
+    schema_init: Callable[[Database], dict[str, Any]] = None,
+    session_cache: bool = False,
+    session_cache_ttl: int = 300,
+):
+    """Create Beforeware that checks for authenticated session and sets up request.state.
+    
+    Args:
+        redirect_path: Where to redirect unauthenticated users
+        session_key: Session key for user ID
+        skip: List of regex patterns to skip auth
+        include_defaults: Include default skip patterns
+        setup_tenant_db: Auto-setup tenant database on request.state
+        schema_init: Optional callback to initialize tables dict.
+                     Signature: (tenant_db: Database) -> dict[str, Table]
+                     Result stored in request.state.tables
+        session_cache: Enable caching user dict in session to reduce DB queries.
+                       Recommended for HTMX-heavy apps. Default: False
+        session_cache_ttl: Cache TTL in seconds. Default: 300 (5 minutes)
+    
+    Returns:
+        Beforeware instance for FastHTML apps
+        
+    Sets on request.state:
+        - user: dict with user_id, email, tenant_id, role, is_owner
+        - tenant_id: str
+        - tenant_db: Database connection
+        - tables: dict of Table objects (if schema_init provided)
+        
+    Example:
+        >>> # Basic usage
+        >>> beforeware = create_auth_beforeware()
+        
+        >>> # With session caching for HTMX apps
+        >>> beforeware = create_auth_beforeware(
+        ...     session_cache=True,
+        ...     session_cache_ttl=300
+        ... )
+        
+        >>> # With schema initialization
+        >>> def get_app_tables(db):
+        ...     return {'users': db.create(User, pk='id')}
+        >>> beforeware = create_auth_beforeware(schema_init=get_app_tables)
+    """
+    skip_patterns = []
+    if include_defaults:
+        skip_patterns.extend(DEFAULT_SKIP_AUTH)
+    if skip:
+        skip_patterns.extend(skip)
+    
+    def check_auth(req, sess):
+        if session_key not in sess:
+            return RedirectResponse(redirect_path, status_code=303)
+        
+        if setup_tenant_db:
+            # Check session cache first (if enabled)
+            cache_hit = False
+            if session_cache:
+                cached = _get_cached_auth(sess, session_cache_ttl)
+                if cached:
+                    cache_hit = True
+                    user = cached['user']
+                    req.state.user = user
+                    req.state.tenant_id = cached['tenant_id']
+                    
+                    # Still need to create tenant_db connection (lightweight)
+                    if user.get('tenant_id') and not user.get('is_sys_admin'):
+                        try:
+                            req.state.tenant_db = get_or_create_tenant_db(user['tenant_id'])
+                        except Exception as e:
+                            logger.error(f"Failed to setup tenant_db from cache: {e}")
+                            req.state.tenant_db = None
+                    else:
+                        req.state.tenant_db = None
+            
+            # Cache miss - do full DB lookup
+            if not cache_hit:
+                user = get_current_user(sess)
+                if user:
+                    req.state.tenant_id = user.get('tenant_id')
+                    req.state.tenant_db = None
+                    
+                    if user.get('tenant_id') and not user.get('is_sys_admin'):
+                        try:
+                            host_db = HostDatabase.from_env()
+                            if verify_membership(host_db, user['user_id'], user['tenant_id']):
+                                req.state.tenant_db = get_or_create_tenant_db(user['tenant_id'])
+                            else:
+                                logger.warning(f"Invalid membership for user {user['user_id']}")
+                        except Exception as e:
+                            logger.error(f"Failed to setup tenant_db: {e}")
+                    
+                    # Derive effective role from session + tenant DB
+                    role = get_user_role(sess, req.state.tenant_db)
+                    user['role'] = role
+                    user['is_owner'] = sess.get('tenant_role') == 'owner'
+                    req.state.user = user
+                    
+                    # Update session cache (if enabled)
+                    if session_cache:
+                        _set_auth_cache(sess, user, user.get('tenant_id'))
+            
+            # Auto-initialize tables if schema_init provided
+            if schema_init and req.state.tenant_db:
+                try:
+                    req.state.tables = schema_init(req.state.tenant_db)
+                except Exception as e:
+                    logger.error(f"Failed to initialize schema: {e}")
+                    req.state.tables = {}
+            else:
+                req.state.tables = {}
+    
+    return Beforeware(check_auth, skip=skip_patterns)
+
+# %% ../nbs/04_utils_auth.ipynb 21
+def get_google_oauth_client():
+    """Initialize Google OAuth client with credentials from environment."""
+    client_id = os.getenv('GOOGLE_CLIENT_ID')
+    client_secret = os.getenv('GOOGLE_CLIENT_SECRET')
+    
+    if not client_id or not client_secret:
+        raise ValueError("Missing GOOGLE_CLIENT_ID or GOOGLE_CLIENT_SECRET in .env")
+    
+    return GoogleAppClient(client_id=client_id, client_secret=client_secret)
+
+# %% ../nbs/04_utils_auth.ipynb 24
+def generate_oauth_state():
+    """Generate cryptographically secure random state token for CSRF protection."""
+    return uuid.uuid4().hex
+
+
+def verify_oauth_state(session: dict, callback_state: str):
+    """Verify OAuth callback state matches stored session state (CSRF protection)."""
+    stored_state = session.get('oauth_state')
+    
+    if not stored_state:
+        raise ValueError("CSRF validation failed: No state in session")
+    
+    if stored_state != callback_state:
+        raise ValueError("CSRF validation failed: State mismatch")
+    
+    session.pop('oauth_state', None)  # One-time use
+
+# %% ../nbs/04_utils_auth.ipynb 28
+def create_or_get_global_user(host_db: HostDatabase, oauth_id: str, email: str, oauth_info: dict = None):
+    """Create or retrieve GlobalUser from host database."""
+    try:
+        host_db.rollback()
+        all_users = host_db.global_users()
+        existing = [u for u in all_users if u.oauth_id == oauth_id]
+        
+        if existing:
+            user = existing[0]
+            user.last_login = timestamp()
+            host_db.global_users.update(user)
+            logger.info(f'User login: {email}', extra={'user_id': user.id, 'email': email})
+            return user
+        
+        new_user = GlobalUser(
+            id=gen_id(),
+            email=email,
+            oauth_id=oauth_id,
+            created_at=timestamp(),
+            last_login=timestamp()
+        )
+        host_db.global_users.insert(new_user)
+        logger.info(f'New user created: {email}', extra={'user_id': new_user.id, 'email': email})
+        return new_user
+        
+    except Exception as e:
+        host_db.rollback()
+        logger.error(f'Failed to create/get user {email}: {e}', exc_info=True)
+        raise
+
+
+def get_user_membership(host_db: HostDatabase, user_id: str):
+    """Get single active membership for user."""
+    host_db.rollback()
+    all_memberships = host_db.memberships()
+    active = [m for m in all_memberships if m.user_id == user_id and m.is_active]
+    return active[0] if active else None
+
+
+def verify_membership(host_db: HostDatabase, user_id: str, tenant_id: str) -> bool:
+    """Verify user has active membership for specific tenant."""
+    host_db.rollback()
+    all_memberships = host_db.memberships()
+    valid = [m for m in all_memberships if m.user_id == user_id and m.tenant_id == tenant_id and m.is_active]
+    return len(valid) > 0
+
+# %% ../nbs/04_utils_auth.ipynb 33
+def provision_new_user(host_db: HostDatabase, global_user: GlobalUser) -> str:
+    """Auto-provision new tenant for first-time user."""
+    tenant_id = gen_id()
+    username = global_user.email.split('@')[0]
+    tenant_name = f"{username}'s Workspace"
+    tenant_db = None
+    
+    try:
+        logger.info(f'Starting tenant provisioning for {global_user.email}',
+                   extra={'tenant_id': tenant_id, 'user_id': global_user.id})
+        
+        # Create physical tenant database and register in catalog
+        tenant_db = get_or_create_tenant_db(tenant_id, tenant_name)
+        
+        # Initialize core tenant schema
+        core_tables = init_tenant_core_schema(tenant_db)
+        
+        # Create TenantUser profile
+        tenant_user = TenantUser(
+            id=global_user.id,
+            display_name=username,
+            local_role='admin',
+            created_at=timestamp()
+        )
+        core_tables['tenant_users'].insert(tenant_user)
+        tenant_db.conn.commit()  # Commit tenant changes
+        
+        # Create membership in host database
+        membership = Membership(
+            id=gen_id(),
+            user_id=global_user.id,
+            tenant_id=tenant_id,
+            profile_id=global_user.id,
+            role='owner',
+            created_at=timestamp()
+        )
+        host_db.memberships.insert(membership)
+        
+        # Log provisioning event
+        audit_log = HostAuditLog(
+            id=gen_id(),
+            actor_user_id=global_user.id,
+            event_type='tenant_provisioned',
+            target_id=tenant_id,
+            details=json.dumps({'tenant_name': tenant_name, 'plan_tier': 'free', 'user_email': global_user.email}),
+            created_at=timestamp()
+        )
+        host_db.audit_logs.insert(audit_log)
+        
+        host_db.commit()
+        logger.info(f'Tenant provisioned: {tenant_name}',
+                   extra={'tenant_id': tenant_id, 'tenant_name': tenant_name, 'user_id': global_user.id})
+        return tenant_id
+        
+    except Exception as e:
+        host_db.rollback()
+        if tenant_db:
+            try:
+                tenant_db.conn.rollback()
+            except Exception:
+                pass
+        logger.error(f'Tenant provisioning failed for {global_user.email}: {e}',
+                    extra={'tenant_id': tenant_id, 'user_id': global_user.id}, exc_info=True)
+        raise Exception(f"Failed to provision tenant for {global_user.email}: {str(e)}") from e
+    finally:
+        # Always close tenant_db connection to prevent connection leaks
+        if tenant_db:
+            try:
+                tenant_db.conn.close()
+                tenant_db.engine.dispose()
+            except Exception:
+                pass
+
+# %% ../nbs/04_utils_auth.ipynb 36
+def create_user_session(session: dict, global_user: GlobalUser, membership: Membership):
+    """Create authenticated session after successful OAuth login."""
+    session['user_id'] = global_user.id
+    session['email'] = global_user.email
+    session['tenant_id'] = membership.tenant_id
+    session['tenant_role'] = membership.role
+    session['is_sys_admin'] = global_user.is_sys_admin
+    session['login_at'] = timestamp()
+
+
+def get_current_user(session: dict) -> dict | None:
+    """Extract current user info from session.
+    
+    Returns:
+        dict with keys: user_id, email, tenant_id, tenant_role, is_sys_admin
+        
+    Note:
+        The 'role' and 'is_owner' fields are added by create_auth_beforeware
+        after deriving the effective role from TenantUser.local_role.
+        Access via request.state.user['role'] in routes.
+    """
+    if 'user_id' not in session:
+        return None
+    
+    return {
+        'user_id': session.get('user_id'),
+        'email': session.get('email'),
+        'tenant_id': session.get('tenant_id'),
+        'tenant_role': session.get('tenant_role'),
+        'is_sys_admin': session.get('is_sys_admin', False)
+    }
+
+
+def clear_session(session: dict):
+    """Clear all session data (logout)."""
+    session.clear()
+
+# %% ../nbs/04_utils_auth.ipynb 41
+def route_user_after_login(global_user: GlobalUser, membership: Membership = None) -> str:
+    """Determine redirect URL based on user type and membership."""
+    if global_user.is_sys_admin:
+        return '/admin/dashboard'
+    if membership:
+        return '/dashboard'
+    raise ValueError(f"User {global_user.email} has no membership")
+
+
+def require_tenant_access(request_or_session):
+    """Get tenant database with membership validation."""
+    # Check if request.state.tenant_db already set by beforeware
+    if hasattr(request_or_session, 'state'):
+        if hasattr(request_or_session.state, 'tenant_db') and request_or_session.state.tenant_db:
+            return request_or_session.state.tenant_db
+        session = getattr(request_or_session, 'session', {})
+    else:
+        session = request_or_session
+    
+    user = get_current_user(session)
+    if not user:
+        raise ValueError("Authentication required")
+    
+    host_db = HostDatabase.from_env()
+    
+    if not verify_membership(host_db, user['user_id'], user['tenant_id']):
+        raise PermissionError(f"Access denied for user {user['user_id']} to tenant {user['tenant_id']}")
+    
+    return get_or_create_tenant_db(user['tenant_id'])
+
+# %% ../nbs/04_utils_auth.ipynb 45
+def handle_login_request(request, session):
+    """Generate Google OAuth URL with CSRF state protection."""
+    logger.debug('Login request initiated')
+    
+    # Generate CSRF state token
+    state = generate_oauth_state()
+    session['oauth_state'] = state
+    
+    # Get OAuth client and generate login link
+    client = get_google_oauth_client()
+    redirect_uri = redir_url(request, '/auth/callback')
+    login_link = client.login_link(redirect_uri=redirect_uri, state=state)
+    
+    return login_link
+
+
+def handle_oauth_callback(code: str, state: str, request, session):
+    """Complete OAuth flow: CSRF verify → user info → provision → session → redirect."""
+    logger.debug('OAuth callback received')
+    
+    # Step 1: CSRF validation (CRITICAL - must be first)
+    verify_oauth_state(session, state)
+    
+    # Step 2: Exchange authorization code for user info
+    client = get_google_oauth_client()
+    redirect_uri = redir_url(request, '/auth/callback')
+    user_info = client.retr_info(code, redirect_uri)
+    
+    # Step 3: Get host database instance (singleton - no need to close)
+    host_db = HostDatabase.from_env()
+    
+    # Step 4: Create or get GlobalUser
+    oauth_id = user_info[client.id_key]  # Google 'sub' field
+    email = user_info.get('email', '')
+    global_user = create_or_get_global_user(host_db, oauth_id, email, user_info)
+    
+    # Step 5: Check for existing membership
+    membership = get_user_membership(host_db, global_user.id)
+    
+    # Step 6: Auto-provision if new user (no membership)
+    if not membership and not global_user.is_sys_admin:
+        tenant_id = provision_new_user(host_db, global_user)
+        membership = get_user_membership(host_db, global_user.id)
+    
+    # Step 7: Create session (skip for sys admin - no tenant)
+    if membership:
+        create_user_session(session, global_user, membership)
+    else:
+        # System admin - minimal session
+        session['user_id'] = global_user.id
+        session['email'] = global_user.email
+        session['is_sys_admin'] = True
+        session['login_at'] = timestamp()
+    
+    # Step 8: Route to appropriate dashboard
+    redirect_url = route_user_after_login(global_user, membership)
+    logger.info(f'OAuth complete, redirecting to {redirect_url}', extra={'email': email})
+    return RedirectResponse(redirect_url, status_code=303)
+
+
+def handle_logout(session):
+    """Clear session and redirect to login page."""
+    email = session.get('email', 'unknown')
+    clear_session(session)
+    logger.info(f'User logged out: {email}')
+    return RedirectResponse('/login', status_code=303)