security/tenant_scope.py

"""
Tenant Scope Enforcement for AegisLM Multi-Tenant System

Provides tenant isolation utilities and context management.
"""

import uuid
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional
from dataclasses import dataclass, field


@dataclass
class TenantScope:
    """
    Represents the current tenant scope for request processing.
    
    This ensures all database queries are properly scoped to the 
    current tenant to prevent cross-tenant data leakage.
    """
    tenant_id: uuid.UUID
    active: bool = True
    
    def filter_query(self, table_name: str) -> str:
        """
        Get the tenant filter SQL for a given table.
        
        This should be used in all raw SQL queries to ensure
        tenant isolation.
        """
        return f"tenant_id = '{self.tenant_id}'"
    
    def get_tenant_filter(self) -> dict:
        """Get the tenant filter as a dictionary for ORM queries."""
        return {"tenant_id": self.tenant_id}


# Thread-local storage for tenant context
_tenant_context: Optional[TenantScope] = None


def get_tenant_scope() -> Optional[TenantScope]:
    """
    Get the current tenant scope from context.
    
    Returns None if no tenant context is set (e.g., for system tasks).
    """
    return _tenant_context


def set_tenant_scope(scope: TenantScope) -> None:
    """
    Set the current tenant scope in context.
    
    This should be called after authentication to set the tenant
    for the current request processing.
    """
    global _tenant_context
    _tenant_context = scope


def clear_tenant_scope() -> None:
    """Clear the current tenant scope."""
    global _tenant_context
    _tenant_context = None


@asynccontextmanager
async def tenant_context(tenant_id: uuid.UUID) -> AsyncGenerator[TenantScope, None]:
    """
    Context manager for tenant scope.
    
    Usage:
        async with tenant_context(request.tenant_id):
            # All queries will be scoped to this tenant
            results = await get_evaluation_runs()
    """
    scope = TenantScope(tenant_id=tenant_id)
    old_context = _tenant_context
    try:
        set_tenant_scope(scope)
        yield scope
    finally:
        clear_tenant_scope()
        if old_context is not None:
            set_tenant_scope(old_context)


def require_tenant(tenant_id: uuid.UUID) -> TenantScope:
    """
    Require a tenant scope to be set.
    
    Raises an exception if no tenant context is active.
    """
    scope = get_tenant_scope()
    if scope is None:
        raise RuntimeError("No tenant context set")
    
    if scope.tenant_id != tenant_id:
        raise PermissionError(f"Access denied: tenant {tenant_id} does not match context {scope.tenant_id}")
    
    return scope


class TenantFilter:
    """
    Utility class for adding tenant filters to queries.
    """
    
    @staticmethod
    def for_evaluation_runs(tenant_id: uuid.UUID) -> dict:
        """Get filter for evaluation_runs table."""
        return {"tenant_id": tenant_id}
    
    @staticmethod
    def for_evaluation_results(tenant_id: uuid.UUID) -> dict:
        """Get filter for evaluation_results table - requires join through runs."""
        return {"run": {"tenant_id": tenant_id}}
    
    @staticmethod
    def for_monitoring_metrics(tenant_id: uuid.UUID) -> dict:
        """Get filter for monitoring_metrics table."""
        return {"tenant_id": tenant_id}
    
    @staticmethod
    def for_alerts(tenant_id: uuid.UUID) -> dict:
        """Get filter for alerts table."""
        return {"tenant_id": tenant_id}
    
    @staticmethod
    def for_release_validations(tenant_id: uuid.UUID) -> dict:
        """Get filter for release_validations table."""
        return {"tenant_id": tenant_id}
    
    @staticmethod
    def for_api_keys(tenant_id: uuid.UUID) -> dict:
        """Get filter for api_keys table."""
        return {"tenant_id": tenant_id}
    
    @staticmethod
    def for_users(tenant_id: uuid.UUID) -> dict:
        """Get filter for users table."""
        return {"tenant_id": tenant_id}


def validate_tenant_access(resource_tenant_id: uuid.UUID, user_tenant_id: uuid.UUID) -> bool:
    """
    Validate that a user can access a resource based on tenant.
    
    This is a safety check to prevent cross-tenant access.
    Returns True if access is allowed, False otherwise.
    """
    return resource_tenant_id == user_tenant_id


def get_safe_tenant_filter(tenant_id: Optional[uuid.UUID]) -> dict:
    """
    Get a tenant filter that handles None gracefully.
    
    If tenant_id is None, returns an empty filter (for system tasks).
    """
    if tenant_id is None:
        return {}
    return {"tenant_id": tenant_id}