kailash 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

kailash/middleware/auth/models.py

@@ -0,0 +1,137 @@
+"""
+Authentication Models for Kailash Middleware
+
+Provides data models for JWT authentication without any circular dependencies.
+These models can be imported anywhere in the codebase safely.
+"""
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import List, Optional
+
+
+@dataclass
+class JWTConfig:
+    """Configuration for JWT authentication supporting both HS256 and RSA algorithms."""
+
+    # Signing configuration
+    algorithm: str = "HS256"  # Default to HS256 for simplicity
+    secret_key: Optional[str] = None  # For HS256
+    use_rsa: bool = False  # Enable RSA mode
+    private_key: Optional[str] = None  # For RSA (PEM format)
+    public_key: Optional[str] = None  # For RSA (PEM format)
+
+    # Token expiration
+    access_token_expire_minutes: int = 15
+    refresh_token_expire_days: int = 7
+
+    # Security settings
+    issuer: str = "kailash-middleware"
+    audience: str = "kailash-api"
+
+    # Key management
+    auto_generate_keys: bool = True
+    key_rotation_days: int = 30  # Only applies to RSA mode
+
+    # Token settings
+    include_user_claims: bool = True
+    include_permissions: bool = True
+    max_refresh_count: int = 10
+
+    # Security features
+    enable_blacklist: bool = True
+    enable_token_cleanup: bool = True
+    cleanup_interval_minutes: int = 60
+
+
+@dataclass
+class TokenPayload:
+    """JWT token payload structure."""
+
+    # Standard claims
+    sub: str  # Subject (user ID)
+    iss: str  # Issuer
+    aud: str  # Audience
+    exp: int  # Expiration time
+    iat: int  # Issued at
+    jti: str  # JWT ID
+
+    # Custom claims
+    tenant_id: Optional[str] = None
+    session_id: Optional[str] = None
+    user_type: str = "user"
+    permissions: List[str] = None
+    roles: List[str] = None
+
+    # Token metadata
+    token_type: str = "access"  # access, refresh
+    refresh_count: int = 0
+
+    def __post_init__(self):
+        if self.permissions is None:
+            self.permissions = []
+        if self.roles is None:
+            self.roles = []
+
+
+@dataclass
+class TokenPair:
+    """Access and refresh token pair."""
+
+    access_token: str
+    refresh_token: str
+    token_type: str = "Bearer"
+    expires_in: int = 0
+    expires_at: Optional[datetime] = None
+    scope: Optional[str] = None
+
+
+@dataclass
+class RefreshTokenData:
+    """Metadata for tracking refresh tokens."""
+
+    jti: str  # Token ID
+    user_id: str
+    tenant_id: Optional[str] = None
+    session_id: Optional[str] = None
+    created_at: datetime = None
+    last_used: Optional[datetime] = None
+    refresh_count: int = 0
+    ip_address: Optional[str] = None
+    user_agent: Optional[str] = None
+
+    def __post_init__(self):
+        if self.created_at is None:
+            self.created_at = datetime.utcnow()
+
+
+@dataclass
+class UserClaims:
+    """User claims for JWT tokens."""
+
+    user_id: str
+    tenant_id: Optional[str] = None
+    email: Optional[str] = None
+    username: Optional[str] = None
+    roles: List[str] = None
+    permissions: List[str] = None
+    metadata: dict = None
+
+    def __post_init__(self):
+        if self.roles is None:
+            self.roles = []
+        if self.permissions is None:
+            self.permissions = []
+        if self.metadata is None:
+            self.metadata = {}
+
+
+@dataclass
+class AuthenticationResult:
+    """Result of authentication attempt."""
+
+    success: bool
+    token_pair: Optional[TokenPair] = None
+    user_claims: Optional[UserClaims] = None
+    error: Optional[str] = None
+    error_code: Optional[str] = None

kailash/middleware/auth/utils.py

@@ -0,0 +1,257 @@
+"""
+Authentication Utilities for Kailash Middleware
+
+Provides helper functions for authentication without circular dependencies.
+"""
+
+import base64
+import secrets
+import string
+from datetime import datetime, timedelta, timezone
+from typing import Dict, Optional, Tuple
+
+
+def generate_secret_key(length: int = 32) -> str:
+    """
+    Generate a secure random secret key for HS256.
+
+    Args:
+        length: Length of the secret key (default: 32)
+
+    Returns:
+        URL-safe base64 encoded secret key
+    """
+    return secrets.token_urlsafe(length)
+
+
+def generate_key_pair() -> Tuple[str, str]:
+    """
+    Generate RSA key pair for RS256.
+
+    Returns:
+        Tuple of (private_key_pem, public_key_pem)
+    """
+    try:
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives import serialization
+        from cryptography.hazmat.primitives.asymmetric import rsa
+
+        # Generate private key
+        private_key = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048, backend=default_backend()
+        )
+
+        # Serialize private key
+        private_pem = private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        ).decode("utf-8")
+
+        # Get public key
+        public_key = private_key.public_key()
+
+        # Serialize public key
+        public_pem = public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        ).decode("utf-8")
+
+        return private_pem, public_pem
+
+    except ImportError:
+        raise ImportError(
+            "RSA key generation requires 'cryptography' package. "
+            "Install with: pip install cryptography"
+        )
+
+
+def calculate_token_expiry(
+    token_type: str = "access", access_minutes: int = 15, refresh_days: int = 7
+) -> datetime:
+    """
+    Calculate token expiration time.
+
+    Args:
+        token_type: Type of token ("access" or "refresh")
+        access_minutes: Minutes until access token expires
+        refresh_days: Days until refresh token expires
+
+    Returns:
+        Expiration datetime in UTC
+    """
+    now = datetime.now(timezone.utc)
+
+    if token_type == "access":
+        return now + timedelta(minutes=access_minutes)
+    else:  # refresh
+        return now + timedelta(days=refresh_days)
+
+
+def is_token_expired(exp_timestamp: int) -> bool:
+    """
+    Check if token has expired based on exp claim.
+
+    Args:
+        exp_timestamp: Expiration timestamp from token
+
+    Returns:
+        True if token has expired
+    """
+    now = datetime.now(timezone.utc)
+    exp_datetime = datetime.fromtimestamp(exp_timestamp, timezone.utc)
+    return now > exp_datetime
+
+
+def generate_jti() -> str:
+    """
+    Generate unique JWT ID.
+
+    Returns:
+        Unique identifier for JWT
+    """
+    return secrets.token_urlsafe(16)
+
+
+def encode_for_jwks(number: int) -> str:
+    """
+    Encode integer for JWKS format.
+
+    Args:
+        number: Integer to encode (e.g., RSA modulus or exponent)
+
+    Returns:
+        Base64url encoded string without padding
+    """
+    byte_length = (number.bit_length() + 7) // 8
+    number_bytes = number.to_bytes(byte_length, "big")
+    return base64.urlsafe_b64encode(number_bytes).decode("ascii").rstrip("=")
+
+
+def validate_algorithm(algorithm: str) -> bool:
+    """
+    Validate JWT algorithm.
+
+    Args:
+        algorithm: Algorithm name
+
+    Returns:
+        True if algorithm is supported
+    """
+    supported = ["HS256", "HS384", "HS512", "RS256", "RS384", "RS512"]
+    return algorithm in supported
+
+
+def parse_bearer_token(authorization_header: str) -> Optional[str]:
+    """
+    Extract token from Authorization header.
+
+    Args:
+        authorization_header: Value of Authorization header
+
+    Returns:
+        Token string or None if invalid format
+    """
+    if not authorization_header:
+        return None
+
+    parts = authorization_header.split()
+
+    if len(parts) != 2 or parts[0].lower() != "bearer":
+        return None
+
+    return parts[1]
+
+
+def generate_random_password(
+    length: int = 16,
+    include_uppercase: bool = True,
+    include_lowercase: bool = True,
+    include_digits: bool = True,
+    include_symbols: bool = True,
+) -> str:
+    """
+    Generate a random password.
+
+    Args:
+        length: Password length
+        include_uppercase: Include uppercase letters
+        include_lowercase: Include lowercase letters
+        include_digits: Include digits
+        include_symbols: Include symbols
+
+    Returns:
+        Random password string
+    """
+    characters = ""
+
+    if include_uppercase:
+        characters += string.ascii_uppercase
+    if include_lowercase:
+        characters += string.ascii_lowercase
+    if include_digits:
+        characters += string.digits
+    if include_symbols:
+        characters += string.punctuation
+
+    if not characters:
+        characters = string.ascii_letters + string.digits
+
+    return "".join(secrets.choice(characters) for _ in range(length))
+
+
+def hash_token_for_storage(token: str) -> str:
+    """
+    Hash token for secure storage (e.g., in blacklist).
+
+    Args:
+        token: JWT token to hash
+
+    Returns:
+        SHA256 hash of token
+    """
+    import hashlib
+
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+def create_jwks_response(
+    public_key_pem: str, key_id: str, algorithm: str = "RS256"
+) -> Dict:
+    """
+    Create JWKS response for public key endpoint.
+
+    Args:
+        public_key_pem: Public key in PEM format
+        key_id: Key identifier
+        algorithm: Algorithm used
+
+    Returns:
+        JWKS formatted response
+    """
+    try:
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives import serialization
+
+        # Load public key
+        public_key = serialization.load_pem_public_key(
+            public_key_pem.encode(), backend=default_backend()
+        )
+
+        # Get public numbers
+        public_numbers = public_key.public_numbers()
+
+        return {
+            "keys": [
+                {
+                    "kty": "RSA",
+                    "kid": key_id,
+                    "use": "sig",
+                    "alg": algorithm,
+                    "n": encode_for_jwks(public_numbers.n),
+                    "e": encode_for_jwks(public_numbers.e),
+                }
+            ]
+        }
+    except Exception:
+        return {"keys": []}

kailash/middleware/communication/api_gateway.py

@@ -33,7 +33,6 @@ from ...nodes.security import CredentialManagerNode
 from ...nodes.transform import DataTransformer
 from ...workflow import Workflow
 from ...workflow.builder import WorkflowBuilder
-from ..auth import KailashJWTAuthManager
 from ..core.agent_ui import AgentUIMiddleware
 from ..core.schema import DynamicSchemaRegistry
 from .events import EventFilter, EventType
@@ -41,8 +40,8 @@ from .realtime import RealtimeMiddleware
 
 logger = logging.getLogger(__name__)
 
-# Use SDK Auth Manager instead of manual security
-auth_manager = KailashJWTAuthManager()
+# Auth manager will be injected via dependency injection
+# This avoids circular imports and allows for flexible auth implementations
 
 
 # Pydantic Models
@@ -135,8 +134,23 @@ class APIGateway:
         enable_docs: bool = True,
         max_sessions: int = 1000,
         enable_auth: bool = True,
+        auth_manager=None,  # Dependency injection for auth
         database_url: str = None,
     ):
+        """
+        Initialize API Gateway with dependency injection support.
+
+        Args:
+            title: API title
+            description: API description
+            version: API version
+            cors_origins: Allowed CORS origins
+            enable_docs: Enable OpenAPI documentation
+            max_sessions: Maximum concurrent sessions
+            enable_auth: Enable authentication
+            auth_manager: Optional auth manager instance (creates default if None and auth enabled)
+            database_url: Optional database URL for persistence
+        """
         self.title = title
         self.version = version
         self.enable_docs = enable_docs
@@ -153,7 +167,21 @@ class APIGateway:
 
         # Initialize auth manager if enabled
         if enable_auth:
-            self.auth_manager = KailashJWTAuthManager(secret_key="api-gateway-secret")
+            if auth_manager is None:
+                # Create default auth manager if none provided
+                # Import here to avoid circular dependency
+                from ..auth import JWTAuthManager
+
+                self.auth_manager = JWTAuthManager(
+                    secret_key="api-gateway-secret",
+                    algorithm="HS256",
+                    issuer="kailash-gateway",
+                    audience="kailash-api",
+                )
+            else:
+                self.auth_manager = auth_manager
+        else:
+            self.auth_manager = None
 
         # Create FastAPI app with lifespan management
         @asynccontextmanager
@@ -774,25 +802,39 @@ class APIGateway:
 
 # Convenience function for quick setup
 def create_gateway(
-    agent_ui_middleware: AgentUIMiddleware = None, **kwargs
+    agent_ui_middleware: AgentUIMiddleware = None, auth_manager=None, **kwargs
 ) -> APIGateway:
     """
-    Create a configured API gateway instance.
+    Create a configured API gateway instance with dependency injection.
 
     Args:
         agent_ui_middleware: Optional existing AgentUIMiddleware instance
+        auth_manager: Optional auth manager instance (e.g., JWTAuthManager)
         **kwargs: Additional arguments for APIGateway initialization
 
     Returns:
         Configured APIGateway instance
 
     Example:
+        >>> from kailash.middleware.auth import JWTAuthManager
+        >>>
+        >>> # Create with custom auth
+        >>> auth = JWTAuthManager(use_rsa=True)
         >>> gateway = create_gateway(
         ...     title="My App Gateway",
-        ...     cors_origins=["http://localhost:3000"]
+        ...     cors_origins=["http://localhost:3000"],
+        ...     auth_manager=auth
         ... )
+        >>>
+        >>> # Or use default auth
+        >>> gateway = create_gateway(title="My App")
+        >>>
         >>> gateway.run(port=8000)
     """
+    # Pass auth_manager to APIGateway
+    if auth_manager is not None:
+        kwargs["auth_manager"] = auth_manager
+
     gateway = APIGateway(**kwargs)
 
     if agent_ui_middleware:

kailash/middleware/core/agent_ui.py

@@ -540,7 +540,114 @@ class AgentUIMiddleware:
         inputs: Dict[str, Any] = None,
         config_overrides: Dict[str, Any] = None,
     ) -> str:
-        """Execute a workflow asynchronously."""
+        """Execute a workflow asynchronously.
+
+        .. deprecated:: 0.5.0
+            Use :meth:`execute` instead. This method will be removed in version 1.0.0.
+        """
+        import warnings
+
+        warnings.warn(
+            "execute_workflow() is deprecated and will be removed in version 1.0.0. "
+            "Use execute() instead for consistency with runtime API.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        session = await self.get_session(session_id)
+        if not session:
+            raise ValueError(f"Session {session_id} not found")
+
+        # Get workflow
+        workflow = None
+        if workflow_id in session.workflows:
+            workflow = session.workflows[workflow_id]
+        elif workflow_id in self.shared_workflows:
+            workflow = self.shared_workflows[workflow_id]
+        else:
+            raise ValueError(f"Workflow {workflow_id} not found")
+
+        # Start execution
+        execution_id = session.start_execution(workflow_id, inputs)
+
+        # Track execution
+        self.active_executions[execution_id] = {
+            "session_id": session_id,
+            "workflow_id": workflow_id,
+            "workflow": workflow,
+            "inputs": inputs or {},
+            "config_overrides": config_overrides or {},
+            "start_time": time.time(),
+        }
+
+        # Persist execution if enabled
+        if self.enable_persistence:
+            try:
+                await self.execution_repo.create(
+                    {
+                        "id": execution_id,
+                        "workflow_id": workflow_id,
+                        "session_id": session_id,
+                        "user_id": session.user_id,
+                        "inputs": inputs,
+                        "metadata": config_overrides,
+                    }
+                )
+            except Exception as e:
+                logger.error(f"Failed to persist execution: {e}")
+
+        # Log execution start
+        logger.info(
+            f"Workflow execution started: {execution_id} for workflow {workflow_id}"
+        )
+
+        # Emit started event
+        await self.event_stream.emit_workflow_started(
+            workflow_id=workflow_id,
+            workflow_name=workflow.name,
+            execution_id=execution_id,
+            user_id=session.user_id,
+            session_id=session_id,
+        )
+
+        # Execute in background
+        asyncio.create_task(self._execute_workflow_async(execution_id))
+
+        self.workflows_executed += 1
+        return execution_id
+
+    async def execute(
+        self,
+        session_id: str,
+        workflow_id: str,
+        inputs: Dict[str, Any] = None,
+        config_overrides: Dict[str, Any] = None,
+    ) -> str:
+        """
+        Execute a workflow asynchronously.
+
+        This is the preferred method for workflow execution, providing consistency
+        with the runtime API.
+
+        Args:
+            session_id: Session identifier
+            workflow_id: Workflow identifier
+            inputs: Optional input parameters for the workflow
+            config_overrides: Optional configuration overrides
+
+        Returns:
+            str: Execution ID for tracking
+
+        Raises:
+            ValueError: If session or workflow not found
+            RuntimeError: If execution fails
+
+        Example:
+            >>> execution_id = await middleware.execute(
+            ...     session_id="sess_123",
+            ...     workflow_id="data_pipeline",
+            ...     inputs={"data": "input.csv"}
+            ... )
+        """
         session = await self.get_session(session_id)
         if not session:
             raise ValueError(f"Session {session_id} not found")

kailash/middleware/database/repositories.py

@@ -56,7 +56,9 @@ class BaseRepository:
         """Execute database query using SDK node."""
         try:
             if self.use_async:
-                result = await self.db_node.execute(query=query, params=params or {})
+                result = await self.db_node.execute_async(
+                    query=query, params=params or {}
+                )
             else:
                 result = self.db_node.execute(query=query, params=params or {})
 

kailash/middleware/mcp/enhanced_server.py

@@ -248,8 +248,8 @@ else:
 
         self.tool_register_workflow.add_node(validator)
         self.tool_register_workflow.add_node(register_handler)
-        self.tool_register_workflow.connect(
-            validator, register_handler, mapping={"result": "validation_result"}
+        self.tool_register_workflow.add_connection(
+            "validate_tool", "result", "register_tool", "validation_result"
         )
 
         # Tool Execution Workflow