raxe 0.4.6__py3-none-any.whl

raxe/integrations/utils.py

@@ -0,0 +1,135 @@
+"""
+Utility functions for safe optional imports.
+
+This module provides patterns for handling optional dependencies
+in a type-safe and user-friendly way.
+"""
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
+@contextmanager
+def optional_import(package_name: str, install_hint: str):
+    """
+    Context manager for optional imports with helpful error messages.
+
+    Usage:
+        with optional_import("langchain", "pip install raxe[langchain]"):
+            from langchain.callbacks import BaseCallbackHandler
+
+    Args:
+        package_name: Name of the package being imported
+        install_hint: Installation command to show on error
+
+    Yields:
+        None
+
+    Raises:
+        ImportError: With helpful message if import fails
+    """
+    try:
+        yield
+    except ImportError as e:
+        raise ImportError(
+            f"Could not import {package_name}. "
+            f"This feature requires additional dependencies.\n"
+            f"Install with: {install_hint}\n"
+            f"Original error: {e}"
+        ) from e
+
+
+def safe_import(
+    module_path: str,
+    attribute: str | None = None,
+    fallback: T | None = None,
+) -> T | Any:
+    """
+    Safely import a module or attribute, returning fallback on failure.
+
+    This is useful when you need a base class or type that may not
+    be available, and want to provide a fallback.
+
+    Args:
+        module_path: Full module path (e.g., "langchain.callbacks")
+        attribute: Optional attribute to get from module
+        fallback: Value to return if import fails
+
+    Returns:
+        The imported module/attribute or fallback value
+
+    Example:
+        # Use object as fallback base class when langchain not installed
+        BaseHandler = safe_import(
+            "langchain.callbacks.base",
+            "BaseCallbackHandler",
+            fallback=object
+        )
+
+        class MyHandler(BaseHandler):
+            ...
+    """
+    try:
+        import importlib
+
+        module = importlib.import_module(module_path)
+        if attribute:
+            return getattr(module, attribute)
+        return module
+    except (ImportError, AttributeError):
+        return fallback
+
+
+def get_version(package: str) -> str | None:
+    """
+    Get the installed version of a package.
+
+    Args:
+        package: Package name
+
+    Returns:
+        Version string or None if not installed
+    """
+    try:
+        from importlib.metadata import version
+
+        return version(package)
+    except Exception:
+        return None
+
+
+def check_version_compatibility(
+    package: str,
+    min_version: str | None = None,
+    max_version: str | None = None,
+) -> tuple[bool, str | None]:
+    """
+    Check if installed package version is compatible.
+
+    Args:
+        package: Package name
+        min_version: Minimum required version (inclusive)
+        max_version: Maximum allowed version (exclusive)
+
+    Returns:
+        Tuple of (is_compatible, installed_version)
+    """
+    from packaging.version import Version
+
+    installed = get_version(package)
+    if installed is None:
+        return False, None
+
+    try:
+        ver = Version(installed)
+        if min_version and ver < Version(min_version):
+            return False, installed
+        if max_version and ver >= Version(max_version):
+            return False, installed
+        return True, installed
+    except Exception:
+        # If version parsing fails, assume compatible
+        return True, installed

raxe/mcp/__init__.py

@@ -0,0 +1,62 @@
+"""
+RAXE MCP (Model Context Protocol) Server.
+
+This module provides an MCP server that exposes RAXE's threat detection
+capabilities to AI assistants like Claude Desktop and IDE extensions.
+
+Installation:
+    pip install raxe[mcp]
+
+Usage (CLI):
+    raxe-mcp                    # Start server with stdio transport
+    raxe-mcp --transport sse    # Start server with SSE transport
+    raxe-mcp --port 8765        # Specify port for HTTP transport
+
+Usage (Python):
+    from raxe.mcp import create_server
+    server = create_server()
+    server.run()
+
+Available Tools:
+    - scan_prompt: Scan a prompt for security threats
+    - get_scan_result: Get detailed result for a previous scan
+    - list_rules: List available detection rules
+    - get_rule_info: Get details about a specific rule
+
+Available Resources:
+    - raxe://rules/{family}: Detection rules by family
+    - raxe://stats: Current scan statistics
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from raxe.integrations.availability import require_integration
+
+if TYPE_CHECKING:
+    from raxe.mcp.server import RaxeMCPServer
+
+__all__ = ["create_server", "run_server", "RaxeMCPServer"]
+
+
+def __getattr__(name: str):
+    """Lazy import to avoid loading MCP dependencies until needed."""
+    if name == "create_server":
+        require_integration("mcp")
+        from raxe.mcp.server import create_server
+
+        return create_server
+
+    if name == "run_server":
+        require_integration("mcp")
+        from raxe.mcp.server import run_server
+
+        return run_server
+
+    if name == "RaxeMCPServer":
+        require_integration("mcp")
+        from raxe.mcp.server import RaxeMCPServer
+
+        return RaxeMCPServer
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

raxe/mcp/cli.py

@@ -0,0 +1,97 @@
+"""
+RAXE MCP Server CLI Entry Point.
+
+This module provides the raxe-mcp command for running the MCP server.
+Requires: pip install raxe[mcp]
+"""
+from __future__ import annotations
+
+import sys
+
+
+def main() -> int:
+    """
+    Main entry point for raxe-mcp command.
+
+    Returns:
+        Exit code (0 for success, non-zero for errors)
+    """
+    # Check for MCP availability before importing heavy modules
+    try:
+        import mcp  # noqa: F401
+    except ImportError:
+        print(
+            "ERROR: MCP dependencies not installed.\n"
+            "Install with: pip install raxe[mcp]",
+            file=sys.stderr,
+        )
+        return 1
+
+    # Parse command line arguments
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        prog="raxe-mcp",
+        description="RAXE MCP Server - AI Security for LLM Assistants",
+    )
+    parser.add_argument(
+        "--transport",
+        choices=["stdio", "sse"],
+        default="stdio",
+        help="Transport protocol (default: stdio)",
+    )
+    parser.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help="Host to bind for SSE transport (default: 127.0.0.1)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8765,
+        help="Port to bind for SSE transport (default: 8765)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose logging",
+    )
+    parser.add_argument(
+        "--version",
+        action="store_true",
+        help="Show version and exit",
+    )
+
+    args = parser.parse_args()
+
+    if args.version:
+        from raxe import __version__
+
+        print(f"raxe-mcp {__version__}")
+        return 0
+
+    # Import server after confirming dependencies are available
+    try:
+        from raxe.mcp.server import run_server
+
+        return run_server(
+            transport=args.transport,
+            host=args.host,
+            port=args.port,
+            verbose=args.verbose,
+        )
+    except KeyboardInterrupt:
+        print("\nServer stopped.")
+        return 0
+    except Exception as e:
+        print(f"ERROR: {e}", file=sys.stderr)
+        if args.verbose:
+            import traceback
+
+            traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

raxe/mcp/server.py

@@ -0,0 +1,409 @@
+"""
+RAXE MCP Server Implementation.
+
+This module implements the Model Context Protocol server for RAXE,
+exposing threat detection capabilities to AI assistants.
+
+The server is sync-first and uses the MCP SDK's synchronous APIs.
+"""
+from __future__ import annotations
+
+import sys
+import time
+from collections import defaultdict
+from typing import Any
+
+# Security limits
+MAX_TEXT_LENGTH = 100_000  # 100KB max input
+MAX_CONTEXT_LENGTH = 1_000  # 1KB max context
+RATE_LIMIT_RPM = 60  # 60 requests per minute per client
+
+# These imports will fail if MCP is not installed
+# The CLI entry point checks for this before importing this module
+from mcp.server import Server
+from mcp.server.stdio import stdio_server
+from mcp.types import (
+    CallToolResult,
+    GetPromptResult,
+    ListPromptsResult,
+    ListResourcesResult,
+    ListToolsResult,
+    Prompt,
+    PromptArgument,
+    PromptMessage,
+    Resource,
+    TextContent,
+    Tool,
+)
+from pydantic import AnyUrl
+
+from raxe import __version__
+from raxe.sdk.client import Raxe
+
+
+class RateLimiter:
+    """Simple in-memory rate limiter for MCP requests."""
+
+    def __init__(self, requests_per_minute: int = RATE_LIMIT_RPM) -> None:
+        self.requests_per_minute = requests_per_minute
+        self._request_times: dict[str, list[float]] = defaultdict(list)
+
+    def allow(self, client_id: str = "default") -> bool:
+        """Check if request is allowed under rate limit.
+
+        Args:
+            client_id: Client identifier for per-client limiting
+
+        Returns:
+            True if request is allowed, False if rate limited
+        """
+        now = time.time()
+        window_start = now - 60  # 1 minute window
+
+        # Clean old requests
+        self._request_times[client_id] = [
+            t for t in self._request_times[client_id] if t > window_start
+        ]
+
+        # Check limit
+        if len(self._request_times[client_id]) >= self.requests_per_minute:
+            return False
+
+        # Record this request
+        self._request_times[client_id].append(now)
+        return True
+
+
+class RaxeMCPServer:
+    """
+    MCP Server for RAXE threat detection.
+
+    Exposes RAXE scanning capabilities as MCP tools that can be
+    called by AI assistants like Claude Desktop.
+    """
+
+    def __init__(self, verbose: bool = False) -> None:
+        """
+        Initialize the MCP server.
+
+        Args:
+            verbose: Enable verbose logging
+        """
+        self.verbose = verbose
+        self.raxe = Raxe()
+        self.server = Server("raxe-security")
+        self.rate_limiter = RateLimiter()
+
+        # Register handlers
+        self._register_tools()
+        self._register_resources()
+        self._register_prompts()
+
+    def _register_tools(self) -> None:
+        """Register MCP tools for threat detection."""
+
+        @self.server.list_tools()
+        async def list_tools() -> ListToolsResult:
+            return ListToolsResult(
+                tools=[
+                    Tool(
+                        name="scan_prompt",
+                        description=(
+                            "Scan a prompt or text for security threats including "
+                            "prompt injection, jailbreak attempts, and data exfiltration. "
+                            "Returns threat detections with severity levels."
+                        ),
+                        inputSchema={
+                            "type": "object",
+                            "properties": {
+                                "text": {
+                                    "type": "string",
+                                    "description": "The text/prompt to scan for threats",
+                                },
+                                "context": {
+                                    "type": "string",
+                                    "description": "Optional context about where this text came from",
+                                },
+                            },
+                            "required": ["text"],
+                        },
+                    ),
+                    Tool(
+                        name="list_threat_families",
+                        description="List all available threat detection families (e.g., PI for Prompt Injection)",
+                        inputSchema={
+                            "type": "object",
+                            "properties": {},
+                        },
+                    ),
+                    Tool(
+                        name="get_rule_info",
+                        description="Get detailed information about a specific detection rule",
+                        inputSchema={
+                            "type": "object",
+                            "properties": {
+                                "rule_id": {
+                                    "type": "string",
+                                    "description": "The rule ID (e.g., 'pi-001', 'de-003')",
+                                },
+                            },
+                            "required": ["rule_id"],
+                        },
+                    ),
+                ]
+            )
+
+        @self.server.call_tool()
+        async def call_tool(name: str, arguments: dict[str, Any]) -> CallToolResult:
+            if name == "scan_prompt":
+                return await self._handle_scan(arguments)
+            elif name == "list_threat_families":
+                return await self._handle_list_families()
+            elif name == "get_rule_info":
+                return await self._handle_rule_info(arguments)
+            else:
+                return CallToolResult(
+                    content=[TextContent(type="text", text=f"Unknown tool: {name}")]
+                )
+
+    async def _handle_scan(self, arguments: dict[str, Any]) -> CallToolResult:
+        """Handle scan_prompt tool call."""
+        text = arguments.get("text", "")
+        context = arguments.get("context", "")
+
+        # Rate limiting
+        if not self.rate_limiter.allow():
+            return CallToolResult(
+                content=[TextContent(type="text", text="Error: Rate limit exceeded. Please try again later.")]
+            )
+
+        # Input validation
+        if not text:
+            return CallToolResult(
+                content=[TextContent(type="text", text="Error: 'text' is required")]
+            )
+
+        if len(text) > MAX_TEXT_LENGTH:
+            return CallToolResult(
+                content=[TextContent(
+                    type="text",
+                    text=f"Error: Input too large. Maximum {MAX_TEXT_LENGTH:,} characters allowed."
+                )]
+            )
+
+        if context and len(context) > MAX_CONTEXT_LENGTH:
+            return CallToolResult(
+                content=[TextContent(
+                    type="text",
+                    text=f"Error: Context too large. Maximum {MAX_CONTEXT_LENGTH:,} characters allowed."
+                )]
+            )
+
+        try:
+            result = self.raxe.scan(text)
+
+            if result.is_safe:
+                response = f"SAFE: No threats detected in the provided text."
+            else:
+                threats = []
+                for detection in result.detections:
+                    threats.append(
+                        f"- [{detection.severity}] {detection.rule_id}: {detection.description}"
+                    )
+                threat_list = "\n".join(threats)
+                response = f"THREATS DETECTED ({len(result.detections)}):\n{threat_list}"
+
+            if context:
+                response = f"Context: {context}\n\n{response}"
+
+            return CallToolResult(
+                content=[TextContent(type="text", text=response)]
+            )
+        except Exception as e:
+            return CallToolResult(
+                content=[TextContent(type="text", text=f"Scan error: {e}")]
+            )
+
+    async def _handle_list_families(self) -> CallToolResult:
+        """Handle list_threat_families tool call."""
+        families = {
+            "PI": "Prompt Injection - Attempts to override or manipulate AI instructions",
+            "JB": "Jailbreak - Attempts to bypass AI safety constraints",
+            "DE": "Data Exfiltration - Attempts to extract sensitive information",
+            "RP": "Role Play - Manipulation through persona adoption",
+            "CS": "Context Switching - Attempts to change conversation context",
+            "SE": "Social Engineering - Psychological manipulation techniques",
+            "TA": "Token Abuse - Exploitation of tokenization behaviors",
+            "CI": "Code Injection - Attempts to inject malicious code",
+        }
+
+        lines = ["Available Threat Families:\n"]
+        for code, desc in families.items():
+            lines.append(f"- {code}: {desc}")
+
+        return CallToolResult(
+            content=[TextContent(type="text", text="\n".join(lines))]
+        )
+
+    async def _handle_rule_info(self, arguments: dict[str, Any]) -> CallToolResult:
+        """Handle get_rule_info tool call."""
+        rule_id = arguments.get("rule_id", "")
+
+        if not rule_id:
+            return CallToolResult(
+                content=[TextContent(type="text", text="Error: 'rule_id' is required")]
+            )
+
+        # Get rule info from RAXE
+        try:
+            from raxe.infrastructure.rules.loader import RuleLoader
+
+            loader = RuleLoader()
+            rules = loader.load_rules()
+
+            for rule in rules:
+                if rule.rule_id.lower() == rule_id.lower():
+                    info = [
+                        f"Rule: {rule.rule_id}",
+                        f"Family: {rule.family}",
+                        f"Severity: {rule.severity}",
+                        f"Description: {rule.description}",
+                    ]
+                    return CallToolResult(
+                        content=[TextContent(type="text", text="\n".join(info))]
+                    )
+
+            return CallToolResult(
+                content=[TextContent(type="text", text=f"Rule not found: {rule_id}")]
+            )
+        except Exception as e:
+            return CallToolResult(
+                content=[TextContent(type="text", text=f"Error: {e}")]
+            )
+
+    def _register_resources(self) -> None:
+        """Register MCP resources for RAXE data."""
+
+        @self.server.list_resources()
+        async def list_resources() -> ListResourcesResult:
+            return ListResourcesResult(
+                resources=[
+                    Resource(
+                        uri=AnyUrl("raxe://version"),
+                        name="RAXE Version",
+                        description="Current RAXE version information",
+                        mimeType="text/plain",
+                    ),
+                    Resource(
+                        uri=AnyUrl("raxe://rules/summary"),
+                        name="Rules Summary",
+                        description="Summary of available detection rules",
+                        mimeType="text/plain",
+                    ),
+                ]
+            )
+
+    def _register_prompts(self) -> None:
+        """Register MCP prompts for common security tasks."""
+
+        @self.server.list_prompts()
+        async def list_prompts() -> ListPromptsResult:
+            return ListPromptsResult(
+                prompts=[
+                    Prompt(
+                        name="security_review",
+                        description="Review text for security concerns before processing",
+                        arguments=[
+                            PromptArgument(
+                                name="text",
+                                description="The text to review",
+                                required=True,
+                            ),
+                        ],
+                    ),
+                ]
+            )
+
+        @self.server.get_prompt()
+        async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult:
+            if name == "security_review":
+                text = (arguments or {}).get("text", "")
+                return GetPromptResult(
+                    description="Security review prompt",
+                    messages=[
+                        PromptMessage(
+                            role="user",
+                            content=TextContent(
+                                type="text",
+                                text=f"Please use the scan_prompt tool to check this text for security threats before processing:\n\n{text}",
+                            ),
+                        ),
+                    ],
+                )
+            return GetPromptResult(description="Unknown prompt", messages=[])
+
+    async def run_async(self) -> None:
+        """Run the server with stdio transport (async)."""
+        async with stdio_server() as (read_stream, write_stream):
+            await self.server.run(
+                read_stream,
+                write_stream,
+                self.server.create_initialization_options(),
+            )
+
+
+def create_server(verbose: bool = False) -> RaxeMCPServer:
+    """
+    Create a new RAXE MCP server instance.
+
+    Args:
+        verbose: Enable verbose logging
+
+    Returns:
+        Configured RaxeMCPServer instance
+    """
+    return RaxeMCPServer(verbose=verbose)
+
+
+def run_server(
+    transport: str = "stdio",
+    host: str = "127.0.0.1",
+    port: int = 8765,
+    verbose: bool = False,
+) -> int:
+    """
+    Run the RAXE MCP server.
+
+    Args:
+        transport: Transport protocol ("stdio" or "sse")
+        host: Host to bind for SSE transport
+        port: Port to bind for SSE transport
+        verbose: Enable verbose logging
+
+    Returns:
+        Exit code (0 for success)
+    """
+    import asyncio
+
+    server = create_server(verbose=verbose)
+
+    if transport == "stdio":
+        if verbose:
+            print(f"Starting RAXE MCP server v{__version__} (stdio transport)", file=sys.stderr)
+        asyncio.run(server.run_async())
+    elif transport == "sse":
+        # SSE transport requires additional setup
+        if verbose:
+            print(
+                f"Starting RAXE MCP server v{__version__} on {host}:{port} (SSE transport)",
+                file=sys.stderr,
+            )
+        # For SSE, we would use starlette/uvicorn
+        # This is a placeholder - full SSE implementation would go here
+        print("SSE transport not yet implemented. Use stdio transport.", file=sys.stderr)
+        return 1
+    else:
+        print(f"Unknown transport: {transport}", file=sys.stderr)
+        return 1
+
+    return 0