kailash 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

kailash/middleware/auth/kailash_jwt_auth.py

@@ -0,0 +1,616 @@
+"""
+Kailash SDK-based JWT Authentication Manager
+
+Built entirely with Kailash SDK components - uses nodes and workflows
+for all authentication operations. This demonstrates enterprise-grade
+authentication using only Kailash patterns.
+"""
+
+import json
+import logging
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import Any, Dict, List, Optional
+
+# Import Kailash SDK components only
+from kailash.nodes.base import Node, NodeParameter, register_node
+from kailash.nodes.code import PythonCodeNode
+from kailash.nodes.data import JSONReaderNode
+from kailash.nodes.logic import MergeNode, SwitchNode
+from kailash.nodes.security import CredentialManagerNode
+from kailash.runtime.local import LocalRuntime
+from kailash.workflow.builder import WorkflowBuilder
+
+logger = logging.getLogger(__name__)
+
+
+@register_node()
+class JWTConfigNode(Node):
+    """Kailash node for JWT configuration management."""
+
+    def __init__(self, name: str = "jwt_config"):
+        super().__init__(name=name)
+        self.algorithm = "HS256"  # Use symmetric for simplicity with Kailash
+        self.access_token_expire_minutes = 15
+        self.refresh_token_expire_days = 7
+        self.issuer = "kailash-middleware"
+        self.audience = "kailash-api"
+
+    def get_parameters(self) -> Dict[str, NodeParameter]:
+        return {
+            "algorithm": NodeParameter(
+                name="algorithm",
+                type=str,
+                required=False,
+                default="HS256",
+                description="JWT signing algorithm",
+            ),
+            "access_expire_minutes": NodeParameter(
+                name="access_expire_minutes",
+                type=int,
+                required=False,
+                default=15,
+                description="Access token expiration in minutes",
+            ),
+            "refresh_expire_days": NodeParameter(
+                name="refresh_expire_days",
+                type=int,
+                required=False,
+                default=7,
+                description="Refresh token expiration in days",
+            ),
+        }
+
+    def run(self, **kwargs) -> Dict[str, Any]:
+        """Return JWT configuration."""
+        return {
+            "config": {
+                "algorithm": getattr(self, "algorithm", "HS256"),
+                "access_token_expire_minutes": getattr(
+                    self, "access_token_expire_minutes", 15
+                ),
+                "refresh_token_expire_days": getattr(
+                    self, "refresh_token_expire_days", 7
+                ),
+                "issuer": getattr(self, "issuer", "kailash-middleware"),
+                "audience": getattr(self, "audience", "kailash-api"),
+            }
+        }
+
+    async def process(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        """Async wrapper for middleware compatibility."""
+        return self.run(**inputs)
+
+
+class TokenGeneratorNode(PythonCodeNode):
+    """Kailash node for generating JWT tokens."""
+
+    def __init__(self, name: str = "token_generator"):
+        super().__init__(
+            name=name,
+            code="""
+import jwt
+import uuid
+from datetime import datetime, timedelta, timezone
+
+def generate_token(user_id, token_type='access', config=None, **claims):
+    '''Generate JWT token using Kailash patterns'''
+    config = config or {}
+
+    now = datetime.now(timezone.utc)
+
+    # Set expiration based on token type
+    if token_type == 'access':
+        expire_minutes = config.get('access_token_expire_minutes', 15)
+        exp = now + timedelta(minutes=expire_minutes)
+    else:  # refresh
+        expire_days = config.get('refresh_token_expire_days', 7)
+        exp = now + timedelta(days=expire_days)
+
+    # Create payload with Kailash conventions
+    payload = {
+        'sub': user_id,
+        'iss': config.get('issuer', 'kailash-middleware'),
+        'aud': config.get('audience', 'kailash-api'),
+        'exp': int(exp.timestamp()),
+        'iat': int(now.timestamp()),
+        'jti': str(uuid.uuid4()),
+        'token_type': token_type,
+        **claims
+    }
+
+    # Use simple secret for Kailash demo (in production, use proper key management)
+    secret = config.get('secret', 'kailash-jwt-secret-key')
+    algorithm = config.get('algorithm', 'HS256')
+
+    token = jwt.encode(payload, secret, algorithm=algorithm)
+
+    return {
+        'token': token,
+        'payload': payload,
+        'expires_at': exp.isoformat(),
+        'token_type': token_type
+    }
+
+# Main execution
+user_id = input_data.get('user_id')
+token_type = input_data.get('token_type', 'access')
+config = input_data.get('config', {})
+claims = input_data.get('claims', {})
+
+if not user_id:
+    result = {'error': 'user_id is required'}
+else:
+    result = generate_token(user_id, token_type, config, **claims)
+""",
+        )
+
+
+class TokenVerifierNode(PythonCodeNode):
+    """Kailash node for verifying JWT tokens."""
+
+    def __init__(self, name: str = "token_verifier"):
+        super().__init__(
+            name=name,
+            code="""
+import jwt
+from datetime import datetime, timezone
+
+def verify_token(token, config=None, blacklisted_tokens=None):
+    '''Verify JWT token using Kailash patterns'''
+    config = config or {}
+    blacklisted_tokens = blacklisted_tokens or set()
+
+    try:
+        # Check if token is blacklisted
+        if token in blacklisted_tokens:
+            return {
+                'valid': False,
+                'error': 'Token has been revoked',
+                'error_type': 'revoked'
+            }
+
+        # Verify token
+        secret = config.get('secret', 'kailash-jwt-secret-key')
+        algorithm = config.get('algorithm', 'HS256')
+
+        payload = jwt.decode(
+            token,
+            secret,
+            algorithms=[algorithm],
+            issuer=config.get('issuer'),
+            audience=config.get('audience')
+        )
+
+        return {
+            'valid': True,
+            'payload': payload,
+            'user_id': payload.get('sub'),
+            'token_type': payload.get('token_type', 'access'),
+            'expires_at': payload.get('exp'),
+            'permissions': payload.get('permissions', []),
+            'roles': payload.get('roles', []),
+            'tenant_id': payload.get('tenant_id'),
+            'session_id': payload.get('session_id')
+        }
+
+    except jwt.ExpiredSignatureError:
+        return {
+            'valid': False,
+            'error': 'Token has expired',
+            'error_type': 'expired'
+        }
+    except jwt.InvalidTokenError as e:
+        return {
+            'valid': False,
+            'error': str(e),
+            'error_type': 'invalid'
+        }
+    except Exception as e:
+        return {
+            'valid': False,
+            'error': f'Token verification failed: {str(e)}',
+            'error_type': 'error'
+        }
+
+# Main execution
+token = input_data.get('token')
+config = input_data.get('config', {})
+blacklisted_tokens = input_data.get('blacklisted_tokens', set())
+
+if not token:
+    result = {'valid': False, 'error': 'Token is required', 'error_type': 'missing'}
+else:
+    result = verify_token(token, config, blacklisted_tokens)
+""",
+        )
+
+
+class RefreshTokenNode(PythonCodeNode):
+    """Kailash node for handling token refresh."""
+
+    def __init__(self, name: str = "refresh_token"):
+        super().__init__(
+            name=name,
+            code="""
+def refresh_access_token(refresh_token_data, config, refresh_tracking):
+    '''Refresh access token using Kailash patterns'''
+
+    if not refresh_token_data.get('valid'):
+        return {
+            'success': False,
+            'error': 'Invalid refresh token',
+            'error_type': 'invalid_refresh'
+        }
+
+    payload = refresh_token_data.get('payload', {})
+
+    # Check token type
+    if payload.get('token_type') != 'refresh':
+        return {
+            'success': False,
+            'error': 'Token is not a refresh token',
+            'error_type': 'wrong_type'
+        }
+
+    jti = payload.get('jti')
+    user_id = payload.get('sub')
+
+    # Check refresh count (Kailash pattern for security)
+    refresh_info = refresh_tracking.get(jti, {})
+    max_refresh_count = config.get('max_refresh_count', 10)
+
+    if refresh_info.get('refresh_count', 0) >= max_refresh_count:
+        return {
+            'success': False,
+            'error': 'Refresh token has exceeded usage limit',
+            'error_type': 'limit_exceeded'
+        }
+
+    # Update refresh tracking
+    refresh_info['refresh_count'] = refresh_info.get('refresh_count', 0) + 1
+    refresh_info['last_used'] = datetime.now().isoformat()
+    refresh_tracking[jti] = refresh_info
+
+    # Prepare new token creation data
+    return {
+        'success': True,
+        'user_id': user_id,
+        'tenant_id': payload.get('tenant_id'),
+        'session_id': payload.get('session_id'),
+        'permissions': payload.get('permissions', []),
+        'roles': payload.get('roles', []),
+        'refresh_tracking': refresh_tracking
+    }
+
+# Main execution
+refresh_token_data = input_data.get('refresh_token_data', {})
+config = input_data.get('config', {})
+refresh_tracking = input_data.get('refresh_tracking', {})
+
+result = refresh_access_token(refresh_token_data, config, refresh_tracking)
+""",
+        )
+
+
+class KailashJWTAuthManager:
+    """
+    JWT Authentication Manager built entirely with Kailash SDK components.
+
+    This demonstrates enterprise-grade authentication using only Kailash
+    nodes and workflows - no external dependencies beyond JWT library.
+    """
+
+    def __init__(self, secret_key: str = "kailash-jwt-secret-key"):
+        self.secret_key = secret_key
+        self.runtime = LocalRuntime()
+
+        # State management using Kailash patterns
+        self.blacklisted_tokens = set()
+        self.refresh_tracking = {}
+
+        # Create authentication workflows
+        self._create_workflows()
+
+    def _get_token_generator_code(self):
+        """Get token generator code for PythonCodeNode."""
+        return """
+import jwt
+import uuid
+from datetime import datetime, timedelta, timezone
+
+def generate_token(user_id, token_type='access', config=None, **claims):
+    '''Generate JWT token using Kailash patterns'''
+    config = config or {}
+
+    now = datetime.now(timezone.utc)
+
+    # Set expiration based on token type
+    if token_type == 'access':
+        expire_minutes = config.get('access_token_expire_minutes', 15)
+        exp = now + timedelta(minutes=expire_minutes)
+    else:  # refresh
+        expire_days = config.get('refresh_token_expire_days', 7)
+        exp = now + timedelta(days=expire_days)
+
+    # Create payload with Kailash conventions
+    payload = {
+        'sub': user_id,
+        'iss': config.get('issuer', 'kailash-middleware'),
+        'aud': config.get('audience', 'kailash-api'),
+        'exp': int(exp.timestamp()),
+        'iat': int(now.timestamp()),
+        'jti': str(uuid.uuid4()),
+        'token_type': token_type,
+        **claims
+    }
+
+    # Use simple secret for Kailash demo (in production, use proper key management)
+    secret = config.get('secret', 'kailash-jwt-secret-key')
+    algorithm = config.get('algorithm', 'HS256')
+
+    token = jwt.encode(payload, secret, algorithm=algorithm)
+
+    return {
+        'token': token,
+        'payload': payload,
+        'expires_at': exp.isoformat(),
+        'token_type': token_type
+    }
+
+# Main execution
+user_id = input_data.get('user_id')
+token_type = input_data.get('token_type', 'access')
+config = input_data.get('config', {})
+claims = input_data.get('claims', {})
+
+if not user_id:
+    result = {'error': 'user_id is required'}
+else:
+    result = generate_token(user_id, token_type, config, **claims)
+"""
+
+    def _get_token_verifier_code(self):
+        """Get token verifier code for PythonCodeNode."""
+        return """
+import jwt
+from datetime import datetime, timezone
+
+def verify_token(token, config=None, blacklisted_tokens=None):
+    config = config or {}
+    blacklisted_tokens = blacklisted_tokens or set()
+
+    try:
+        if token in blacklisted_tokens:
+            return {'valid': False, 'error': 'Token has been revoked', 'error_type': 'revoked'}
+
+        secret = config.get('secret', 'kailash-jwt-secret-key')
+        algorithm = config.get('algorithm', 'HS256')
+
+        payload = jwt.decode(token, secret, algorithms=[algorithm])
+
+        return {
+            'valid': True,
+            'payload': payload,
+            'user_id': payload.get('sub'),
+            'token_type': payload.get('token_type', 'access')
+        }
+
+    except jwt.ExpiredSignatureError:
+        return {'valid': False, 'error': 'Token has expired', 'error_type': 'expired'}
+    except jwt.InvalidTokenError as e:
+        return {'valid': False, 'error': str(e), 'error_type': 'invalid'}
+    except Exception as e:
+        return {'valid': False, 'error': f'Token verification failed: {str(e)}', 'error_type': 'error'}
+
+token = input_data.get('token')
+config = input_data.get('config', {})
+blacklisted_tokens = input_data.get('blacklisted_tokens', set())
+
+if not token:
+    result = {'valid': False, 'error': 'Token is required', 'error_type': 'missing'}
+else:
+    result = verify_token(token, config, blacklisted_tokens)
+"""
+
+    def _create_workflows(self):
+        """Create Kailash workflows for authentication operations."""
+
+        # Token Creation Workflow
+        self.token_workflow = WorkflowBuilder()
+
+        config_id = self.token_workflow.add_node(
+            "JWTConfigNode", "jwt_config", {"name": "jwt_config"}
+        )
+        token_id = self.token_workflow.add_node(
+            "PythonCodeNode",
+            "generate_token",
+            {"name": "generate_token", "code": self._get_token_generator_code()},
+        )
+
+        self.token_workflow.add_connection(config_id, "config", token_id, "config")
+
+        # Token Verification Workflow
+        self.verify_workflow = WorkflowBuilder()
+
+        config_verify_id = self.verify_workflow.add_node(
+            "JWTConfigNode", "jwt_config_verify", {"name": "jwt_config_verify"}
+        )
+        verify_id = self.verify_workflow.add_node(
+            "PythonCodeNode",
+            "verify_token",
+            {"name": "verify_token", "code": self._get_token_verifier_code()},
+        )
+
+        self.verify_workflow.add_connection(
+            config_verify_id, "config", verify_id, "config"
+        )
+
+        # For simplicity, we'll use direct verification in methods instead of complex workflows
+
+    def create_access_token(
+        self,
+        user_id: str,
+        tenant_id: str = None,
+        session_id: str = None,
+        permissions: List[str] = None,
+        roles: List[str] = None,
+    ) -> str:
+        """Create access token using Kailash workflow."""
+
+        claims = {}
+        if tenant_id:
+            claims["tenant_id"] = tenant_id
+        if session_id:
+            claims["session_id"] = session_id
+        if permissions:
+            claims["permissions"] = permissions
+        if roles:
+            claims["roles"] = roles
+
+        inputs = {
+            "user_id": user_id,
+            "token_type": "access",
+            "claims": claims,
+            "config": {"secret": self.secret_key},
+        }
+
+        # Execute Kailash workflow
+        workflow = self.token_workflow.build()
+        results, _ = self.runtime.execute(workflow, parameters=inputs)
+
+        return results.get("generate_token", {}).get("token")
+
+    def create_refresh_token(
+        self, user_id: str, tenant_id: str = None, session_id: str = None
+    ) -> str:
+        """Create refresh token using Kailash workflow."""
+
+        claims = {}
+        if tenant_id:
+            claims["tenant_id"] = tenant_id
+        if session_id:
+            claims["session_id"] = session_id
+
+        inputs = {
+            "user_id": user_id,
+            "token_type": "refresh",
+            "claims": claims,
+            "config": {"secret": self.secret_key},
+        }
+
+        # Execute Kailash workflow
+        workflow = self.token_workflow.build()
+        results, _ = self.runtime.execute(workflow, parameters=inputs)
+
+        token_result = results.get("generate_token", {})
+        token = token_result.get("token")
+
+        # Track refresh token in Kailash pattern
+        if token:
+            payload = token_result.get("payload", {})
+            jti = payload.get("jti")
+            if jti:
+                self.refresh_tracking[jti] = {
+                    "user_id": user_id,
+                    "tenant_id": tenant_id,
+                    "session_id": session_id,
+                    "created_at": datetime.now(timezone.utc).isoformat(),
+                    "refresh_count": 0,
+                }
+
+        return token
+
+    def create_token_pair(
+        self,
+        user_id: str,
+        tenant_id: str = None,
+        session_id: str = None,
+        permissions: List[str] = None,
+        roles: List[str] = None,
+    ) -> Dict[str, Any]:
+        """Create access and refresh token pair using Kailash workflows."""
+
+        access_token = self.create_access_token(
+            user_id, tenant_id, session_id, permissions, roles
+        )
+        refresh_token = self.create_refresh_token(user_id, tenant_id, session_id)
+
+        return {
+            "access_token": access_token,
+            "refresh_token": refresh_token,
+            "token_type": "Bearer",
+            "expires_in": 15 * 60,  # 15 minutes in seconds
+            "scope": "kailash-api",
+        }
+
+    def verify_token(self, token: str) -> Dict[str, Any]:
+        """Verify token using Kailash workflow."""
+
+        inputs = {
+            "token": token,
+            "config": {"secret": self.secret_key},
+            "blacklisted_tokens": self.blacklisted_tokens,
+        }
+
+        # Execute Kailash verification workflow
+        workflow = self.verify_workflow.build()
+        results, _ = self.runtime.execute(workflow, parameters=inputs)
+
+        return results.get("verify_token", {})
+
+    def refresh_access_token(self, refresh_token: str) -> Dict[str, Any]:
+        """Refresh access token using simplified Kailash pattern."""
+
+        # Verify the refresh token
+        verify_result = self.verify_token(refresh_token)
+
+        if not verify_result.get("valid"):
+            return {"success": False, "error": "Invalid refresh token"}
+
+        payload = verify_result.get("payload", {})
+
+        if payload.get("token_type") != "refresh":
+            return {"success": False, "error": "Token is not a refresh token"}
+
+        # Create new token pair
+        return self.create_token_pair(
+            user_id=payload.get("sub"),
+            tenant_id=payload.get("tenant_id"),
+            session_id=payload.get("session_id"),
+            permissions=payload.get("permissions", []),
+            roles=payload.get("roles", []),
+        )
+
+    def revoke_token(self, token: str):
+        """Revoke token by adding to blacklist (Kailash pattern)."""
+        self.blacklisted_tokens.add(token)
+        logger.info(f"Revoked token (length: {len(token)})")
+
+    def revoke_all_user_tokens(self, user_id: str):
+        """Revoke all tokens for a user (Kailash pattern)."""
+        # Remove all refresh tokens for user
+        to_remove = []
+        for jti, data in self.refresh_tracking.items():
+            if data.get("user_id") == user_id:
+                to_remove.append(jti)
+
+        for jti in to_remove:
+            del self.refresh_tracking[jti]
+
+        logger.info(f"Revoked all tokens for user {user_id}")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get authentication statistics (Kailash pattern)."""
+        return {
+            "active_refresh_tokens": len(self.refresh_tracking),
+            "blacklisted_tokens": len(self.blacklisted_tokens),
+            "auth_manager": "KailashJWTAuthManager",
+            "workflow_runtime": "LocalRuntime",
+            "nodes_used": [
+                "JWTConfigNode",
+                "TokenGeneratorNode",
+                "TokenVerifierNode",
+                "RefreshTokenNode",
+            ],
+        }

kailash/middleware/communication/__init__.py

@@ -0,0 +1,37 @@
+"""
+Communication Layer for Kailash Middleware
+
+Handles all external communication including REST APIs, WebSockets,
+events, and AI chat interfaces.
+"""
+
+from .ai_chat import AIChatMiddleware, ChatMessage, WorkflowGenerator
+from .api_gateway import APIGateway, create_gateway
+from .events import (
+    EventPriority,
+    EventStream,
+    EventType,
+    NodeEvent,
+    UIEvent,
+    WorkflowEvent,
+)
+from .realtime import RealtimeMiddleware
+
+__all__ = [
+    # Event system
+    "EventStream",
+    "EventType",
+    "EventPriority",
+    "WorkflowEvent",
+    "NodeEvent",
+    "UIEvent",
+    # Real-time communication
+    "RealtimeMiddleware",
+    # API Gateway
+    "APIGateway",
+    "create_gateway",
+    # AI Chat
+    "AIChatMiddleware",
+    "ChatMessage",
+    "WorkflowGenerator",
+]