dtSpark 1.0.4__py3-none-any.whl

dtSpark/safety/llm_service.py

@@ -0,0 +1,111 @@
+"""
+LLM service wrapper for prompt inspection.
+
+Provides unified interface to multiple LLM providers (AWS Bedrock, Ollama, Anthropic Direct)
+specifically for prompt inspection and security analysis tasks.
+
+
+"""
+
+import logging
+from typing import Optional, Dict, List, Any
+
+
+class InspectionLLMService:
+    """
+    LLM service wrapper for prompt inspection with multi-provider support.
+    """
+
+    def __init__(self, config: Dict, provider_manager: Optional[Any] = None):
+        """
+        Initialise inspection LLM service.
+
+        Args:
+            config: LLM inspection configuration
+            provider_manager: Provider manager with access to all LLM providers
+        """
+        self.config = config
+        self.provider_manager = provider_manager
+        self.model_id = config.get('model')
+        self.provider_name = config.get('provider')  # Optional: AWS Bedrock, Ollama, Anthropic Direct
+        self.max_tokens = config.get('max_tokens', 500)
+
+        # Determine which provider to use
+        self.provider = None
+        if provider_manager:
+            self.provider = self._get_provider()
+
+        if self.provider:
+            logging.info(f"Inspection LLM service initialised: provider={self.provider_name}, model={self.model_id}")
+        else:
+            logging.warning("Inspection LLM service could not initialise provider")
+
+    def _get_provider(self):
+        """
+        Get the configured LLM provider.
+
+        Returns:
+            Provider instance or None
+        """
+        try:
+            # If specific provider requested, use that
+            if self.provider_name:
+                if self.provider_name.lower() == 'aws bedrock':
+                    return self.provider_manager.get_bedrock_provider()
+                elif self.provider_name.lower() == 'ollama':
+                    return self.provider_manager.get_ollama_provider()
+                elif self.provider_name.lower() == 'anthropic direct':
+                    return self.provider_manager.get_anthropic_provider()
+
+            # Otherwise, auto-detect based on model ID
+            return self.provider_manager.get_provider_for_model(self.model_id)
+
+        except Exception as e:
+            logging.error(f"Failed to get LLM provider for inspection: {e}")
+            return None
+
+    def invoke_model(self, messages: List[Dict], max_tokens: Optional[int] = None,
+                    temperature: float = 0.1) -> Optional[Dict]:
+        """
+        Invoke LLM for prompt analysis.
+
+        Args:
+            messages: List of message dicts with role and content
+            max_tokens: Maximum tokens to generate (default from config)
+            temperature: Temperature for generation (default 0.1 for consistency)
+
+        Returns:
+            Response dict with content, or None on failure
+        """
+        if not self.provider:
+            logging.error("No LLM provider available for inspection")
+            return None
+
+        try:
+            # Use configured max_tokens if not specified
+            if max_tokens is None:
+                max_tokens = self.max_tokens
+
+            # Invoke via provider
+            response = self.provider.invoke_model(
+                model_id=self.model_id,
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                tools=None  # No tools for inspection
+            )
+
+            return response
+
+        except Exception as e:
+            logging.error(f"Failed to invoke LLM for inspection: {e}")
+            return None
+
+    def is_available(self) -> bool:
+        """
+        Check if LLM service is available for inspection.
+
+        Returns:
+            True if provider and model are configured
+        """
+        return self.provider is not None and self.model_id is not None

dtSpark/safety/patterns.py

@@ -0,0 +1,229 @@
+"""
+Attack pattern definitions for prompt inspection.
+
+This module contains regex patterns and keyword lists for detecting:
+- Prompt injection attempts
+- Jailbreak attempts
+- Code injection patterns
+- Excessive content
+- Other security risks
+
+
+"""
+
+import re
+from typing import Dict, List, Tuple, Optional
+
+# Prompt injection patterns - attempts to override system instructions
+PROMPT_INJECTION_PATTERNS = [
+    r'ignore\s+(previous|all|above|earlier|prior)\s+(instructions?|prompts?|rules?|directives?)',
+    r'disregard\s+(previous|all|above|earlier|prior)\s+(instructions?|prompts?|rules?)',
+    r'forget\s+(previous|all|above|earlier|prior)\s+(instructions?|prompts?|rules?)',
+    r'new\s+(instruction|directive|command|task|role):\s*',
+    r'system:\s*(you\s+)?(are|must|should|will)\s+now',
+    r'override\s+(previous|all|system)\s+(instructions?|settings?|rules?)',
+    r'</s>|<\|im_end\|>|<\|endoftext\|>|<\|system\|>',  # Model control tokens
+    r'\[INST\]|\[/INST\]',  # Instruction markers
+    r'<\|assistant\|>|<\|user\|>',  # Role markers
+]
+
+# Jailbreak patterns - attempts to bypass safety guidelines
+JAILBREAK_PATTERNS = [
+    r'\b(DAN|Developer\s+Mode|Sudo\s+Mode|God\s+Mode|Admin\s+Mode)\b',
+    r'pretend\s+you\s+(are|have)\s+no\s+(restrictions?|limitations?|rules?|ethics?|guidelines?)',
+    r'roleplay\s+as\s+(an?\s+)?(unrestricted|unfiltered|uncensored)',
+    r'act\s+as\s+if\s+you\s+(have\s+)?no\s+(restrictions?|limitations?|rules?)',
+    r'bypass\s+(safety|content|moderation|filter)',
+    r'without\s+(any\s+)?(ethical|moral|safety)\s+(concerns?|considerations?|restrictions?)',
+    r'evil\s+(ai|assistant|mode)',
+    r'jailbreak|jailbroken',
+]
+
+# Code injection patterns - malicious code attempts
+CODE_INJECTION_PATTERNS = [
+    r';\s*rm\s+-rf',
+    r';\s*DROP\s+TABLE',
+    r';\s*DELETE\s+FROM',
+    r'<script[^>]*>.*?</script>',
+    r'\$\([^)]+\)|`[^`]+`',  # Command substitution
+    r'\|\s*bash|\|\s*sh|\|\s*zsh',
+    r'&&\s*(rm|del|format|wget|curl)\s',
+    r'eval\s*\(',
+    r'exec\s*\(',
+    r'system\s*\(',
+    r'__import__\s*\(',
+]
+
+# PII patterns - potential personally identifiable information
+PII_PATTERNS = [
+    r'\b\d{3}-\d{2}-\d{4}\b',  # SSN pattern
+    r'\b\d{16}\b',  # Credit card pattern
+    r'\b[A-Z]{2}\d{6,8}\b',  # Passport pattern
+    r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b',  # IP address
+    r'password\s*[:=]\s*\S+',  # Password disclosure
+    r'api[_-]?key\s*[:=]\s*\S+',  # API key
+    r'secret\s*[:=]\s*\S+',  # Secret disclosure
+]
+
+# Excessive repetition pattern
+EXCESSIVE_REPETITION_PATTERN = r'(.{10,}?)\1{5,}'  # Same pattern repeated 5+ times
+
+
+class PatternMatcher:
+    """
+    Pattern-based detection for prompt security issues.
+    """
+
+    def __init__(self):
+        """Initialise pattern matcher with compiled regex patterns."""
+        self.prompt_injection_regex = [re.compile(p, re.IGNORECASE) for p in PROMPT_INJECTION_PATTERNS]
+        self.jailbreak_regex = [re.compile(p, re.IGNORECASE) for p in JAILBREAK_PATTERNS]
+        self.code_injection_regex = [re.compile(p, re.IGNORECASE) for p in CODE_INJECTION_PATTERNS]
+        self.pii_regex = [re.compile(p, re.IGNORECASE) for p in PII_PATTERNS]
+        self.repetition_regex = re.compile(EXCESSIVE_REPETITION_PATTERN, re.IGNORECASE)
+
+    def check_prompt_injection(self, text: str) -> Tuple[bool, Optional[str]]:
+        """
+        Check for prompt injection patterns.
+
+        Args:
+            text: Text to analyse
+
+        Returns:
+            Tuple of (detected, matched_pattern)
+        """
+        for pattern in self.prompt_injection_regex:
+            match = pattern.search(text)
+            if match:
+                return True, match.group(0)
+        return False, None
+
+    def check_jailbreak(self, text: str) -> Tuple[bool, Optional[str]]:
+        """
+        Check for jailbreak attempt patterns.
+
+        Args:
+            text: Text to analyse
+
+        Returns:
+            Tuple of (detected, matched_pattern)
+        """
+        for pattern in self.jailbreak_regex:
+            match = pattern.search(text)
+            if match:
+                return True, match.group(0)
+        return False, None
+
+    def check_code_injection(self, text: str) -> Tuple[bool, Optional[str]]:
+        """
+        Check for code injection patterns.
+
+        Args:
+            text: Text to analyse
+
+        Returns:
+            Tuple of (detected, matched_pattern)
+        """
+        for pattern in self.code_injection_regex:
+            match = pattern.search(text)
+            if match:
+                return True, match.group(0)
+        return False, None
+
+    def check_pii(self, text: str) -> Tuple[bool, Optional[str]]:
+        """
+        Check for potential PII exposure.
+
+        Args:
+            text: Text to analyse
+
+        Returns:
+            Tuple of (detected, matched_pattern)
+        """
+        for pattern in self.pii_regex:
+            match = pattern.search(text)
+            if match:
+                # Return redacted version
+                return True, "[REDACTED]"
+        return False, None
+
+    def check_excessive_repetition(self, text: str) -> bool:
+        """
+        Check for excessive repetition (potential DoS).
+
+        Args:
+            text: Text to analyse
+
+        Returns:
+            True if excessive repetition detected
+        """
+        return bool(self.repetition_regex.search(text))
+
+    def scan_all(self, text: str, config: Dict) -> Dict[str, any]:
+        """
+        Run all configured pattern checks.
+
+        Args:
+            text: Text to analyse
+            config: Configuration dict with enabled checks
+
+        Returns:
+            Dict with scan results
+        """
+        results = {
+            'violations': [],
+            'severity': 'none',
+            'detected_patterns': []
+        }
+
+        # Check prompt injection
+        if config.get('check_prompt_injection', True):
+            detected, pattern = self.check_prompt_injection(text)
+            if detected:
+                results['violations'].append('prompt_injection')
+                results['detected_patterns'].append(pattern)
+                results['severity'] = 'high'
+
+        # Check jailbreak
+        if config.get('check_jailbreak', True):
+            detected, pattern = self.check_jailbreak(text)
+            if detected:
+                results['violations'].append('jailbreak')
+                results['detected_patterns'].append(pattern)
+                if results['severity'] == 'none':
+                    results['severity'] = 'high'
+
+        # Check code injection
+        if config.get('check_code_injection', True):
+            detected, pattern = self.check_code_injection(text)
+            if detected:
+                results['violations'].append('code_injection')
+                results['detected_patterns'].append(pattern)
+                results['severity'] = 'critical'
+
+        # Check PII
+        if config.get('check_pii', False):
+            detected, pattern = self.check_pii(text)
+            if detected:
+                results['violations'].append('pii_exposure')
+                results['detected_patterns'].append(pattern)
+                if results['severity'] in ['none', 'low']:
+                    results['severity'] = 'medium'
+
+        # Check excessive length
+        if config.get('check_excessive_length', True):
+            max_length = config.get('max_prompt_length', 50000)
+            if len(text) > max_length:
+                results['violations'].append('excessive_length')
+                results['detected_patterns'].append(f'Length: {len(text)} > {max_length}')
+                if results['severity'] == 'none':
+                    results['severity'] = 'medium'
+
+        # Check excessive repetition
+        if self.check_excessive_repetition(text):
+            results['violations'].append('excessive_repetition')
+            results['detected_patterns'].append('Detected repetitive pattern')
+            if results['severity'] == 'none':
+                results['severity'] = 'low'
+
+        return results