Guardrails

from abc import ABC, abstractmethod
from typing import Dict, Any, Optional, List

class GuardrailProvider(ABC):
    @abstractmethod
    def get_name(self) -> str:
        """Get provider name (e.g., 'enkrypt', 'openai')"""
        pass

    @abstractmethod
    def get_version(self) -> str:
        """Get provider version"""
        pass

    @abstractmethod
    def create_input_guardrail(
        self, config: Dict[str, Any]
    ) -> Optional[InputGuardrail]:
        """Create input guardrail instance"""
        pass

    @abstractmethod
    def create_output_guardrail(
        self, config: Dict[str, Any]
    ) -> Optional[OutputGuardrail]:
        """Create output guardrail instance"""
        pass

    def create_pii_handler(
        self, config: Dict[str, Any]
    ) -> Optional[PIIHandler]:
        """Create PII handler (optional)"""
        return None

class EnkryptInputGuardrail:
    def __init__(self, config: Dict[str, Any], api_key: str, base_url: str):
        self.api_key = api_key
        self.base_url = base_url
        self.policy_name = config.get("policy_name", "")
        self.block_list = config.get("block", [])
        self.guardrail_url = f"{base_url}/guardrails/policy/detect"

    async def validate(self, request: GuardrailRequest) -> GuardrailResponse:
        # Prepare API request
        payload = {"text": request.content}
        headers = {
            "X-Enkrypt-Policy": self.policy_name,
            "apikey": self.api_key,
            "Content-Type": "application/json",
            "X-Enkrypt-Source-Name": "mcp-gateway",
            "X-Enkrypt-Source-Event": "pre-tool"
        }

        # Call Enkrypt API
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.guardrail_url, json=payload, headers=headers
            ) as response:
                resp_json = await response.json()

        # Parse violations and return result
        violations = self._parse_violations(resp_json)
        return GuardrailResponse(
            is_safe=not violations,
            action=GuardrailAction.BLOCK if violations else GuardrailAction.ALLOW,
            violations=violations
        )

class OpenAIInputGuardrail:
    def __init__(self, config: Dict[str, Any]):
        self.api_key = config.get("api_key", "")
        self.threshold = config.get("threshold", 0.7)
        self.block_categories = config.get(
            "block_categories", ["hate", "violence", "sexual"]
        )

    async def validate(self, request: GuardrailRequest) -> GuardrailResponse:
        async with httpx.AsyncClient() as client:
            response = await client.post(
                "https://api.openai.com/v1/moderations",
                headers={"Authorization": f"Bearer {self.api_key}"},
                json={"input": request.content}
            )

            result = response.json()
            moderation_result = result["results"][0]

            violations = []
            for category, flagged in moderation_result["categories"].items():
                if flagged and category in self.block_categories:
                    score = moderation_result["category_scores"][category]
                    if score >= self.threshold:
                        violations.append(
                            GuardrailViolation(
                                violation_type=ViolationType.TOXIC_CONTENT,
                                severity=score,
                                message=f"Content flagged for {category}",
                                action=GuardrailAction.BLOCK,
                                metadata={"category": category, "score": score}
                            )
                        )

            return GuardrailResponse(
                is_safe=len(violations) == 0,
                action=GuardrailAction.ALLOW if not violations else GuardrailAction.BLOCK,
                violations=violations
            )

class CustomKeywordGuardrail:
    def __init__(self, config: Dict[str, Any]):
        self.blocked_keywords = config.get("blocked_keywords", [])
        self.case_sensitive = config.get("case_sensitive", False)

    async def validate(self, request: GuardrailRequest) -> GuardrailResponse:
        content = request.content
        if not self.case_sensitive:
            content = content.lower()
            blocked_keywords = [kw.lower() for kw in self.blocked_keywords]
        else:
            blocked_keywords = self.blocked_keywords

        violations = []
        for keyword in blocked_keywords:
            if keyword in content:
                violations.append(
                    GuardrailViolation(
                        violation_type=ViolationType.KEYWORD_VIOLATION,
                        severity=0.8,
                        message=f"Blocked keyword detected: {keyword}",
                        action=GuardrailAction.BLOCK,
                        metadata={"keyword": keyword}
                    )
                )

        return GuardrailResponse(
            is_safe=len(violations) == 0,
            action=GuardrailAction.ALLOW if not violations else GuardrailAction.BLOCK,
            violations=violations
        )

{
  "plugins": {
    "guardrails": {
      "provider": "enkrypt",
      "config": {
        "api_key": "YOUR_ENKRYPT_API_KEY",
        "base_url": "https://api.enkryptai.com"
      }
    }
  }
}

# Original request
"Contact John at [email protected] or call 555-1234"

# Redacted (sent to MCP server)
"Contact John at [EMAIL_1] or call [PHONE_1]"

# Mapping stored
{
  "[EMAIL_1]": "[email protected]",
  "[PHONE_1]": "555-1234"
}

# Response received
"I've sent an email to [EMAIL_1]"

# De-anonymized (returned to client)
"I've sent an email to [email protected]"

from secure_mcp_gateway.plugins.guardrails.base import (
    GuardrailProvider,
    InputGuardrail,
    GuardrailRequest,
    GuardrailResponse
)

class MyInputGuardrail:
    async def validate(self, request: GuardrailRequest) -> GuardrailResponse:
        # Implement your custom validation logic
        if "forbidden" in request.content:
            return GuardrailResponse(
                is_safe=False,
                action=GuardrailAction.BLOCK,
                violations=[...]
            )
        return GuardrailResponse(is_safe=True, action=GuardrailAction.ALLOW, violations=[])

class MyGuardrailProvider(GuardrailProvider):
    def get_name(self) -> str:
        return "my_custom"

    def create_input_guardrail(self, config):
        return MyInputGuardrail()