#!/usr/bin/env python3
"""
Response Masker - Redacts secrets from MCP server responses before forwarding to AI
"""

import re
import hashlib
from typing import Dict, List, Tuple, Any, Optional
import logging

logger = logging.getLogger('langprotect-gateway')


class ResponseMasker:
    """Masks secrets in MCP server responses before forwarding to AI models"""
    
    # Comprehensive secret detection patterns
    # Format: (regex_pattern, secret_type, risk_score)
    # IMPORTANT: Order matters - more specific patterns should come BEFORE generic ones
    SECRET_PATTERNS = [
        # AWS Credentials (specific patterns first)
        (r'AKIA[0-9A-Z]{16}', 'AWS_ACCESS_KEY', 100),
        (r'aws_secret_access_key\s*[=:]\s*[A-Za-z0-9/+=]{40}', 'AWS_SECRET_KEY', 100),
        (r'aws_session_token\s*[=:]\s*[A-Za-z0-9/+=]{100,}', 'AWS_SESSION_TOKEN', 100),
        
        # Private Keys (PEM format)
        (r'-----BEGIN (?:RSA |OPENSSH |EC |DSA |ENCRYPTED )?PRIVATE KEY-----[^-]*-----END (?:RSA |OPENSSH |EC |DSA |ENCRYPTED )?PRIVATE KEY-----', 'PRIVATE_KEY', 100),
        (r'-----BEGIN CERTIFICATE-----[^-]*-----END CERTIFICATE-----', 'CERTIFICATE', 80),
        (r'-----BEGIN PGP PRIVATE KEY BLOCK-----[^-]*-----END PGP PRIVATE KEY BLOCK-----', 'PGP_PRIVATE_KEY', 100),
        
        # Cloud Provider Keys (BEFORE generic patterns)
        (r'(?<![A-Z_])AIza[0-9A-Za-z_\-]{35}', 'GOOGLE_API_KEY', 95),  # Negative lookbehind to avoid matching in variable names
        (r'sk_live_[0-9A-Za-z]{24,99}', 'STRIPE_LIVE_KEY', 100),  # Extended length range
        (r'sk_test_[0-9A-Za-z]{24,99}', 'STRIPE_TEST_KEY', 70),
        (r'rk_live_[0-9A-Za-z]{24,}', 'STRIPE_RESTRICTED_KEY', 95),
        
        # GitHub Tokens (BEFORE generic token pattern)
        (r'gh[pousr]_[A-Za-z0-9]{36,255}', 'GITHUB_TOKEN', 95),
        (r'github_pat_[A-Za-z0-9]{22}_[A-Za-z0-9]{59}', 'GITHUB_PAT', 95),
        
        # SSH Keys (public keys - need substantial base64)
        (r'ssh-rsa\s+[A-Za-z0-9+/]{200,}[=]{0,3}(?:\s|$)', 'SSH_RSA_PUBLIC_KEY', 70),
        (r'ssh-ed25519\s+[A-Za-z0-9+/]{68}(?:\s|$)', 'SSH_ED25519_PUBLIC_KEY', 70),
        
        # Password Fields (context-aware - must have = or : and value in quotes or after space)
        (r'(?<!#)(?<![A-Za-z])\bpassword\s*[=:]\s*["\']([^"\'\s]{8,})["\']', 'PASSWORD', 85),
        
        # Generic API Keys and Tokens (AFTER specific patterns)
        (r'["\']?api[_-]?key["\']?\s*[=:]\s*["\']?([A-Za-z0-9_\-]{20,})["\']?', 'API_KEY', 90),
        (r'["\']?token["\']?\s*[=:]\s*["\']?([A-Za-z0-9_\-\.]{20,})["\']?', 'TOKEN', 90),
        (r'["\']?secret["\']?\s*[=:]\s*["\']?([A-Za-z0-9_\-]{16,})["\']?', 'SECRET', 85),
        
        # Database Connection Strings
        (r'(postgres|mysql|mongodb|redis)://[^:]+:[^@]+@[\w\.\-]+(?::\d+)?(?:/[\w\-]+)?', 'DB_CONNECTION_STRING', 95),
        (r'Server=[\w\.\-]+;Database=[\w\-]+;User Id=[\w\-]+;Password=[^;]+', 'MSSQL_CONNECTION', 95),
        
        # JWT Tokens
        (r'eyJ[A-Za-z0-9_-]+\.eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+', 'JWT', 85),
        
        # OAuth Tokens
        (r'access_token["\']?\s*[=:]\s*["\']?([A-Za-z0-9_\-\.]{20,})["\']?', 'OAUTH_ACCESS_TOKEN', 90),
        (r'refresh_token["\']?\s*[=:]\s*["\']?([A-Za-z0-9_\-\.]{20,})["\']?', 'OAUTH_REFRESH_TOKEN', 90),
        (r'Bearer\s+([A-Za-z0-9_\-\.]{20,})', 'BEARER_TOKEN', 90),
        
        # Kubernetes Secrets (base64 encoded in YAML)
        (r'apiVersion:\s*v1\s*\nkind:\s*Secret\s*\ndata:\s*\n(?:\s+[\w\-]+:\s+[A-Za-z0-9+/=]+\n?)+', 'K8S_SECRET', 90),
        
        # Environment Variable Assignment (dangerous patterns)
        (r'export\s+(?:AWS_|DB_|API_|SECRET_|TOKEN_)[A-Z_]+=["\']?([A-Za-z0-9_\-\+/=]{16,})["\']?', 'ENV_VAR_SECRET', 85),
    ]
    
    def __init__(self, enable_entropy_detection: bool = True, entropy_threshold: float = 4.5):
        """
        Initialize the response masker.
        
        Args:
            enable_entropy_detection: Enable high-entropy string detection for unknown secret formats
            entropy_threshold: Minimum Shannon entropy for flagging potential secrets (default: 4.5)
        """
        self.enable_entropy_detection = enable_entropy_detection
        self.entropy_threshold = entropy_threshold
        logger.info(f"Response masker initialized (entropy_detection={enable_entropy_detection}, threshold={entropy_threshold})")
    
    def mask(self, content: str, context: Optional[Dict] = None) -> Tuple[str, List[Dict]]:
        """
        Mask secrets in content and return masked content + metadata.
        
        Args:
            content: The text content to scan for secrets
            context: Optional context information (file path, tool name, etc.)
        
        Returns:
            (masked_content, mask_events)
            mask_events = [{'type': 'AWS_KEY', 'hash': 'abc123...', 'risk_score': 100}, ...]
        """
        if not content or not isinstance(content, str):
            return content, []
        
        masked_content = content
        mask_events = []
        
        # Apply pattern-based detection
        for pattern, secret_type, risk_score in self.SECRET_PATTERNS:
            try:
                matches = list(re.finditer(pattern, content, re.IGNORECASE | re.MULTILINE | re.DOTALL))
                
                for match in matches:
                    original = match.group(0)
                    
                    # Skip if too short (likely false positive)
                    if len(original) < 10 and secret_type not in ['JWT', 'AWS_ACCESS_KEY']:
                        continue
                    
                    # Create hash for audit (NEVER log actual secret)
                    secret_hash = hashlib.sha256(original.encode()).hexdigest()[:16]
                    
                    # Create redaction placeholder
                    placeholder = f"<REDACTED:{secret_type}:{secret_hash}>"
                    
                    # Replace in content (handle multiple occurrences)
                    masked_content = masked_content.replace(original, placeholder)
                    
                    # Record mask event
                    mask_events.append({
                        'type': secret_type,
                        'hash': secret_hash,
                        'risk_score': risk_score,
                        'pattern': pattern[:60],  # Truncate for logging
                        'location': match.span(),
                        'length': len(original),
                        'context': context or {}
                    })
                    
                    logger.info(f"Masked {secret_type} (hash={secret_hash}, len={len(original)})")
            
            except Exception as e:
                logger.warning(f"Error applying pattern {secret_type}: {e}")
                continue
        
        # Optional: High-entropy string detection for unknown secret formats
        if self.enable_entropy_detection and not mask_events:
            entropy_matches = self._detect_high_entropy_strings(content, masked_content)
            for entropy_match in entropy_matches:
                masked_content = entropy_match['masked_content']
                mask_events.append(entropy_match['event'])
        
        return masked_content, mask_events
    
    def _detect_high_entropy_strings(self, original_content: str, current_masked: str) -> List[Dict]:
        """
        Detect high-entropy strings that might be unknown secret formats.
        Uses Shannon entropy to find random-looking strings.
        """
        results = []
        
        # Find candidate strings (alphanumeric, 16+ chars)
        pattern = r'\b[A-Za-z0-9_\-+=]{16,64}\b'
        matches = re.finditer(pattern, original_content)
        
        for match in matches:
            candidate = match.group(0)
            
            # Skip if already masked
            if candidate not in current_masked:
                continue
            
            # Calculate Shannon entropy
            entropy = self._calculate_shannon_entropy(candidate)
            
            if entropy >= self.entropy_threshold:
                secret_hash = hashlib.sha256(candidate.encode()).hexdigest()[:16]
                placeholder = f"<REDACTED:HIGH_ENTROPY:{secret_hash}>"
                current_masked = current_masked.replace(candidate, placeholder)
                
                results.append({
                    'masked_content': current_masked,
                    'event': {
                        'type': 'HIGH_ENTROPY_STRING',
                        'hash': secret_hash,
                        'risk_score': 70,
                        'pattern': 'entropy_detection',
                        'location': match.span(),
                        'length': len(candidate),
                        'entropy': round(entropy, 2),
                        'context': {}
                    }
                })
                
                logger.info(f"Masked high-entropy string (entropy={entropy:.2f}, len={len(candidate)})")
        
        return results
    
    def _calculate_shannon_entropy(self, data: str) -> float:
        """Calculate Shannon entropy of a string (bits per character)"""
        if not data:
            return 0.0
        
        import math
        from collections import Counter
        
        # Count character frequencies
        counter = Counter(data)
        length = len(data)
        
        # Calculate entropy
        entropy = 0.0
        for count in counter.values():
            probability = count / length
            if probability > 0:
                entropy -= probability * math.log2(probability)
        
        return entropy
    
    def mask_mcp_response(self, mcp_response: Dict, mask_events: List[Dict]) -> Dict:
        """
        Apply masking to an MCP JSON-RPC response structure.
        
        Handles multiple response formats:
        - Simple text result: {"result": "text content"}
        - Structured content: {"result": {"content": "..."}}
        - Content array: {"result": {"contents": [{"text": "..."}, ...]}}
        
        Args:
            mcp_response: The JSON-RPC response from MCP server
            mask_events: List to accumulate mask events (modified in-place)
        
        Returns:
            Modified MCP response with secrets masked
        """
        if not isinstance(mcp_response, dict):
            return mcp_response
        
        # Handle error responses (no masking needed)
        if 'error' in mcp_response:
            return mcp_response
        
        # Process result field
        if 'result' in mcp_response:
            result = mcp_response['result']
            
            # Case 1: Simple string result
            if isinstance(result, str):
                masked, events = self.mask(result)
                mcp_response['result'] = masked
                mask_events.extend(events)
            
            # Case 2: Dictionary result
            elif isinstance(result, dict):
                # Check for 'content' field (common in MCP responses)
                if 'content' in result:
                    content_value = result['content']
                    if isinstance(content_value, str):
                        masked, events = self.mask(content_value)
                        result['content'] = masked
                        mask_events.extend(events)
                    elif isinstance(content_value, list):
                        # Array of content items
                        for i, item in enumerate(content_value):
                            if isinstance(item, dict) and 'text' in item:
                                masked, events = self.mask(item['text'])
                                item['text'] = masked
                                mask_events.extend(events)
                
                # Check for 'contents' field (array format)
                if 'contents' in result and isinstance(result['contents'], list):
                    for item in result['contents']:
                        if isinstance(item, dict):
                            if 'text' in item and isinstance(item['text'], str):
                                masked, events = self.mask(item['text'])
                                item['text'] = masked
                                mask_events.extend(events)
                            # Also check nested content fields
                            if 'content' in item and isinstance(item['content'], str):
                                masked, events = self.mask(item['content'])
                                item['content'] = masked
                                mask_events.extend(events)
                
                # Generic fallback: scan all string values in result dict
                # (but skip metadata fields that shouldn't contain secrets)
                skip_fields = {'type', 'mimeType', 'mime_type', 'name', 'id', 'method', 'jsonrpc'}
                for key, value in result.items():
                    if isinstance(value, str) and key not in skip_fields and len(value) > 10:
                        masked, events = self.mask(value)
                        if events:  # Only replace if secrets were found
                            result[key] = masked
                            mask_events.extend(events)
        
        return mcp_response
    
    def should_mask_tool(self, tool_name: str) -> bool:
        """
        Determine if a tool's response should be masked based on tool name.
        
        High-risk tools that commonly return sensitive data:
        - File reading tools
        - Environment variable tools
        - Configuration retrieval tools
        """
        high_risk_tools = [
            'read_file', 'read_text_file', 'read_multiple_files',
            'get_file_info', 'list_directory', 'search_files',
            'get_env', 'list_env', 'read_env',
            'read_config', 'get_config', 'show_config',
            'cat', 'grep', 'search'
        ]
        
        tool_lower = tool_name.lower()
        return any(risk_tool in tool_lower for risk_tool in high_risk_tools)


# Singleton instance for reuse
_masker_instance: Optional[ResponseMasker] = None


def get_masker(enable_entropy: bool = True, entropy_threshold: float = 4.5) -> ResponseMasker:
    """Get or create the singleton response masker instance"""
    global _masker_instance
    if _masker_instance is None:
        _masker_instance = ResponseMasker(
            enable_entropy_detection=enable_entropy,
            entropy_threshold=entropy_threshold
        )
    return _masker_instance
