#!/usr/bin/env python3
"""
Unit tests for Response Masker
"""

import unittest
import json
from langprotect_mcp_gateway.response_masker import ResponseMasker, get_masker


class TestResponseMasker(unittest.TestCase):
    
    def setUp(self):
        """Initialize masker for each test"""
        self.masker = ResponseMasker(enable_entropy_detection=False)
    
    def test_mask_aws_access_key(self):
        """Test AWS access key detection"""
        content = "aws_access_key_id=AKIAIOSFODNN7EXAMPLE"
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("AKIAIOSFODNN7EXAMPLE", masked)
        self.assertIn("<REDACTED:AWS_ACCESS_KEY:", masked)
        self.assertEqual(len(events), 1)
        self.assertEqual(events[0]['type'], 'AWS_ACCESS_KEY')
        self.assertEqual(events[0]['risk_score'], 100)
    
    def test_mask_aws_secret_key(self):
        """Test AWS secret key detection"""
        content = "aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("wJalrXUtnFEMI", masked)
        self.assertIn("<REDACTED:AWS_SECRET_KEY:", masked)
        self.assertGreaterEqual(len(events), 1)
    
    def test_mask_private_key(self):
        """Test PEM private key detection"""
        content = """
        -----BEGIN RSA PRIVATE KEY-----
        MIIEpAIBAAKCAQEA0Z8Yjf3qTQXZ3...
        -----END RSA PRIVATE KEY-----
        """
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("-----BEGIN RSA PRIVATE KEY-----", masked)
        self.assertIn("<REDACTED:PRIVATE_KEY:", masked)
        self.assertEqual(len(events), 1)
        self.assertEqual(events[0]['type'], 'PRIVATE_KEY')
    
    def test_mask_jwt_token(self):
        """Test JWT token detection"""
        content = "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9", masked)
        self.assertIn("<REDACTED:JWT:", masked)
    
    def test_mask_database_connection_string(self):
        """Test database connection string masking"""
        content = "postgres://admin:SuperSecret123@db.example.com:5432/mydb"
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("SuperSecret123", masked)
        self.assertIn("<REDACTED:DB_CONNECTION_STRING:", masked)
    
    def test_mask_api_key_in_json(self):
        """Test API key detection in JSON"""
        content = '{"api_key": "sk_live_abc123def456ghi789jkl012", "name": "test"}'
        masked, events = self.masker.mask(content)
        
        # Ensure JSON structure preserved
        self.assertIn('"name"', masked)
        self.assertIn('"test"', masked)
        self.assertIn("<REDACTED:", masked)
    
    def test_mask_github_token(self):
        """Test GitHub token detection"""
        content = "GITHUB_TOKEN=ghp_abc123XYZ789def456GHI123jkl456MNO789"
        masked, events = self.masker.mask(content)
        
        self.assertNotIn("ghp_abc123XYZ789", masked)
        self.assertIn("<REDACTED:GITHUB_TOKEN:", masked)
    
    def test_mask_stripe_keys(self):
        """Test Stripe API key detection"""
        # Live key (high risk)
        content_live = "sk_live_51HabcdefghijklmnopQR"
        masked_live, events_live = self.masker.mask(content_live)
        self.assertIn("<REDACTED:STRIPE_LIVE_KEY:", masked_live)
        self.assertEqual(events_live[0]['risk_score'], 100)
        
        # Test key (lower risk)
        content_test = "sk_test_51HabcdefghijklmnopQR"
        masked_test, events_test = self.masker.mask(content_test)
        self.assertIn("<REDACTED:STRIPE_TEST_KEY:", masked_test)
        self.assertEqual(events_test[0]['risk_score'], 70)
    
    def test_mask_multiple_secrets(self):
        """Test masking multiple different secrets"""
        content = """
        AWS_KEY=AKIAIOSFODNN7EXAMPLE
        DB_URL=postgres://user:pass@host/db
        API_TOKEN=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.signature
        """
        masked, events = self.masker.mask(content)
        
        # Should detect at least 3 secrets
        self.assertGreaterEqual(len(events), 3)
        
        # Check all are masked
        self.assertNotIn("AKIAIOSFODNN7EXAMPLE", masked)
        self.assertNotIn("postgres://user:pass", masked)
    
    def test_mask_preserves_structure(self):
        """Test that masking preserves document structure"""
        content = '''
        # Configuration File
        [database]
        host = localhost
        password = MySecretPassword123
        
        [api]
        key = sk_test_abcdef123456
        '''
        masked, events = self.masker.mask(content)
        
        # Structure should be preserved
        self.assertIn("# Configuration File", masked)
        self.assertIn("[database]", masked)
        self.assertIn("[api]", masked)
        
        # Secrets should be masked
        self.assertNotIn("MySecretPassword123", masked)
        self.assertNotIn("sk_test_abcdef123456", masked)
    
    def test_mask_mcp_response_simple(self):
        """Test masking MCP response with simple result"""
        response = {
            'jsonrpc': '2.0',
            'id': 1,
            'result': 'AWS_KEY=AKIAIOSFODNN7EXAMPLE'
        }
        
        mask_events = []
        masked_response = self.masker.mask_mcp_response(response, mask_events)
        
        self.assertNotIn('AKIAIOSFODNN7EXAMPLE', masked_response['result'])
        self.assertIn('<REDACTED:AWS_ACCESS_KEY:', masked_response['result'])
        self.assertEqual(len(mask_events), 1)
    
    def test_mask_mcp_response_structured(self):
        """Test masking MCP response with structured content"""
        response = {
            'jsonrpc': '2.0',
            'id': 1,
            'result': {
                'content': 'db_password=SuperSecret123',
                'mimeType': 'text/plain'
            }
        }
        
        mask_events = []
        masked_response = self.masker.mask_mcp_response(response, mask_events)
        
        self.assertNotIn('SuperSecret123', str(masked_response))
        self.assertIn('<REDACTED:', masked_response['result']['content'])
    
    def test_mask_mcp_response_contents_array(self):
        """Test masking MCP response with contents array"""
        response = {
            'jsonrpc': '2.0',
            'id': 1,
            'result': {
                'contents': [
                    {'type': 'text', 'text': 'API_KEY=abc123def456ghi789jkl012mno345'},
                    {'type': 'text', 'text': 'No secrets here'}
                ]
            }
        }
        
        mask_events = []
        masked_response = self.masker.mask_mcp_response(response, mask_events)
        
        # First item should be masked
        self.assertIn('<REDACTED:', masked_response['result']['contents'][0]['text'])
        
        # Second item should be unchanged
        self.assertEqual(masked_response['result']['contents'][1]['text'], 'No secrets here')
    
    def test_mask_error_responses_unchanged(self):
        """Test that error responses are not modified"""
        response = {
            'jsonrpc': '2.0',
            'id': 1,
            'error': {
                'code': -32000,
                'message': 'Error with AWS_KEY=AKIAIOSFODNN7EXAMPLE'
            }
        }
        
        mask_events = []
        masked_response = self.masker.mask_mcp_response(response, mask_events)
        
        # Error responses should not be masked (to preserve error details)
        self.assertEqual(response, masked_response)
        self.assertEqual(len(mask_events), 0)
    
    def test_should_mask_tool_detection(self):
        """Test tool name detection for selective masking"""
        # High-risk tools
        self.assertTrue(self.masker.should_mask_tool('read_file'))
        self.assertTrue(self.masker.should_mask_tool('read_text_file'))
        self.assertTrue(self.masker.should_mask_tool('get_env'))
        self.assertTrue(self.masker.should_mask_tool('list_directory'))
        
        # Low-risk tools (should not mask)
        self.assertFalse(self.masker.should_mask_tool('add'))
        self.assertFalse(self.masker.should_mask_tool('multiply'))
        self.assertFalse(self.masker.should_mask_tool('get_weather'))
    
    def test_no_false_positives_on_safe_content(self):
        """Test that safe content is not incorrectly masked"""
        safe_content = """
        This is a normal document with no secrets.
        It contains regular text, numbers like 12345,
        and common words like 'password' in documentation.
        """
        masked, events = self.masker.mask(safe_content)
        
        # No events should be triggered
        self.assertEqual(len(events), 0)
        # Content should be unchanged
        self.assertEqual(safe_content, masked)
    
    def test_high_entropy_detection(self):
        """Test high-entropy string detection (optional feature)"""
        masker_with_entropy = ResponseMasker(enable_entropy_detection=True, entropy_threshold=4.5)
        
        # High-entropy string (looks like a secret)
        content = "secret_value = xK9pL2mN5qR8tV3wY6zA"
        masked, events = masker_with_entropy.mask(content)
        
        # Should detect as high-entropy
        if events:  # Only if entropy detection works
            self.assertTrue(any(e['type'] == 'HIGH_ENTROPY_STRING' for e in events))
    
    def test_mask_ssh_keys(self):
        """Test SSH key detection"""
        content = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDZzKz... user@host"
        masked, events = self.masker.mask(content)
        
        self.assertIn("<REDACTED:SSH_RSA_PUBLIC_KEY:", masked)
    
    def test_mask_google_api_key(self):
        """Test Google API key detection"""
        content = "GOOGLE_API_KEY=AIzaSyD1234567890abcdefghijklmnopqrst"
        masked, events = self.masker.mask(content)
        
        self.assertIn("<REDACTED:GOOGLE_API_KEY:", masked)
    
    def test_singleton_masker(self):
        """Test singleton pattern for masker"""
        masker1 = get_masker()
        masker2 = get_masker()
        
        # Should return same instance
        self.assertIs(masker1, masker2)


if __name__ == '__main__':
    unittest.main()
