"""Corrected tests for exploit generator module with actual interfaces."""

import os
import sys
from unittest.mock import Mock, patch

import pytest

# Add the src directory to the path to import modules
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from adversary_mcp_server.credential_manager import CredentialManager, SecurityConfig
from adversary_mcp_server.exploit_generator import (
    ExploitContext,
    ExploitGenerationError,
    ExploitGenerator,
    LLMExploitGenerator,
    SafetyFilter,
    TemplateEngine,
)
from adversary_mcp_server.threat_engine import Category, Severity, ThreatMatch


class TestExploitContext:
    """Test ExploitContext class."""

    def test_exploit_context_initialization(self):
        """Test ExploitContext initialization."""
        threat = ThreatMatch(
            rule_id="test_rule",
            rule_name="Test Rule",
            description="Test",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
        )

        context = ExploitContext(threat, "test code", {"key": "value"})

        assert context.threat_match == threat
        assert context.source_code == "test code"
        assert context.additional_context == {"key": "value"}

    def test_exploit_context_defaults(self):
        """Test ExploitContext with default values."""
        threat = ThreatMatch(
            rule_id="test_rule",
            rule_name="Test Rule",
            description="Test",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
        )

        context = ExploitContext(threat)

        assert context.threat_match == threat
        assert context.source_code == ""
        assert context.additional_context == {}


class TestSafetyFilterCorrected:
    """Test SafetyFilter with actual implementation."""

    def test_safety_filter_initialization(self):
        """Test SafetyFilter initialization."""
        filter = SafetyFilter()
        assert filter is not None
        assert hasattr(filter, "dangerous_patterns")
        assert hasattr(filter, "safe_replacements")

    def test_filter_dangerous_commands(self):
        """Test filtering of dangerous commands."""
        filter = SafetyFilter()

        # Test some dangerous commands
        dangerous_exploits = [
            "rm -rf /",
            "format c:",
            "shutdown now",
            "dd if=/dev/zero of=/dev/sda",
        ]

        for exploit in dangerous_exploits:
            filtered = filter.filter_exploit(exploit)
            # Should be filtered or modified
            assert "[REDACTED" in filtered or filtered != exploit

    def test_filter_safe_replacements(self):
        """Test safe replacements."""
        filter = SafetyFilter()

        # Test safe replacements
        exploit = "cat /etc/passwd"
        filtered = filter.filter_exploit(exploit)
        assert "/tmp/example_file" in filtered

    def test_is_safe_exploit(self):
        """Test is_safe_exploit method."""
        filter = SafetyFilter()

        # Safe exploit
        assert filter.is_safe_exploit("' OR '1'='1' --")

        # Dangerous exploit
        assert not filter.is_safe_exploit("rm -rf /")


class TestTemplateEngineCorrected:
    """Test TemplateEngine with actual implementation."""

    def test_template_engine_initialization(self):
        """Test TemplateEngine initialization."""
        engine = TemplateEngine()
        assert engine is not None
        assert hasattr(engine, "templates")
        assert hasattr(engine, "env")

    def test_template_engine_has_default_templates(self):
        """Test that template engine has default templates."""
        engine = TemplateEngine()

        # Check for common vulnerability templates
        assert "sql_injection" in engine.templates
        assert "command_injection" in engine.templates
        assert "xss_dom" in engine.templates
        assert "deserialization" in engine.templates
        assert "path_traversal" in engine.templates

    def test_generate_from_template(self):
        """Test generating exploit from template."""
        engine = TemplateEngine()

        context = {
            "vulnerable_param": "user_id",
            "original_query": "SELECT * FROM users WHERE id = ?",
        }

        result = engine.generate_from_template("sql_injection", context)
        assert isinstance(result, str)
        assert len(result) > 0
        assert "user_id" in result

    def test_add_template(self):
        """Test adding custom template."""
        engine = TemplateEngine()

        custom_template = "Custom exploit for {{ param }}"
        engine.add_template("custom_test", custom_template)

        # Should be able to use the custom template
        result = engine.generate_from_template("custom_test", {"param": "test"})
        assert "Custom exploit for test" in result

    def test_template_with_missing_context(self):
        """Test template with missing context variables."""
        engine = TemplateEngine()

        # Should handle missing context gracefully
        result = engine.generate_from_template("sql_injection", {})
        assert isinstance(result, str)


class TestLLMExploitGeneratorCorrected:
    """Test LLMExploitGenerator with actual implementation."""

    def test_llm_generator_initialization(self):
        """Test LLM generator initialization."""
        config = SecurityConfig(openai_api_key="sk-test123")
        generator = LLMExploitGenerator(config)

        assert generator.config == config

    def test_llm_generator_without_api_key(self):
        """Test LLM generator without API key."""
        config = SecurityConfig(openai_api_key="")
        generator = LLMExploitGenerator(config)

        assert generator.config == config

    @patch("adversary_mcp_server.exploit_generator.openai")
    def test_generate_exploit_with_llm(self, mock_openai):
        """Test exploit generation with LLM."""
        config = SecurityConfig(openai_api_key="sk-test123")
        generator = LLMExploitGenerator(config)

        # Mock OpenAI client
        mock_client = Mock()
        mock_openai.OpenAI.return_value = mock_client

        # Mock OpenAI response - ensure content is a string
        mock_response = Mock()
        mock_response.choices = [Mock()]
        mock_response.choices[0].message = Mock()
        mock_response.choices[0].message.content = "Generated exploit content"
        mock_client.chat.completions.create.return_value = mock_response

        # Also ensure the actual client is set on the generator
        generator.client = mock_client

        threat = ThreatMatch(
            rule_id="sql_injection",
            rule_name="SQL Injection",
            description="SQL injection vulnerability",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
        )

        context = ExploitContext(threat, "test code")
        result = generator.generate_exploit(context)

        assert isinstance(result, list)
        assert len(result) > 0

    def test_create_prompt(self):
        """Test prompt creation."""
        config = SecurityConfig()
        generator = LLMExploitGenerator(config)

        threat = ThreatMatch(
            rule_id="xss",
            rule_name="XSS",
            description="XSS vulnerability",
            category=Category.XSS,
            severity=Severity.MEDIUM,
            file_path="test.js",
            line_number=1,
        )

        context = ExploitContext(threat, "document.innerHTML = userInput")
        prompt = generator._create_prompt(context)

        assert isinstance(prompt, str)
        assert "XSS" in prompt
        assert "userInput" in prompt

    def test_parse_exploits(self):
        """Test parsing exploits from text."""
        config = SecurityConfig()
        generator = LLMExploitGenerator(config)

        text = """
        Here are some exploits:
        
        ```
        ' OR '1'='1' --
        ```
        
        ```
        ' UNION SELECT * FROM users --
        ```
        """

        exploits = generator._parse_exploits(text)
        assert isinstance(exploits, list)
        assert len(exploits) >= 2


class TestExploitGeneratorCorrected:
    """Test ExploitGenerator with actual implementation."""

    def test_exploit_generator_initialization(self):
        """Test ExploitGenerator initialization."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        assert generator.credential_manager == mock_manager
        assert hasattr(generator, "safety_filter")
        assert hasattr(generator, "template_engine")
        assert hasattr(generator, "llm_generator")

    def test_generate_exploits_with_template(self):
        """Test exploit generation using templates."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        threat = ThreatMatch(
            rule_id="sql_injection",
            rule_name="SQL Injection",
            description="SQL injection vulnerability",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
            code_snippet="SELECT * FROM users WHERE id = " + "user_input",
        )

        # Generate exploits without LLM
        result = generator.generate_exploits(threat, "test code", use_llm=False)

        assert isinstance(result, list)
        assert len(result) > 0
        assert all(isinstance(exploit, str) for exploit in result)

    def test_generate_exploits_with_llm(self):
        """Test exploit generation with LLM (mocked)."""
        mock_manager = Mock()
        mock_config = SecurityConfig(openai_api_key="sk-test123")
        mock_manager.load_config.return_value = mock_config

        with patch("adversary_mcp_server.exploit_generator.openai") as mock_openai:
            # Mock successful API response
            mock_response = Mock()
            mock_response.choices = [Mock()]
            mock_response.choices[0].message.content = "Generated exploit"
            mock_openai.ChatCompletion.create.return_value = mock_response

            generator = ExploitGenerator(mock_manager)

            threat = ThreatMatch(
                rule_id="xss",
                rule_name="XSS",
                description="XSS vulnerability",
                category=Category.XSS,
                severity=Severity.MEDIUM,
                file_path="test.js",
                line_number=1,
            )

            result = generator.generate_exploits(threat, "test code", use_llm=True)

            assert isinstance(result, list)

    def test_is_llm_available(self):
        """Test LLM availability check."""
        mock_manager = Mock()

        # With API key
        mock_config = SecurityConfig(openai_api_key="sk-test123")
        mock_manager.load_config.return_value = mock_config
        generator = ExploitGenerator(mock_manager)
        assert generator.is_llm_available() is True

        # Without API key
        mock_config = SecurityConfig(openai_api_key="")
        mock_manager.load_config.return_value = mock_config
        generator = ExploitGenerator(mock_manager)
        assert generator.is_llm_available() is False

    def test_get_exploit_metadata(self):
        """Test getting exploit metadata."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        threat = ThreatMatch(
            rule_id="command_injection",
            rule_name="Command Injection",
            description="Command injection vulnerability",
            category=Category.INJECTION,
            severity=Severity.CRITICAL,
            file_path="test.py",
            line_number=1,
        )

        metadata = generator.get_exploit_metadata(threat)

        assert isinstance(metadata, dict)
        assert "category" in metadata
        assert "severity" in metadata
        # Check for actual keys returned by the method
        assert "llm_available" in metadata or "available_templates" in metadata

    def test_add_custom_template(self):
        """Test adding custom template."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        custom_template = "Custom {{ vulnerability }} exploit"
        generator.add_custom_template("custom_test", custom_template)

        # Should be added to template engine
        assert "custom_test" in generator.template_engine.templates

    def test_generate_from_templates_method(self):
        """Test _generate_from_templates method."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        threat = ThreatMatch(
            rule_id="path_traversal",
            rule_name="Path Traversal",
            description="Path traversal vulnerability",
            category=Category.LFI,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
        )

        result = generator._generate_from_templates(threat, "test code")

        assert isinstance(result, list)

    def test_get_available_templates(self):
        """Test getting available templates."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        templates = generator._get_available_templates(Category.INJECTION)

        assert isinstance(templates, list)

    def test_exploit_generation_error(self):
        """Test ExploitGenerationError."""
        error = ExploitGenerationError("Test error")
        assert str(error) == "Test error"
        assert isinstance(error, Exception)

    def test_safety_filtering_integration(self):
        """Test that safety filtering is applied."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        # Create a threat that might generate dangerous content
        threat = ThreatMatch(
            rule_id="command_injection",
            rule_name="Command Injection",
            description="Command injection vulnerability",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
        )

        result = generator.generate_exploits(
            threat, "os.system(user_input)", use_llm=False
        )

        # Results should be filtered for safety
        assert isinstance(result, list)
        for exploit in result:
            assert isinstance(exploit, str)
            # Should not contain extremely dangerous patterns
            assert "rm -rf /" not in exploit

    def test_different_vulnerability_categories(self):
        """Test exploit generation for different vulnerability categories."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        categories = [
            Category.INJECTION,
            Category.XSS,
            Category.DESERIALIZATION,
            Category.LFI,
        ]

        for category in categories:
            threat = ThreatMatch(
                rule_id=f"test_{category.value}",
                rule_name=f"Test {category.value}",
                description=f"Test {category.value} vulnerability",
                category=category,
                severity=Severity.MEDIUM,
                file_path="test.py",
                line_number=1,
            )

            result = generator.generate_exploits(threat, "test code", use_llm=False)
            assert isinstance(result, list)

    def test_exploit_generation_with_context(self):
        """Test exploit generation with different contexts."""
        mock_manager = Mock()
        mock_config = SecurityConfig()
        mock_manager.load_config.return_value = mock_config

        generator = ExploitGenerator(mock_manager)

        threat = ThreatMatch(
            rule_id="sql_injection",
            rule_name="SQL Injection",
            description="SQL injection vulnerability",
            category=Category.INJECTION,
            severity=Severity.HIGH,
            file_path="test.py",
            line_number=1,
            code_snippet="SELECT * FROM users WHERE id = " + "user_input",
        )

        # Test with different source code contexts
        contexts = [
            "SELECT * FROM users WHERE id = user_input",
            "query = 'SELECT * FROM products WHERE name = ' + search_term",
            "database.execute('SELECT * FROM orders WHERE customer_id = ' + customer_id)",
        ]

        for context in contexts:
            result = generator.generate_exploits(threat, context, use_llm=False)
            assert isinstance(result, list)
            assert len(result) > 0
