"""
BREACH.AI - Script Generator

Gives the agent the power to write custom exploitation scripts on-the-fly.
This is what makes BREACH.AI truly autonomous - it can adapt to any situation
by generating custom code to exploit unique vulnerabilities.

The agent can:
1. Write custom exploitation scripts when standard attacks fail
2. Create data extraction scripts for specific database schemas
3. Generate credential harvesters for custom auth systems
4. Build custom fuzzers for unusual input formats
5. Create proof-of-concept scripts for the report
"""

import asyncio
import json
import os
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional

from anthropic import Anthropic

from breach.utils.logger import logger


class ScriptLanguage(Enum):
    """Supported scripting languages."""
    PYTHON = "python"
    BASH = "bash"
    JAVASCRIPT = "javascript"
    SQL = "sql"


class ScriptPurpose(Enum):
    """What the script is for."""
    EXPLOITATION = "exploitation"
    DATA_EXTRACTION = "data_extraction"
    ENUMERATION = "enumeration"
    FUZZING = "fuzzing"
    CREDENTIAL_HARVEST = "credential_harvest"
    PROOF_OF_CONCEPT = "poc"
    CUSTOM_PAYLOAD = "payload"
    WAF_BYPASS = "waf_bypass"
    CHAIN_EXPLOIT = "chain_exploit"


@dataclass
class GeneratedScript:
    """A script generated by the agent."""
    language: ScriptLanguage
    purpose: ScriptPurpose
    code: str
    description: str
    target: str

    # Execution details
    dependencies: list[str] = field(default_factory=list)
    timeout_seconds: int = 60
    requires_network: bool = True

    # Safety
    is_destructive: bool = False
    risk_level: str = "medium"

    # Results
    executed: bool = False
    output: Optional[str] = None
    success: bool = False
    error: Optional[str] = None


@dataclass
class ScriptResult:
    """Result of script execution."""
    success: bool
    output: str
    error: Optional[str] = None
    data_extracted: Optional[Any] = None
    execution_time: float = 0.0


# System prompt for script generation
SCRIPT_GEN_PROMPT = """You are the script-writing module of BREACH.AI, an autonomous security assessment agent.

Your job is to write custom exploitation scripts when standard attacks don't work or when we need to extract specific data.

CURRENT CONTEXT:
Target: {target}
Vulnerability Found: {vulnerability}
Current Access Level: {access_level}
Goal: {goal}

AVAILABLE INFORMATION:
{context}

CONSTRAINTS:
- Scripts must be self-contained and executable
- Include all necessary imports/dependencies
- Handle errors gracefully
- Never cause permanent damage (read-only where possible)
- Output results in a parseable format (JSON preferred)
- Include timeout handling
- Maximum execution time: {timeout} seconds

SCRIPT REQUIREMENTS:
Language: {language}
Purpose: {purpose}

Write a script that accomplishes the goal. The script should:
1. Be production-quality and handle edge cases
2. Output clear results that can be parsed
3. Include comments explaining key sections
4. Be safe to execute in a sandboxed environment

Return your response in this exact JSON format:
{{
    "code": "the full script code here",
    "description": "what this script does",
    "dependencies": ["list", "of", "required", "packages"],
    "expected_output": "description of what success looks like",
    "risk_level": "low|medium|high",
    "is_destructive": false,
    "execution_notes": "any special notes about running this"
}}

IMPORTANT: Return ONLY the JSON, no other text."""


EXPLOITATION_PROMPT = """
Write an exploitation script for the following vulnerability:

Vulnerability Type: {vuln_type}
Target Endpoint: {endpoint}
Vulnerable Parameter: {parameter}
Current Payload: {current_payload}
Error/Response: {response}

The standard payload didn't work. Analyze the response and write a custom
exploitation script that bypasses any protections and successfully exploits
this vulnerability.

Consider:
- WAF bypass techniques
- Encoding variations
- Alternative injection points
- Time-based or blind techniques if direct doesn't work
"""


DATA_EXTRACTION_PROMPT = """
Write a data extraction script to prove the severity of this breach:

Access Type: {access_type}
Database/System: {system}
Schema Info: {schema}
Goal: Extract sample data as evidence

The script should:
1. Connect using the credentials/access we have
2. Enumerate what data is available
3. Extract a SAMPLE (not everything) to prove access
4. Format the output as evidence for the report
5. Count total records without downloading all

We want to prove impact without being destructive.
"""


WAF_BYPASS_PROMPT = """
Our payload was blocked by a WAF. Generate a bypass script.

Target: {target}
Original Payload: {blocked_payload}
WAF Response: {waf_response}
Vulnerability Type: {vuln_type}

Analyze the WAF behavior and generate a script that:
1. Identifies the WAF type if possible
2. Tries multiple bypass techniques:
   - Case variation
   - Encoding (URL, double URL, Unicode, hex)
   - Comment insertion
   - Alternate syntax
   - HTTP parameter pollution
   - Chunked encoding
   - Null bytes
3. Finds a working bypass
4. Reports which technique worked

Return as JSON with the script in the "code" field.
"""


class ScriptGenerator:
    """
    Generates custom scripts on-the-fly using Claude.

    This is the "infinite adaptability" power of BREACH.AI.
    When standard attacks fail, we write custom ones.
    """

    def __init__(self, model: str = "claude-sonnet-4-20250514"):
        self.client = Anthropic()
        self.model = model
        self.generated_scripts: list[GeneratedScript] = []
        self.sandbox_dir = tempfile.mkdtemp(prefix="breach_scripts_")

    async def generate_exploitation_script(
        self,
        target: str,
        vuln_type: str,
        endpoint: str,
        parameter: str,
        current_payload: str,
        response: str,
        context: dict
    ) -> GeneratedScript:
        """Generate a custom exploitation script when standard attacks fail."""
        prompt = SCRIPT_GEN_PROMPT.format(
            target=target,
            vulnerability=f"{vuln_type} in {parameter} at {endpoint}",
            access_level=context.get("access_level", "none"),
            goal=f"Successfully exploit {vuln_type} vulnerability",
            context=json.dumps(context, indent=2),
            timeout=60,
            language="python",
            purpose="exploitation"
        )

        prompt += "\n\n" + EXPLOITATION_PROMPT.format(
            vuln_type=vuln_type,
            endpoint=endpoint,
            parameter=parameter,
            current_payload=current_payload,
            response=response[:2000]
        )

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.EXPLOITATION,
            target
        )

    async def generate_data_extraction_script(
        self,
        target: str,
        access_type: str,
        system: str,
        credentials: dict,
        schema_info: str,
        context: dict
    ) -> GeneratedScript:
        """Generate a script to extract data as proof of breach."""
        prompt = SCRIPT_GEN_PROMPT.format(
            target=target,
            vulnerability=f"Full {access_type} access to {system}",
            access_level=context.get("access_level", "database"),
            goal="Extract sample data as evidence of breach severity",
            context=json.dumps({**context, "credentials": "REDACTED"}, indent=2),
            timeout=120,
            language="python",
            purpose="data_extraction"
        )

        prompt += "\n\n" + DATA_EXTRACTION_PROMPT.format(
            access_type=access_type,
            system=system,
            schema=schema_info
        )

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.DATA_EXTRACTION,
            target
        )

    async def generate_waf_bypass_script(
        self,
        target: str,
        blocked_payload: str,
        waf_response: str,
        vuln_type: str,
        context: dict
    ) -> GeneratedScript:
        """Generate a script to bypass WAF protection."""
        prompt = WAF_BYPASS_PROMPT.format(
            target=target,
            blocked_payload=blocked_payload,
            waf_response=waf_response[:1000],
            vuln_type=vuln_type
        )

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.WAF_BYPASS,
            target
        )

    async def generate_custom_fuzzer(
        self,
        target: str,
        endpoint: str,
        parameter: str,
        input_format: str,
        context: dict
    ) -> GeneratedScript:
        """Generate a custom fuzzer for unusual input formats."""
        prompt = SCRIPT_GEN_PROMPT.format(
            target=target,
            vulnerability=f"Potential vulnerability in {parameter} at {endpoint}",
            access_level=context.get("access_level", "none"),
            goal=f"Fuzz {parameter} with format-specific payloads",
            context=json.dumps(context, indent=2),
            timeout=300,
            language="python",
            purpose="fuzzing"
        )

        prompt += f"""

Write a custom fuzzer for this unusual input format:

Endpoint: {endpoint}
Parameter: {parameter}
Input Format: {input_format}

The fuzzer should:
1. Understand the input format
2. Generate valid-looking but malicious inputs
3. Test boundary conditions
4. Try format-specific injection techniques
5. Report any anomalous responses
"""

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.FUZZING,
            target
        )

    async def generate_chain_script(
        self,
        vulnerabilities: list[dict],
        goal: str,
        context: dict
    ) -> GeneratedScript:
        """Generate a script that chains multiple vulnerabilities together."""
        prompt = f"""
Generate a script that chains these vulnerabilities together:

VULNERABILITIES FOUND:
{json.dumps(vulnerabilities, indent=2)}

GOAL: {goal}

Create a script that:
1. Exploits vulnerability #1 to get initial access/data
2. Uses that access to exploit vulnerability #2
3. Chains through all vulnerabilities
4. Achieves the final goal
5. Documents each step for the report

This chain should demonstrate how multiple "medium" issues
combine into a critical security failure.

Return as JSON with the script in the "code" field.
"""

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.CHAIN_EXPLOIT,
            context.get("target", "unknown")
        )

    async def generate_poc_script(
        self,
        vulnerability: dict,
        context: dict
    ) -> GeneratedScript:
        """Generate a clean proof-of-concept script for the report."""
        prompt = f"""
Generate a clean, well-documented proof-of-concept script for this vulnerability:

Vulnerability: {vulnerability.get('title', 'Unknown')}
Type: {vulnerability.get('vuln_type', 'Unknown')}
Target: {vulnerability.get('target', 'Unknown')}
Endpoint: {vulnerability.get('endpoint', 'Unknown')}
Parameter: {vulnerability.get('parameter', 'Unknown')}

Details: {vulnerability.get('details', 'N/A')}

Create a standalone Python script that:
1. Clearly demonstrates the vulnerability
2. Is well-commented for non-technical readers
3. Shows the exact request/response
4. Can be run by the client to verify the issue
5. Is safe to include in a security report

The script should be professional and suitable for inclusion in a
penetration testing report.

Return as JSON with the script in the "code" field.
"""

        return await self._generate_and_parse(
            prompt,
            ScriptLanguage.PYTHON,
            ScriptPurpose.PROOF_OF_CONCEPT,
            vulnerability.get('target', 'unknown')
        )

    async def _generate_and_parse(
        self,
        prompt: str,
        language: ScriptLanguage,
        purpose: ScriptPurpose,
        target: str
    ) -> GeneratedScript:
        """Generate script using Claude and parse the response."""
        try:
            response = self.client.messages.create(
                model=self.model,
                max_tokens=4000,
                messages=[
                    {"role": "user", "content": prompt}
                ]
            )

            content = response.content[0].text

            # Parse JSON response
            if "```json" in content:
                start = content.find("```json") + 7
                end = content.find("```", start)
                content = content[start:end]
            elif "```" in content:
                start = content.find("```") + 3
                end = content.find("```", start)
                content = content[start:end]

            data = json.loads(content.strip())

            script = GeneratedScript(
                language=language,
                purpose=purpose,
                code=data.get("code", ""),
                description=data.get("description", ""),
                target=target,
                dependencies=data.get("dependencies", []),
                timeout_seconds=60,
                is_destructive=data.get("is_destructive", False),
                risk_level=data.get("risk_level", "medium"),
            )

            self.generated_scripts.append(script)
            logger.info(f"Generated {purpose.value} script: {script.description[:50]}...")

            return script

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse script response: {e}")
            return GeneratedScript(
                language=language,
                purpose=purpose,
                code="# Generation failed",
                description=f"Failed to generate: {str(e)}",
                target=target,
                dependencies=[],
                error=str(e)
            )
        except Exception as e:
            logger.error(f"Script generation failed: {e}")
            return GeneratedScript(
                language=language,
                purpose=purpose,
                code="# Generation failed",
                description=f"Failed to generate: {str(e)}",
                target=target,
                dependencies=[],
                error=str(e)
            )

    async def execute_script(
        self,
        script: GeneratedScript,
        context: dict = None
    ) -> ScriptResult:
        """Execute a generated script in a sandboxed environment."""
        if not script.code or script.code.startswith("# Generation failed"):
            return ScriptResult(
                success=False,
                output="",
                error="Script generation failed"
            )

        logger.info(f"Executing script: {script.description[:50]}...")

        start_time = time.time()

        try:
            if script.language == ScriptLanguage.PYTHON:
                script_path = os.path.join(self.sandbox_dir, "script.py")
                with open(script_path, 'w') as f:
                    f.write(script.code)

                # Install dependencies if needed
                if script.dependencies:
                    for dep in script.dependencies:
                        subprocess.run(
                            ["pip", "install", dep, "-q"],
                            capture_output=True,
                            timeout=30
                        )

                # Execute with timeout
                result = subprocess.run(
                    ["python", script_path],
                    capture_output=True,
                    text=True,
                    timeout=script.timeout_seconds,
                    cwd=self.sandbox_dir,
                    env={**os.environ, "PYTHONUNBUFFERED": "1"}
                )

                execution_time = time.time() - start_time

                script.executed = True
                script.output = result.stdout

                if result.returncode == 0:
                    script.success = True

                    # Try to parse JSON output
                    data_extracted = None
                    try:
                        data_extracted = json.loads(result.stdout)
                    except json.JSONDecodeError:
                        pass

                    return ScriptResult(
                        success=True,
                        output=result.stdout,
                        data_extracted=data_extracted,
                        execution_time=execution_time
                    )
                else:
                    script.error = result.stderr
                    return ScriptResult(
                        success=False,
                        output=result.stdout,
                        error=result.stderr,
                        execution_time=execution_time
                    )

            elif script.language == ScriptLanguage.BASH:
                script_path = os.path.join(self.sandbox_dir, "script.sh")
                with open(script_path, 'w') as f:
                    f.write(script.code)

                result = subprocess.run(
                    ["bash", script_path],
                    capture_output=True,
                    text=True,
                    timeout=script.timeout_seconds,
                    cwd=self.sandbox_dir
                )

                execution_time = time.time() - start_time

                return ScriptResult(
                    success=result.returncode == 0,
                    output=result.stdout,
                    error=result.stderr if result.returncode != 0 else None,
                    execution_time=execution_time
                )

            else:
                return ScriptResult(
                    success=False,
                    output="",
                    error=f"Unsupported language: {script.language.value}"
                )

        except subprocess.TimeoutExpired:
            return ScriptResult(
                success=False,
                output="",
                error=f"Script timed out after {script.timeout_seconds} seconds"
            )
        except Exception as e:
            return ScriptResult(
                success=False,
                output="",
                error=str(e)
            )

    def get_all_scripts(self) -> list[GeneratedScript]:
        """Get all scripts generated during this session."""
        return self.generated_scripts

    def get_successful_scripts(self) -> list[GeneratedScript]:
        """Get scripts that executed successfully."""
        return [s for s in self.generated_scripts if s.success]

    def cleanup(self):
        """Clean up temporary files."""
        import shutil
        try:
            shutil.rmtree(self.sandbox_dir)
        except Exception as e:
            logger.warning(f"Failed to cleanup sandbox: {e}")


class ScriptingCapability:
    """
    Adds scripting capability to the agent brain.

    This is called when:
    1. Standard attacks fail and we need custom approach
    2. We need to extract data in a specific format
    3. We want to chain vulnerabilities
    4. We need PoC scripts for the report
    """

    def __init__(self):
        self.generator = ScriptGenerator()

    async def handle_failed_attack(
        self,
        attack_type: str,
        target: str,
        endpoint: str,
        parameter: str,
        payload: str,
        response: str,
        context: dict
    ) -> Optional[ScriptResult]:
        """Called when a standard attack fails. Generates and executes a custom script."""
        logger.info(f"Standard {attack_type} attack failed. Generating custom script...")

        # Check if it looks like WAF blocking
        waf_indicators = ["blocked", "forbidden", "waf", "firewall", "security"]
        if any(ind in response.lower() for ind in waf_indicators):
            script = await self.generator.generate_waf_bypass_script(
                target=target,
                blocked_payload=payload,
                waf_response=response,
                vuln_type=attack_type,
                context=context
            )
        else:
            script = await self.generator.generate_exploitation_script(
                target=target,
                vuln_type=attack_type,
                endpoint=endpoint,
                parameter=parameter,
                current_payload=payload,
                response=response,
                context=context
            )

        if script.code and not script.error:
            result = await self.generator.execute_script(script, context)
            return result

        return None

    async def extract_breach_evidence(
        self,
        access_type: str,
        system: str,
        credentials: dict,
        context: dict
    ) -> Optional[ScriptResult]:
        """Generate and run a data extraction script for evidence."""
        logger.info(f"Generating data extraction script for {system}...")

        script = await self.generator.generate_data_extraction_script(
            target=context.get("target", "unknown"),
            access_type=access_type,
            system=system,
            credentials=credentials,
            schema_info=context.get("schema", "unknown"),
            context=context
        )

        if script.code and not script.error:
            return await self.generator.execute_script(script, context)

        return None

    async def create_exploit_chain(
        self,
        vulnerabilities: list[dict],
        goal: str,
        context: dict
    ) -> Optional[ScriptResult]:
        """Create a script that chains vulnerabilities together."""
        logger.info("Generating exploit chain script...")

        script = await self.generator.generate_chain_script(
            vulnerabilities=vulnerabilities,
            goal=goal,
            context=context
        )

        if script.code and not script.error:
            return await self.generator.execute_script(script, context)

        return None

    async def generate_report_pocs(
        self,
        findings: list[dict]
    ) -> list[GeneratedScript]:
        """Generate PoC scripts for all findings (for the report)."""
        pocs = []

        for finding in findings:
            if finding.get("severity") in ["critical", "high"]:
                script = await self.generator.generate_poc_script(
                    vulnerability=finding,
                    context={}
                )
                if script.code and not script.error:
                    pocs.append(script)

        return pocs
