security.py

import re
from typing import List

class SecurityValidator:
    @staticmethod
    def validate_input(text: str, max_length: int = 1000) -> bool:
        """Validate input text for security"""
        if len(text) > max_length:
            raise ValueError(f"Input too long. Max {max_length} characters.")
        
        # Check for potential prompt injection
        injection_patterns = [
            r"ignore previous",
            r"system prompt",
            r"forget your instructions"
        ]
        
        for pattern in injection_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                raise SecurityError("Potential prompt injection detected")
        
        return True
    
    @staticmethod
    def sanitize_output(text: str) -> str:
        """Sanitize model output"""
        # Remove any potentially harmful content
        harmful_patterns = [
            r"<script.*?>.*?</script>",
            r"<iframe.*?>.*?</iframe>"
        ]
        
        for pattern in harmful_patterns:
            text = re.sub(pattern, "", text, flags=re.IGNORECASE | re.DOTALL)
        
        return text

class SecurityError(Exception):
    pass