mcp-bastion-python 1.0.1__py3-none-any.whl

mcp_bastion/__init__.py

@@ -0,0 +1,14 @@
+"""
+MCP-Bastion: Security middleware for Model Context Protocol servers.
+"""
+
+from mcp_bastion.middleware import MCPBastionMiddleware
+from mcp_bastion.base import Middleware, MiddlewareContext, compose_middleware
+
+__all__ = [
+    "MCPBastionMiddleware",
+    "Middleware",
+    "MiddlewareContext",
+    "compose_middleware",
+]
+__version__ = "1.0.0"

mcp_bastion/base.py

@@ -0,0 +1,98 @@
+"""
+Base middleware abstractions for MCP-Bastion.
+
+Middleware base class, MiddlewareContext dataclass, compose_middleware.
+"""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from typing import Any, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass
+class MiddlewareContext(Generic[T]):
+    """Context for middleware chain: message, metadata, request_id, session_id."""
+
+    message: T
+    metadata: dict[str, Any] = field(default_factory=dict)
+    request_id: str | None = None
+    session_id: str | None = None
+
+    def copy(self, **kwargs: Any) -> "MiddlewareContext[T]":
+        """Create a copy with updated fields."""
+        data = {
+            "message": self.message,
+            "metadata": dict(self.metadata),
+            "request_id": self.request_id,
+            "session_id": self.session_id,
+        }
+        data.update(kwargs)
+        return MiddlewareContext(**data)
+
+
+CallNext = Callable[[MiddlewareContext[T]], Awaitable[Any]]
+
+
+class Middleware(Generic[T]):
+    """Base class for MCP middleware. Override on_message, on_call_tool, on_read_resource."""
+
+    async def __call__(
+        self,
+        context: MiddlewareContext[T],
+        call_next: CallNext[T],
+    ) -> Any:
+        return await self.on_message(context, call_next)
+
+    async def on_message(
+        self,
+        context: MiddlewareContext[T],
+        call_next: CallNext[T],
+    ) -> Any:
+        """Handle any message. Override for generic processing."""
+        return await call_next(context)
+
+    async def on_call_tool(
+        self,
+        context: MiddlewareContext[T],
+        call_next: CallNext[T],
+    ) -> Any:
+        """Handle tool calls. Override for tool-specific processing."""
+        return await self.on_message(context, call_next)
+
+    async def on_read_resource(
+        self,
+        context: MiddlewareContext[T],
+        call_next: CallNext[T],
+    ) -> Any:
+        """Handle resource reads. Override for resource-specific processing."""
+        return await self.on_message(context, call_next)
+
+
+def compose_middleware(
+    *middleware: Middleware[Any],
+) -> Callable[[MiddlewareContext[Any], CallNext[Any]], Awaitable[Any]]:
+    """Compose middleware. First in list = outermost."""
+    if not middleware:
+        async def passthrough(ctx: MiddlewareContext[Any], call_next: CallNext[Any]) -> Any:
+            return await call_next(ctx)
+        return passthrough
+
+    async def composed(
+        context: MiddlewareContext[Any],
+        call_next: CallNext[Any],
+    ) -> Any:
+        index = 0
+
+        async def next_handler(ctx: MiddlewareContext[Any]) -> Any:
+            nonlocal index
+            if index >= len(middleware):
+                return await call_next(ctx)
+            mw = middleware[index]
+            index += 1
+            return await mw(ctx, next_handler)
+
+        return await next_handler(context)
+
+    return composed

mcp_bastion/errors.py

@@ -0,0 +1,42 @@
+"""
+MCP-compliant error types for security policy violations.
+"""
+
+from __future__ import annotations
+
+
+class MCPBastionError(Exception):
+    """Base exception for MCP-Bastion security violations."""
+
+    def __init__(self, message: str, code: int = -32000) -> None:
+        super().__init__(message)
+        self.message = message
+        self.code = code
+
+    def to_mcp_error(self) -> dict:
+        """Format as MCP/JSON-RPC error object."""
+        return {
+            "code": self.code,
+            "message": self.message,
+        }
+
+
+class PromptInjectionError(MCPBastionError):
+    """Raised when prompt injection or jailbreak is detected."""
+
+    def __init__(self, message: str = "Request blocked: potential prompt injection detected") -> None:
+        super().__init__(message, code=-32001)
+
+
+class RateLimitExceededError(MCPBastionError):
+    """Raised when rate limit or iteration cap is exceeded."""
+
+    def __init__(self, message: str = "Request blocked: rate limit exceeded") -> None:
+        super().__init__(message, code=-32002)
+
+
+class TokenBudgetExceededError(MCPBastionError):
+    """Raised when FinOps token budget is exhausted."""
+
+    def __init__(self, message: str = "Request blocked: token budget exhausted") -> None:
+        super().__init__(message, code=-32003)

mcp_bastion/middleware.py

@@ -0,0 +1,229 @@
+"""
+MCP-Bastion security middleware.
+
+Intercepts CallToolRequest and ReadResourceResult for prompt injection,
+PII redaction, and rate limiting.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any
+
+from mcp_bastion.base import CallNext, Middleware, MiddlewareContext
+from mcp_bastion.errors import PromptInjectionError, RateLimitExceededError
+from mcp_bastion.pillars.pii_redaction import PIIRedactor
+from mcp_bastion.pillars.prompt_guard import PromptGuardEngine
+from mcp_bastion.pillars.rate_limit import TokenBucketRateLimiter
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_text_from_value(value: Any) -> str:
+    """Flatten args to string for injection check."""
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (int, float, bool)):
+        return str(value)
+    if isinstance(value, dict):
+        return " ".join(_extract_text_from_value(v) for v in value.values())
+    if isinstance(value, (list, tuple)):
+        return " ".join(_extract_text_from_value(v) for v in value)
+    return str(value)
+
+
+def _is_call_tool_request(message: Any) -> bool:
+    """True if message is tools/call."""
+    if hasattr(message, "root"):
+        msg = message.root
+    else:
+        msg = message
+    if hasattr(msg, "method") and getattr(msg, "method", None) == "tools/call":
+        return True
+    if isinstance(msg, dict) and msg.get("method") == "tools/call":
+        return True
+    return False
+
+
+def _is_read_resource_result(message: Any) -> bool:
+    """True if message has resource contents."""
+    if message is None:
+        return False
+    if hasattr(message, "contents"):
+        return True
+    if hasattr(message, "root"):
+        msg = message.root
+    else:
+        msg = message
+    if isinstance(msg, dict):
+        result = msg.get("result") or msg.get("params") or msg
+        if isinstance(result, dict) and ("contents" in result or "content" in result):
+            return True
+        if hasattr(result, "contents"):
+            return True
+    return False
+
+
+def _get_params(message: Any) -> dict | None:
+    """Extract params from message."""
+    if hasattr(message, "root"):
+        msg = message.root
+    else:
+        msg = message
+    if isinstance(msg, dict):
+        return msg.get("params") or msg.get("result")
+    if hasattr(msg, "params"):
+        return getattr(msg.params, "__dict__", None) or {}
+    return None
+
+
+def _get_request_id(message: Any) -> str | None:
+    """Extract request ID from message."""
+    if hasattr(message, "root"):
+        msg = message.root
+    else:
+        msg = message
+    if isinstance(msg, dict):
+        return str(msg.get("id", "")) or None
+    if hasattr(msg, "id"):
+        return str(getattr(msg, "id", "")) or None
+    return None
+
+
+def _get_content_from_result(result: Any) -> list[dict[str, Any]] | None:
+    """Extract content list from result for PII redaction."""
+    if result is None:
+        return None
+    payload = result
+    if isinstance(result, dict) and "result" in result:
+        payload = result["result"]
+    if hasattr(payload, "contents"):
+        items = payload.contents
+    elif isinstance(payload, dict) and "contents" in payload:
+        items = payload["contents"]
+    elif isinstance(payload, dict) and "content" in payload:
+        items = payload["content"]
+    else:
+        return None
+    if not isinstance(items, list):
+        return None
+    out = []
+    for item in items:
+        if hasattr(item, "model_dump"):
+            out.append(item.model_dump())
+        elif isinstance(item, dict):
+            out.append(dict(item))
+        else:
+            out.append({"type": "text", "text": str(item)})
+    return out
+
+
+def _set_content_in_result(result: Any, content: list[dict[str, Any]]) -> None:
+    """Replace content in result after redaction."""
+    payload = result
+    if isinstance(result, dict) and "result" in result:
+        payload = result["result"]
+    if hasattr(payload, "contents"):
+        payload.contents = content
+    elif isinstance(payload, dict):
+        if "contents" in payload:
+            payload["contents"] = content
+        if "content" in payload:
+            payload["content"] = content
+
+
+class MCPBastionMiddleware(Middleware[Any]):
+    def __init__(
+        self,
+        prompt_guard: PromptGuardEngine | None = None,
+        pii_redactor: PIIRedactor | None = None,
+        rate_limiter: TokenBucketRateLimiter | None = None,
+        enable_prompt_guard: bool = True,
+        enable_pii_redaction: bool = True,
+        enable_rate_limit: bool = True,
+    ) -> None:
+        self.prompt_guard = prompt_guard or PromptGuardEngine()
+        self.pii_redactor = pii_redactor or PIIRedactor()
+        self.rate_limiter = rate_limiter or TokenBucketRateLimiter()
+        self.enable_prompt_guard = enable_prompt_guard
+        self.enable_pii_redaction = enable_pii_redaction
+        self.enable_rate_limit = enable_rate_limit
+
+    async def __call__(
+        self,
+        context: MiddlewareContext[Any],
+        call_next: CallNext[Any],
+    ) -> Any:
+        """Run security checks, then call_next."""
+        start = time.perf_counter()
+        msg = context.message
+
+        try:
+            if _is_call_tool_request(msg):
+                return await self._handle_call_tool(context, call_next)
+            result = await call_next(context)
+            if result is not None and _is_read_resource_result(result):
+                result = self._redact_result_content(result)
+            return result
+        finally:
+            elapsed_ms = (time.perf_counter() - start) * 1000
+            context.metadata["elapsed_ms"] = round(elapsed_ms, 2)
+            logger.debug("request done elapsed_ms=%.2f", elapsed_ms)
+
+    async def _handle_call_tool(
+        self,
+        context: MiddlewareContext[Any],
+        call_next: CallNext[Any],
+    ) -> Any:
+        """Apply prompt guard and rate limit before tool execution."""
+        msg = context.message
+        params = _get_params(msg)
+        request_id = _get_request_id(msg) or context.request_id
+        session_id = context.session_id
+
+        if self.enable_rate_limit:
+            allowed, err = self.rate_limiter.check_iteration(
+                request_id=request_id,
+                session_id=session_id,
+            )
+            if not allowed:
+                logger.warning("rate_limit_blocked request_id=%s session_id=%s reason=%s", request_id, session_id, err)
+                raise RateLimitExceededError(err or "Rate limit exceeded")
+
+        if self.enable_prompt_guard and params:
+            arguments = params.get("arguments") or params
+            if isinstance(arguments, str):
+                try:
+                    arguments = json.loads(arguments)
+                except json.JSONDecodeError:
+                    arguments = {"raw": arguments}
+            text = _extract_text_from_value(arguments)
+            if text and self.prompt_guard.is_malicious(text):
+                logger.warning("prompt_injection_blocked request_id=%s", request_id)
+                raise PromptInjectionError()
+
+        self.rate_limiter.consume_iteration(
+            request_id=request_id,
+            session_id=session_id,
+        )
+
+        result = await call_next(context)
+
+        if self.enable_pii_redaction and result is not None:
+            result = self._redact_result_content(result)
+
+        return result
+
+    def _redact_result_content(self, result: Any) -> Any:
+        """Redact PII from result content items."""
+        content = _get_content_from_result(result)
+        if not content:
+            return result
+        redacted = self.pii_redactor.redact_content_items(content)
+        _set_content_in_result(result, redacted)
+        return result
+

mcp_bastion/pillars/__init__.py

@@ -0,0 +1,11 @@
+"""Security pillars for MCP-Bastion."""
+
+from mcp_bastion.pillars.prompt_guard import PromptGuardEngine
+from mcp_bastion.pillars.pii_redaction import PIIRedactor
+from mcp_bastion.pillars.rate_limit import TokenBucketRateLimiter
+
+__all__ = [
+    "PromptGuardEngine",
+    "PIIRedactor",
+    "TokenBucketRateLimiter",
+]

mcp_bastion/pillars/pii_redaction.py

@@ -0,0 +1,105 @@
+"""
+PII redaction via Microsoft Presidio.
+
+presidio-analyzer, presidio-anonymizer, spaCy. Sanitizes TextContent.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class PIIRedactor:
+    """Presidio + spaCy. Sanitizes TextContent."""
+
+    def __init__(
+        self,
+        entities: list[str] | None = None,
+        language: str = "en",
+    ) -> None:
+        self.entities = entities or [
+            "PERSON",
+            "EMAIL_ADDRESS",
+            "PHONE_NUMBER",
+            "CREDIT_CARD",
+            "US_SSN",
+            "US_PASSPORT",
+            "MEDICAL_LICENSE",
+            "IBAN_CODE",
+        ]
+        self.language = language
+        self._analyzer = None
+        self._anonymizer = None
+
+    def _ensure_loaded(self) -> None:
+        """Lazy-load Presidio components with optimized spaCy config."""
+        if self._analyzer is not None:
+            return
+        try:
+            from presidio_analyzer import AnalyzerEngine
+            from presidio_analyzer.nlp_engine import NlpEngineProvider
+            from presidio_anonymizer import AnonymizerEngine
+
+            config = {
+                "nlp_engine_name": "spacy",
+                "models": [{"lang_code": self.language, "model_name": "en_core_web_sm"}],
+            }
+            provider = NlpEngineProvider(nlp_configuration=config)
+            nlp_engine = provider.create_engine()
+
+            self._analyzer = AnalyzerEngine(nlp_engine=nlp_engine, supported_languages=[self.language])
+            self._anonymizer = AnonymizerEngine()
+        except Exception as e:
+            logger.warning("Presidio load failed: %s. PII redaction disabled.", e)
+            raise
+
+    def redact_text(self, text: str) -> str:
+        """
+        Analyze and anonymize PII in the given text.
+
+        Returns sanitized text with detected entities replaced by placeholders.
+        """
+        if not text or not isinstance(text, str):
+            return text
+
+        try:
+            self._ensure_loaded()
+            results = self._analyzer.analyze(
+                text=text,
+                language=self.language,
+                entities=self.entities,
+            )
+            if not results:
+                return text
+            logger.debug("redacted %d entities", len(results))
+            anonymized = self._anonymizer.anonymize(text=text, analyzer_results=results)
+            return anonymized.text
+        except Exception as e:
+            logger.warning("PII redaction failed: %s. Returning original text.", e)
+            return text
+
+    def redact_content_items(self, content: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """
+        Redact PII from MCP content items.
+
+        Processes TextContent items; other types are passed through unchanged.
+        """
+        if not content:
+            return content
+
+        result = []
+        for item in content:
+            if not isinstance(item, dict):
+                result.append(item)
+                continue
+            if item.get("type") == "text" and "text" in item:
+                result.append({
+                    **item,
+                    "text": self.redact_text(str(item["text"])),
+                })
+            else:
+                result.append(item)
+        return result

mcp_bastion/pillars/prompt_guard.py

@@ -0,0 +1,97 @@
+"""
+Prompt injection detection via Llama Prompt Guard 2.
+
+Uses meta-llama/Llama-Prompt-Guard-2-86M, temperature-adjusted softmax.
+Blocks when malicious probability exceeds threshold (default 0.85).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+
+logger = logging.getLogger(__name__)
+
+MALICIOUS_THRESHOLD = 0.85
+TEMPERATURE = 0.1
+MODEL_ID = "meta-llama/Llama-Prompt-Guard-2-86M"
+
+
+class PromptGuardEngine:
+    """PromptGuard 86M, temperature softmax. CPU/GPU."""
+
+    def __init__(
+        self,
+        threshold: float = MALICIOUS_THRESHOLD,
+        temperature: float = TEMPERATURE,
+        model_id: str = MODEL_ID,
+        device: str | None = None,
+    ) -> None:
+        self.threshold = threshold
+        self.temperature = temperature
+        self.model_id = model_id
+        self._model = None
+        self._tokenizer = None
+        self._device = device
+
+    def _ensure_loaded(self) -> None:
+        """Lazy-load model and tokenizer."""
+        if self._model is not None:
+            return
+        try:
+            import torch
+            from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+            self._model = AutoModelForSequenceClassification.from_pretrained(self.model_id)
+
+            if self._device is None:
+                self._device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._model = self._model.to(self._device)
+            self._model.eval()
+            logger.info("PromptGuard loaded model=%s device=%s", self.model_id, self._device)
+        except Exception as e:
+            logger.warning("PromptGuard model load failed: %s. Injection check disabled.", e)
+            raise
+
+    def _temperature_adjusted_softmax(self, logits: "torch.Tensor") -> "torch.Tensor":
+        """Temperature scaling before softmax."""
+        import torch
+        scaled = logits / self.temperature
+        return torch.softmax(scaled, dim=-1)
+
+    def score(self, text: str) -> float:
+        """Malicious probability 0-1. Above threshold = block."""
+        if not text or not text.strip():
+            return 0.0
+
+        self._ensure_loaded()
+        import torch
+
+        inputs = self._tokenizer(
+            text[:512],
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+        ).to(self._device)
+
+        with torch.no_grad():
+            outputs = self._model(**inputs)
+            probs = self._temperature_adjusted_softmax(outputs.logits)
+            probs_np = probs.cpu().numpy()
+
+        label2id = self._model.config.label2id
+        malicious_id = label2id.get("MALICIOUS", label2id.get("malicious", 1))
+        return float(probs_np[0][malicious_id])
+
+    def is_malicious(self, text: str) -> bool:
+        """True if score >= threshold."""
+        try:
+            score = self.score(text)
+            return score >= self.threshold
+        except Exception as e:
+            logger.warning("PromptGuard inference failed: %s. Allowing request.", e)
+            return False

mcp_bastion/pillars/rate_limit.py

@@ -0,0 +1,106 @@
+"""
+Rate limiting: token bucket per session.
+
+Max 15 iterations, 60s timeout, optional token budget (50k default).
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_ITERATIONS = 15
+DEFAULT_TIMEOUT_SECONDS = 60
+DEFAULT_TOKEN_BUDGET = 50_000
+
+
+@dataclass
+class SessionState:
+    """Per-session rate limit state."""
+
+    iterations: int = 0
+    started_at: float = field(default_factory=time.monotonic)
+    tokens_used: int = 0
+
+
+class TokenBucketRateLimiter:
+    """Token bucket per session. Iteration cap, timeout, token budget."""
+
+    def __init__(
+        self,
+        max_iterations: int = DEFAULT_MAX_ITERATIONS,
+        timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+        token_budget: int = DEFAULT_TOKEN_BUDGET,
+    ) -> None:
+        self.max_iterations = max_iterations
+        self.timeout_seconds = timeout_seconds
+        self.token_budget = token_budget
+        self._sessions: dict[str, SessionState] = defaultdict(SessionState)
+
+    def _get_session_id(self, request_id: str | None, session_id: str | None) -> str:
+        """Resolve session key from request or session ID."""
+        return session_id or request_id or "default"
+
+    def _cleanup_expired(self, session_key: str) -> None:
+        """Remove session if it has exceeded the global timeout."""
+        state = self._sessions.get(session_key)
+        if state is None:
+            return
+        elapsed = time.monotonic() - state.started_at
+        if elapsed > self.timeout_seconds:
+            del self._sessions[session_key]
+
+    def check_iteration(
+        self,
+        request_id: str | None = None,
+        session_id: str | None = None,
+    ) -> tuple[bool, str | None]:
+        """
+        Check if another iteration is allowed.
+
+        Returns (allowed, error_message). If allowed is False, error_message
+        describes the violation.
+        """
+        key = self._get_session_id(request_id, session_id)
+        self._cleanup_expired(key)
+
+        state = self._sessions[key]
+        elapsed = time.monotonic() - state.started_at
+
+        if elapsed > self.timeout_seconds:
+            del self._sessions[key]
+            return False, "Session timeout exceeded (60s limit)"
+
+        if state.iterations >= self.max_iterations:
+            return False, f"Maximum iterations exceeded ({self.max_iterations} limit)"
+
+        if state.tokens_used >= self.token_budget:
+            return False, f"Token budget exhausted ({self.token_budget} limit)"
+
+        return True, None
+
+    def consume_iteration(
+        self,
+        request_id: str | None = None,
+        session_id: str | None = None,
+        tokens: int = 0,
+    ) -> None:
+        """Record one iteration and optional token consumption."""
+        key = self._get_session_id(request_id, session_id)
+        state = self._sessions[key]
+        state.iterations += 1
+        state.tokens_used += tokens
+
+    def reset_session(
+        self,
+        request_id: str | None = None,
+        session_id: str | None = None,
+    ) -> None:
+        """Reset session state (e.g., on new request)."""
+        key = self._get_session_id(request_id, session_id)
+        if key in self._sessions:
+            del self._sessions[key]

mcp_bastion_python-1.0.1.dist-info/METADATA

@@ -0,0 +1,506 @@
+Metadata-Version: 2.4
+Name: mcp-bastion-python
+Version: 1.0.1
+Summary: Security middleware for MCP servers protecting LLM agents from prompt injection, resource exhaustion, and PII leakage
+Project-URL: Homepage, https://github.com/mcp-bastion/mcp-bastion
+Project-URL: Repository, https://github.com/mcp-bastion/mcp-bastion
+Project-URL: Documentation, https://github.com/mcp-bastion/mcp-bastion#readme
+Author: Viquar Khan
+License-Expression: MIT
+License-File: NOTICE
+Keywords: llm,mcp,middleware,pii,prompt-injection,security
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Security
+Requires-Python: >=3.10
+Requires-Dist: mcp>=1.0.0
+Requires-Dist: presidio-analyzer>=2.2.0
+Requires-Dist: presidio-anonymizer>=2.2.0
+Requires-Dist: spacy>=3.5.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: transformers>=4.30.0
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest>=7.0.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# MCP-Bastion
+
+**Enterprise-Grade Security Middleware for the Model Context Protocol**
+
+**Author:** Viquar Khan
+
+> Releases are published automatically to npm and PyPI via GitHub Actions when tags are pushed.
+
+The Model Context Protocol (MCP) has rapidly become the universally accepted standard for connecting AI agents to enterprise databases and APIs. However, this connectivity introduces a massive new attack surface: unpredictable, non-deterministic agentic behavior.
+
+MCP-Bastion is a lightweight, drop-in security middleware designed to wrap around any existing Python or TypeScript MCP server. Instead of relying on passive logging, human-in-the-loop approvals, or third-party APIs, MCP-Bastion provides an active, 100% local defense layer. It intercepts standard JSON-RPC traffic to stop threats before they cross the enterprise boundary.
+
+Under 5ms proxy overhead. MCP-Bastion provides:
+
+- **Prompt Injection Defense:** Meta PromptGuard runs locally to block adversarial payloads and jailbreaks.
+- **PII Redaction:** Uses Microsoft Presidio to detect and mask PII before it reaches the LLM context.
+- **Infinite Loop Protection:** Token buckets and cycle detection stop runaway agents from burning API budget.
+
+Secure your MCP server without changing business logic.
+
+---
+
+## Core Features
+
+**Zero-Click Prompt Injection Prevention**
+
+Integrates Meta's PromptGuard model locally to detect and block malicious payloads, jailbreaks, and adversarial tokenization before they reach your external tools.
+
+**PII Redaction**
+
+Microsoft Presidio scans outbound tool results and masks PII (redaction, substitution, generalization).
+
+**Infinite Loop and Denial of Wallet Protection**
+
+Implements stateful cycle detection and configurable FinOps token-bucket algorithms to automatically terminate runaway agents and prevent massive API bill overruns.
+
+**100% Local Execution (Data Privacy)**
+
+All security classification and data redaction happen entirely within the local memory space of your server. Sensitive data never leaves your enterprise network for third-party safety evaluations.
+
+**Low Latency**
+
+Drop-in middleware, under 5ms overhead.
+
+**Framework Integration**
+
+Hooks into MCP SDKs (TypeScript, Python) and FastMCP via standard middleware. No business logic changes.
+
+---
+
+## Why MCP-Bastion (Competitive Comparison)
+
+Early security packages (mcp-guardian, mcp-shield) focus on logging or static scanning. MCP-Bastion adds an active defense layer.
+
+### 1. Active Defense vs. Passive Logging
+
+| The Competition | MCP-Bastion |
+|-----------------|-------------|
+| Tools like mcp-guardian focus on tracing, logging, human-in-the-loop approvals. | Automated interception. MCP-Bastion scrubs PII before it leaves the server. |
+
+### 2. Local Inference vs. Third-Party APIs
+
+| The Competition | MCP-Bastion |
+|-----------------|-------------|
+| Many guardrail proxies send prompts to external APIs (e.g. OpenAI moderation) to check for malice. | PromptGuard-86M and Presidio run locally. Data stays on your network. |
+
+### 3. Stateful Denial of Wallet Protection
+
+| The Competition | MCP-Bastion |
+|-----------------|-------------|
+| Most tools focus on static vulns or basic rate limits. | Tracks tool call history per session. Stops runaway loops before they burn API budget. |
+
+### 4. Drop-in Middleware vs. Standalone Gateway
+
+| The Competition | MCP-Bastion |
+|-----------------|-------------|
+| Some solutions need standalone proxy servers. | Library hooks into `server.setRequestHandler` (TS) or middleware (Python). No extra infra. |
+
+---
+
+## Structure
+
+| Path | Description |
+|------|-------------|
+| `src/mcp_bastion/` | Python package: PromptGuard, Presidio, rate limiting |
+| `packages/core/` | TypeScript package: rate limiting; ML via Python sidecar |
+| `examples/` | Python examples: basic middleware, full demo ([examples/README.md](examples/README.md)) |
+| `scripts/validate_checklist.py` | Enterprise validation runner |
+| `VALIDATION_CHECKLIST.md` | Validation guide and MCP Inspector steps |
+| `SETUP_GUIDE.md` | Setup, config, and validation |
+
+## Installation
+
+**Python**
+
+```bash
+uv add mcp-bastion-python
+# or
+pip install mcp-bastion-python
+```
+
+**TypeScript**
+
+```bash
+npm install @mcp-bastion/core
+```
+
+## Developer Guide
+
+Integration examples for Python and TypeScript.
+
+---
+
+### Quick Start (Python)
+
+Add MCP-Bastion to an existing MCP server in three steps:
+
+```python
+from mcp_bastion import MCPBastionMiddleware, compose_middleware
+
+# 1. Create the security middleware
+bastion = MCPBastionMiddleware(
+    enable_prompt_guard=True,
+    enable_pii_redaction=True,
+    enable_rate_limit=True,
+)
+
+# 2. Compose with your middleware chain (Bastion runs first)
+middleware = compose_middleware(bastion)
+
+# 3. Pass the composed middleware to your MCP server
+# (integration depends on your server framework)
+```
+
+**Examples:**
+
+| Example | Description |
+|---------|-------------|
+| `examples/python_server_example.py` | Basic middleware chain |
+| `examples/full_demo.py` | All features: add, PII, rate limit, prompt injection |
+
+```bash
+# Windows: $env:PYTHONPATH="src"; python examples/full_demo.py
+# Linux/Mac: PYTHONPATH=src python examples/full_demo.py
+```
+
+**Enterprise validation:**
+
+```bash
+PYTHONPATH=src python scripts/validate_checklist.py
+```
+
+See `VALIDATION_CHECKLIST.md` and `SETUP_GUIDE.md`.
+
+---
+
+### Python Tutorial: FastMCP Server
+
+FastMCP server with MCP-Bastion.
+
+**Step 1: Install dependencies**
+
+```bash
+pip install mcp mcp-bastion-python
+```
+
+**Step 2: Create your server file** (`server.py`)
+
+```python
+from mcp.server.fastmcp import FastMCP
+from mcp_bastion import MCPBastionMiddleware, compose_middleware
+
+# Create the MCP server
+mcp = FastMCP("My Secure Server")
+
+# Create MCP-Bastion middleware
+# It intercepts tool calls and resource reads before they execute
+bastion = MCPBastionMiddleware(
+    enable_prompt_guard=True,   # Block malicious prompts via PromptGuard
+    enable_pii_redaction=True,  # Mask PII in outgoing content
+    enable_rate_limit=True,     # Cap at 15 iterations, 60s timeout
+)
+
+# Compose middleware chain (pass to your server's middleware config if supported)
+middleware = compose_middleware(bastion)
+
+# Register a tool (protected when middleware is wired into your server)
+@mcp.tool()
+def get_weather(city: str) -> str:
+    """Get weather for a city."""
+    return f"Weather in {city}: 22C, sunny"
+
+# Resource (PII redacted)
+@mcp.resource("user://profile/{user_id}")
+def get_profile(user_id: str) -> str:
+    """Get user profile. PII redacted."""
+    return f"User {user_id}: John Doe, SSN 123-45-6789, john@example.com"
+
+if __name__ == "__main__":
+    mcp.run(transport="streamable-http")
+```
+
+**Step 3: Run the server**
+
+```bash
+python server.py
+```
+
+MCP-Bastion:
+- Scans tool args for prompt injection
+- Redacts PII from resource responses
+- Blocks sessions over 15 calls or 60s
+
+---
+
+### Python: Custom Rate Limits
+
+Custom config example:
+
+```python
+from mcp_bastion import MCPBastionMiddleware
+from mcp_bastion.pillars.rate_limit import TokenBucketRateLimiter
+from mcp_bastion.pillars.prompt_guard import PromptGuardEngine
+
+# Stricter limits
+rate_limiter = TokenBucketRateLimiter(
+    max_iterations=10,
+    timeout_seconds=30,
+    token_budget=25_000,
+)
+
+# Higher threshold = fewer blocks, more risk
+prompt_guard = PromptGuardEngine(threshold=0.92)
+
+bastion = MCPBastionMiddleware(
+    prompt_guard=prompt_guard,
+    rate_limiter=rate_limiter,
+    enable_prompt_guard=True,
+    enable_pii_redaction=True,
+    enable_rate_limit=True,
+)
+
+# Disable PII redaction if your data has no PII
+bastion_no_pii = MCPBastionMiddleware(enable_pii_redaction=False)
+```
+
+---
+
+### Python: Custom Middleware
+
+Extend `Middleware` to add logging, metrics, or custom logic:
+
+```python
+from mcp_bastion.base import Middleware, MiddlewareContext, compose_middleware
+
+class LoggingMiddleware(Middleware):
+    async def on_message(self, context, call_next):
+        result = await call_next(context)
+        # log method, elapsed, etc.
+        return result
+
+middleware = compose_middleware(bastion, LoggingMiddleware())
+```
+
+See `examples/full_demo.py` for a complete example.
+
+---
+
+### TypeScript: Wrap an MCP Server
+
+**Step 1: Install dependencies**
+
+```bash
+npm install @modelcontextprotocol/sdk @mcp-bastion/core
+```
+
+**Step 2: Create your server** (`server.ts`)
+
+```typescript
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import {
+  wrapWithMcpBastion,
+  wrapCallToolHandler,
+} from "@mcp-bastion/core";
+
+const server = new Server({ name: "my-mcp-server", version: "1.0.0" });
+
+// Wrap the server with MCP-Bastion (rate limiting only by default)
+// For prompt injection and PII, run the Python sidecar and set sidecarUrl
+wrapWithMcpBastion(server, {
+  enableRateLimit: true,
+  maxIterations: 15,
+  timeoutMs: 60_000,
+  // Optional: enable ML features via Python sidecar
+  sidecarUrl: process.env.MCP_BASTION_SIDECAR || "",
+  enablePromptGuard: !!process.env.MCP_BASTION_SIDECAR,
+  enablePiiRedaction: !!process.env.MCP_BASTION_SIDECAR,
+});
+
+// Register tools (handlers are automatically wrapped)
+server.setRequestHandler("tools/call" as any, async (request) => {
+  if (request.params?.name === "get_weather") {
+    return {
+      content: [{ type: "text", text: "Sunny, 22C" }],
+      isError: false,
+    };
+  }
+  throw new Error("Unknown tool");
+});
+
+async function main() {
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+}
+
+main();
+```
+
+**Step 3: Run with rate limiting only**
+
+```bash
+npx tsx server.ts
+```
+
+**Step 4: Run with full ML features (Python sidecar)**
+
+For prompt injection and PII redaction, run a Python HTTP service that exposes `/prompt-guard` and `/pii-redact` endpoints (see the Python package for sidecar implementation). Then:
+
+```bash
+# Start the Python sidecar, then the TypeScript server
+MCP_BASTION_SIDECAR=http://localhost:8000 npx tsx server.ts
+```
+
+---
+
+### TypeScript: Wrap Individual Handlers
+
+Wrap specific handlers only:
+
+```typescript
+import {
+  wrapCallToolHandler,
+  wrapReadResourceHandler,
+} from "@mcp-bastion/core";
+import {
+  CallToolRequestSchema,
+  ReadResourceRequestSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+
+// Wrap only the tool handler
+const safeToolHandler = wrapCallToolHandler(
+  async (request) => {
+    // Your tool logic
+    return { content: [{ type: "text", text: "OK" }], isError: false };
+  },
+  { enableRateLimit: true, maxIterations: 10 }
+);
+
+// Wrap only the resource handler (for PII redaction)
+const safeResourceHandler = wrapReadResourceHandler(
+  async (request) => {
+    const contents = await fetchResource(request.params.uri);
+    return { contents };
+  },
+  { sidecarUrl: "http://localhost:8000", enablePiiRedaction: true }
+);
+
+server.setRequestHandler(CallToolRequestSchema, safeToolHandler);
+server.setRequestHandler(ReadResourceRequestSchema, safeResourceHandler);
+```
+
+---
+
+### Configuration Reference
+
+| Option | Python | TypeScript | Default | Description |
+|--------|--------|------------|---------|-------------|
+| `enable_prompt_guard` | Yes | Yes | `True` (Python) / `False` (TS) | Block malicious prompts via PromptGuard |
+| `enable_pii_redaction` | Yes | Yes | `True` (Python) / `False` (TS) | Mask PII in outgoing content |
+| `enable_rate_limit` | Yes | Yes | `True` | Enforce iteration and timeout caps |
+| `max_iterations` | Via `TokenBucketRateLimiter` | Yes | 15 | Max tool calls per session |
+| `timeout_seconds` / `timeoutMs` | Via `TokenBucketRateLimiter` | Yes | 60 | Session timeout |
+| `token_budget` | Via `TokenBucketRateLimiter` | - | 50,000 | FinOps token cap per request |
+| `sidecarUrl` | - | Yes | `""` | Python sidecar URL for ML features |
+| `threshold` | Via `PromptGuardEngine` | - | 0.85 | Malicious probability cutoff |
+| `setLogLevel` | - | Yes | `"info"` | TypeScript: `"debug"` \| `"info"` \| `"warn"` \| `"error"` |
+
+---
+
+### Error Handling
+
+When MCP-Bastion blocks a request, it returns standard MCP/JSON-RPC errors:
+
+| Code | Exception | When |
+|------|-----------|------|
+| -32001 | `PromptInjectionError` | Tool args contain jailbreak/injection |
+| -32002 | `RateLimitExceededError` | Session exceeds iteration or timeout limit |
+| -32003 | `TokenBudgetExceededError` | Session exceeds token budget |
+
+```python
+# Python: exceptions
+from mcp_bastion.errors import (
+    PromptInjectionError,
+    RateLimitExceededError,
+    TokenBudgetExceededError,
+)
+import logging
+logger = logging.getLogger(__name__)
+
+try:
+    result = await middleware(context, call_next)
+except PromptInjectionError as e:
+    logger.warning("blocked: %s", e.to_mcp_error())
+except RateLimitExceededError as e:
+    logger.warning("blocked: %s", e.to_mcp_error())
+except TokenBudgetExceededError as e:
+    logger.warning("blocked: %s", e.to_mcp_error())
+```
+
+```typescript
+// TypeScript: handlers return isError: true
+import { logger, setLogLevel } from "@mcp-bastion/core";
+setLogLevel("debug");  // optional: "debug" | "info" | "warn" | "error"
+const result = await guardedHandler(request);
+if (result.isError) {
+  logger.error("blocked", result.content);
+}
+```
+
+---
+
+### Testing
+
+MCP Inspector:
+
+```bash
+# Start your guarded server
+python server.py   # or: npx tsx server.ts
+
+# In another terminal, launch the Inspector
+npx -y @modelcontextprotocol/inspector
+```
+
+Connect via HTTP (`http://localhost:8000/mcp`) or stdio, then:
+1. List tools and call one with benign arguments (should succeed)
+2. Call a tool with "Ignore previous instructions" (should be blocked)
+3. Trigger 16+ tool calls in one session (should hit rate limit)
+
+---
+
+## Testing
+
+```bash
+# Python (PYTHONPATH=src on Windows: $env:PYTHONPATH="src")
+pytest tests/ -v
+
+# TypeScript
+npm run test --workspace=@mcp-bastion/core
+
+# Full validation checklist (build, pillars, latency)
+PYTHONPATH=src python scripts/validate_checklist.py
+
+# MCP Inspector (manual)
+npx -y @modelcontextprotocol/inspector
+```
+
+## Third-Party Components
+
+See `NOTICE` for licenses. MCP-Bastion uses Meta Llama Prompt Guard 2 (Llama 4 Community License) and Microsoft Presidio.
+
+## License
+
+MIT

mcp_bastion_python-1.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+mcp_bastion/__init__.py,sha256=izXzfw8A1AXqS-b22XPTtTSolsdMluCjPsUugd6CLoQ,347
+mcp_bastion/base.py,sha256=zZ0YG0tVn01cS7osjJDgBqK_6PVFme3GQ6tmhyfhITo,2910
+mcp_bastion/errors.py,sha256=oMGb3cjjSo3sOof8Hv9DdkucczhOAV2LY74xeEAkiXY,1295
+mcp_bastion/middleware.py,sha256=_FHFfINBP8Qmdl6KliEQbBIZIxlprO8-lkMoK0N0I_o,7799
+mcp_bastion/pillars/__init__.py,sha256=ybeqOYWVTGT3PEnquAz8nXSrofMV1F64H9eTC4P7wuU,317
+mcp_bastion/pillars/pii_redaction.py,sha256=U6jl33qwV53q7I14McRD7tvRadhpHz39gecMJBsps0c,3326
+mcp_bastion/pillars/prompt_guard.py,sha256=wR50TikDcnEjyzWz5Qut2QN523fW8ZIwUNfNlB14lqY,3148
+mcp_bastion/pillars/rate_limit.py,sha256=vp_r8TpNVx9nyVzgf96h6EW7O3YmX_tfkHyaRLlxj6M,3328
+mcp_bastion_python-1.0.1.dist-info/METADATA,sha256=Qp4fkViR2Vva4XVC1ai2JIidleFAjsGWTfeH1Ihf3Wg,15464
+mcp_bastion_python-1.0.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mcp_bastion_python-1.0.1.dist-info/licenses/NOTICE,sha256=_DlzQBhNBsf8mK-N55MCiz9juUyYjIaEPxZlMjFvPmc,273
+mcp_bastion_python-1.0.1.dist-info/RECORD,,

mcp_bastion_python-1.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mcp_bastion_python-1.0.1.dist-info/licenses/NOTICE

@@ -0,0 +1,5 @@
+MCP-Bastion uses the following third-party components:
+
+Llama Prompt Guard 2 (meta-llama/Llama-Prompt-Guard-2-86M)
+  Llama 4 is licensed under the Llama 4 Community License, Copyright (c) Meta Platforms, Inc. All Rights Reserved.
+  See: https://www.llama.com/docs/overview