exploit_generator.py

#!/usr/bin/env python3
"""
Enhanced Exploit PoC Auto-Generator for Counterscarp Engine.

Automatically generates Foundry test exploits from vulnerability findings.
Supports local template-based generation and optional LLM enhancement.

Features:
- Pattern-to-template mapping from vulnerability rule IDs to exploit templates
- Contract state inference for proper test setup
- Assertion oracle for meaningful test assertions
- Batch generation with severity filtering
- Output validation via compilation checking
- Multi-LLM support (OpenAI GPT-4, Anthropic Claude)

Environment Variables:
    OPENAI_API_KEY: OpenAI API key for GPT-4 enhancement (optional)
    ANTHROPIC_API_KEY: Anthropic API key for Claude enhancement (optional)

Example:
    >>> from exploit_generator import ExploitGenerator
    >>> generator = ExploitGenerator()
    >>> result = generator.generate_exploit(finding, contract_source)
"""

from __future__ import annotations

import os
import re
import json
import argparse
import subprocess
from typing import Dict, Any, List, Optional, Tuple, cast
from dataclasses import dataclass, field
from pathlib import Path

import logging
from path_security import sanitize_cli_path
from counterscarp_core.severity import SEVERITY_RANK

# Import exceptions (core module — must always be available)
from exceptions import CounterscarpError, CounterscarpConfigError, CounterscarpValidationError

# Import logger with fallback
try:
    from logger import get_logger
    LOGGER_AVAILABLE = True
except ImportError:
    LOGGER_AVAILABLE = False
    def get_logger(name: str) -> logging.Logger:
        return logging.getLogger(name)

# Initialize logger
logger: logging.Logger = get_logger(__name__)

# Optional LLM imports - graceful degradation
try:
    import openai  # noqa: F401
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False

try:
    import anthropic  # noqa: F401
    ANTHROPIC_AVAILABLE = False  # noqa: F841
except ImportError:
    ANTHROPIC_AVAILABLE = False


# ============================================================================
# Pattern-to-Template Mapping
# ============================================================================

# Mapping from vulnerability rule IDs to exploit template files
RULE_TO_TEMPLATE_MAPPING: Dict[str, str] = {
    # Reentrancy patterns
    "REENTRANCY_PATTERN": "reentrancy.sol",
    "FLASH_LOAN_REENTRANCY": "reentrancy.sol",
    "UNCHECKED_EXTERNAL_CALL": "reentrancy.sol",
    
    # Oracle manipulation patterns
    "ORACLE_STALENESS_CHECK": "oracle_manipulation.sol",
    "PRICE_STALENESS_PATTERN": "oracle_manipulation.sol",
    "MISSING_SLIPPAGE_PROTECTION": "oracle_manipulation.sol",
    
    # Access control patterns
    "MISSING_ACCESS_CONTROL": "access_control.sol",
    "EMERGENCY_WITHDRAW_PUBLIC": "access_control.sol",
    "UPGRADE_FUNCTION": "access_control.sol",
    "FAKE_RENOUNCE_OWNER_ZERO": "access_control.sol",
    "ARBITRARY_EXTERNAL_CALL": "access_control.sol",
    
    # Integer overflow/precision patterns
    "UNSAFE_CAST": "integer_overflow.sol",
    "DIVIDE_BEFORE_MULTIPLY": "integer_overflow.sol",
    "STRICT_BALANCE_EQUALITY": "integer_overflow.sol",
    
    # Front-running patterns
    "TX_ORIGIN_USAGE": "front_running.sol",
    
    # Flash loan patterns
    "UNSAFE_ERC20_TRANSFER": "flash_loan.sol",
    "HIDDEN_MINT": "flash_loan.sol",
    "MSG_VALUE_LOOP": "flash_loan.sol",
    
    # Delegate call patterns
    "DELEGATE_CALL_INJECTION": "delegate_call.sol",
    "ARBITRARY_DELEGATECALL": "delegate_call.sol",
    "DELEGATECALL_TO_UNTRUSTED": "delegate_call.sol",
    "DELEGATE_CALL": "delegate_call.sol",
    "DELEGATE-*": "delegate_call.sol",
    
    # Proxy/upgrade patterns
    "PROXY_STORAGE_COLLISION": "storage_collision.sol",
    "UPGRADEABLE_PROXY_VULNERABILITY": "storage_collision.sol",
    "STORAGE_COLLISION": "storage_collision.sol",
    "UPGRADE-*": "storage_collision.sol",
    "PROXY-*": "storage_collision.sol",

    # Slither detector names — direct mappings to prevent "No template found" warnings
    "DELEGATECALL_USAGE": "delegate_call.sol",
    "controlled-delegatecall": "delegate_call.sol",
    "incorrect-return": "access_control.sol",
    "reentrancy-eth": "reentrancy.sol",
    "reentrancy-no-eth": "reentrancy.sol",
    "reentrancy-benign": "reentrancy.sol",
    "reentrancy-events": "reentrancy.sol",
    "unchecked-transfer": "flash_loan.sol",
    "arbitrary-send-eth": "access_control.sol",
    "suicidal": "access_control.sol",
    "uninitialized-state": "storage_collision.sol",
    "uninitialized-local": "storage_collision.sol",
    "tx-origin": "front_running.sol",
    "weak-prng": "front_running.sol",
    "divide-before-multiply": "integer_overflow.sol",
    "tautology": "integer_overflow.sol",
    "msg-value-loop": "flash_loan.sol",
    "locked-ether": "access_control.sol",
    "shadowing-state": "storage_collision.sol",
}

# Severity order for sorting — imported from canonical source
SEVERITY_ORDER = SEVERITY_RANK


# ============================================================================
# Data Classes
# ============================================================================

@dataclass
class ExploitResult:
    """Result of a single exploit generation.

    Attributes:
        finding: The vulnerability finding that was processed.
        output_path: Path to the generated exploit file.
        status: Generation status (success, failed, skipped).
        message: Additional information or error message.
    """
    finding: Dict[str, Any]
    output_path: Optional[str]
    status: str
    message: str


@dataclass
class ValidationResult:
    """Result of exploit validation.
    
    Attributes:
        valid: Whether the exploit is valid/compiles.
        errors: List of compilation errors if any.
    """
    valid: bool
    errors: List[str] = field(default_factory=list)


@dataclass
class ContractState:
    """Inferred contract state for exploit setup.
    
    Attributes:
        tokens: List of ERC20 tokens needed.
        nfts: List of ERC721 tokens needed.
        requires_liquidity: Whether liquidity provision is needed.
        oracle_mock: Whether oracle mock setup is needed.
        roles: List of role/access assignments needed.
        constructor_params: Constructor parameters.
        initial_balances: Initial token balances to set up.
    """
    tokens: List[Dict[str, Any]] = field(default_factory=list)
    nfts: List[Dict[str, Any]] = field(default_factory=list)
    requires_liquidity: bool = False
    oracle_mock: bool = False
    roles: List[Tuple[str, str]] = field(default_factory=list)  # (role, account)
    constructor_params: List[str] = field(default_factory=list)
    initial_balances: Dict[str, int] = field(default_factory=dict)


# ============================================================================
# Exploit Generator Class
# ============================================================================

class ExploitGenerator:
    """Enhanced exploit PoC generator for Counterscarp Engine.
    
    This class provides comprehensive exploit generation capabilities including
n    template-based generation, contract state inference, assertion generation,
    batch processing, and LLM enhancement.
    
    Attributes:
        config: Configuration dictionary for exploit generation.
        template_dir: Directory containing exploit templates.
        output_dir: Directory for generated exploits.
        llm_backend: LLM backend to use ("none", "openai", "anthropic").
    """
    
    def __init__(
        self,
        config: Optional[Dict[str, Any]] = None,
        template_dir: str = "exploit_templates/",
        output_dir: str = "exploits/",
        llm_backend: str = "none"
    ):
        """Initialize the exploit generator.
        
        Args:
            config: Optional configuration dictionary.
            template_dir: Directory containing exploit templates.
            output_dir: Directory for generated exploits.
            llm_backend: LLM backend to use (none, openai, anthropic).
        """
        self.config = config or {}
        self.template_dir = Path(template_dir)
        # Use a fixed safe output root to avoid arbitrary write paths.
        self.output_dir = (Path.cwd() / "exploits").resolve()
        self.llm_backend = llm_backend.lower()
        
        # Ensure output directory exists
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Check API keys
        self.openai_key = os.getenv("OPENAI_API_KEY")
        self.anthropic_key = os.getenv("ANTHROPIC_API_KEY")
        
        logger.info(f"ExploitGenerator initialized with backend: {self.llm_backend}")
    
    def _get_template_for_rule(self, rule_id: str) -> Optional[str]:
        """Get the template file for a given rule ID.
        
        Args:
            rule_id: The vulnerability rule ID.
            
        Returns:
            Path to template file or None if no mapping exists.
        """
        template_file = RULE_TO_TEMPLATE_MAPPING.get(rule_id)
        if not template_file:
            # Try partial matching
            for rule_pattern, template in RULE_TO_TEMPLATE_MAPPING.items():
                if rule_pattern in rule_id or rule_id in rule_pattern:
                    template_file = template
                    break  # noqa: WPS220
        
        if template_file:
            template_path = self.template_dir / template_file
            if template_path.exists():
                return str(template_path)
        
        return None
    
    def _load_template(self, template_path: str) -> str:
        """Load a template file.
        
        Args:
            template_path: Path to the template file.
            
        Returns:
            Template content as string.
            
        Raises:
            FileNotFoundError: If template file doesn't exist.
        """
        with open(template_path, 'r', encoding='utf-8') as f:
            return f.read()
    
    def infer_contract_state(self, contract_source: str, finding: Dict[str, Any]) -> ContractState:
        """Analyze the vulnerable contract to determine required state for exploit.
        
        Args:
            contract_source: Source code of the vulnerable contract.
            finding: The vulnerability finding dictionary.
            
        Returns:
            ContractState object with inferred requirements.
        """
        state = ContractState()
        
        if not contract_source:
            return state
        
        # Detect ERC20 tokens
        if re.search(r'import.*ERC20|IERC20|transfer\(|transferFrom\(', contract_source):
            state.tokens.append({
                "type": "ERC20",
                "name": "TestToken",
                "symbol": "TEST",
                "decimals": 18
            })
        
        # Detect ERC721 tokens
        if re.search(r'import.*ERC721|IERC721|safeTransferFrom.*uint256', contract_source):
            state.nfts.append({
                "type": "ERC721",
                "name": "TestNFT",
                "symbol": "TNFT"
            })
        
        # Detect oracle usage
        if re.search(r'oracle|price|latestAnswer|latestRoundData', contract_source, re.IGNORECASE):
            state.oracle_mock = True
        
        # Detect access control
        if re.search(r'onlyOwner|onlyRole|AccessControl| Ownable', contract_source):
            state.roles.append(("ADMIN", "owner"))
        
        # Detect liquidity requirements
        if re.search(r'liquidity|pool|pair|swap|DEX|AMM', contract_source, re.IGNORECASE):
            state.requires_liquidity = True
        
        # Detect constructor parameters
        ctor_match = re.search(r'constructor\((.*?)\)', contract_source, re.DOTALL)
        if ctor_match:
            params_str = ctor_match.group(1)
            # Extract parameter names
            params = re.findall(r'\w+\s+(\w+)(?:,|$)', params_str)
            state.constructor_params = params
        
        # Detect initial balance requirements
        if finding.get('rule_id') in ['REENTRANCY_PATTERN', 'FLASH_LOAN_REENTRANCY', 'EMERGENCY_WITHDRAW_PUBLIC']:
            state.initial_balances['eth'] = 100  # 100 ETH
        
        logger.debug(f"Inferred contract state: {state}")
        return state
    
    def generate_assertions(
        self,
        finding: Dict[str, Any],
        state: ContractState
    ) -> List[str]:
        """Generate meaningful test assertions for the exploit.
        
        Args:
            finding: The vulnerability finding dictionary.
            state: The inferred contract state.
            
        Returns:
            List of assertion code strings.
        """
        assertions = []
        rule_id = finding.get('rule_id', '')
        
        # Balance change assertions
        if rule_id in ['REENTRANCY_PATTERN', 'FLASH_LOAN_REENTRANCY', 'EMERGENCY_WITHDRAW_PUBLIC']:
            assertions.append("// Balance change assertions")
            assertions.append("assertGt(attacker.balance, initialAttackerBalance, \"Attacker should gain ETH\");")
            assertions.append("assertLt(address(vulnerable).balance, initialContractBalance, \"Contract should lose ETH\");")
        
        if rule_id in ['HIDDEN_MINT', 'MISSING_ACCESS_CONTROL']:
            assertions.append("// Token balance assertions")
            assertions.append("assertGt(token.balanceOf(attacker), initialAttackerTokenBalance, \"Attacker should gain tokens\");")
        
        # State transition checks
        if rule_id in ['ORACLE_STALENESS_CHECK', 'PRICE_STALENESS_PATTERN']:
            assertions.append("// Price/oracle assertions")
            assertions.append("assertEq(vulnerable.getPrice(), manipulatedPrice, \"Price should be manipulated\");")
        
        # Event emission checks
        if rule_id in ['UNSAFE_ERC20_TRANSFER', 'UNCHECKED_EXTERNAL_CALL']:
            assertions.append("// Event emission checks")
            assertions.append("vm.expectEmit(true, true, false, true);")
            assertions.append("emit Transfer(address(vulnerable), attacker, amount);")
        
        # Revert expectations
        if rule_id in ['UNSAFE_CAST', 'DIVIDE_BEFORE_MULTIPLY']:
            assertions.append("// Expect arithmetic issues")
            assertions.append("// vm.expectRevert(); // Uncomment if overflow protection exists")
        
        # Access control assertions
        if rule_id in ['MISSING_ACCESS_CONTROL', 'TX_ORIGIN_USAGE']:
            assertions.append("// Access control assertions")
            assertions.append("assertEq(vulnerable.owner(), attacker, \"Attacker should gain ownership\");")
        
        logger.debug(f"Generated {len(assertions)} assertions for {rule_id}")
        return assertions
    
    def _fill_template(
        self,
        template: str,
        finding: Dict[str, Any],
        state: ContractState,
        contract_source: str
    ) -> str:
        """Fill template placeholders with actual values.
        
        Args:
            template: The template content.
            finding: The vulnerability finding.
            state: The inferred contract state.
            contract_source: The vulnerable contract source.
            
        Returns:
            Filled template content.
        """
        # Extract contract name from source or use default
        contract_name = "VulnerableContract"
        contract_match = re.search(r'contract\s+(\w+)', contract_source)
        if contract_match:
            contract_name = contract_match.group(1)
        
        # Generate assertions
        assertions = self.generate_assertions(finding, state)
        assertions_str = '\n        '.join(assertions)
        
        # Replace placeholders
        filled = template
        
        # Replace PLACEHOLDER comments with actual code
        filled = re.sub(
            r'// PLACEHOLDER:.*\n\s*// import "path/to/VulnerableContract.sol";',
            f'import "../src/{contract_name}.sol";',
            filled
        )
        
        filled = re.sub(
            r'// PLACEHOLDER:.*\n\s*// (\w+) public vulnerable;',
            f'{contract_name} public vulnerable;',
            filled
        )
        
        filled = re.sub(
            r'// PLACEHOLDER: Deploy vulnerable contract\n\s*// vulnerable = new \w+\(\);',
            f'vulnerable = new {contract_name}();',
            filled
        )
        
        # Add assertions section
        filled = re.sub(
            r'// PLACEHOLDER: Assertions.*\n(\s*)// assert.*\n',
            f'// Assertions\n\\1{assertions_str}\n',
            filled
        )
        
        # --- Educational header + contract-specific context ---
        # Extract contract name from the finding's file path (e.g. "SedaCoreV1.sol" → "SedaCoreV1")
        finding_file: str = finding.get('file', 'Unknown')
        finding_contract_name = Path(finding_file).stem if finding_file != 'Unknown' else 'Unknown'

        # Try to extract function name from line_text / code_snippet via regex
        code_context: str = finding.get('line_text', finding.get('code_snippet', '')) or ''
        func_match = re.search(r'function\s+(\w+)', code_context)
        function_name: str = func_match.group(1) if func_match else 'N/A'

        # Truncate code snippet to 200 chars
        code_snippet_display = (code_context[:200] + '...') if len(code_context) > 200 else code_context
        # Escape */ to avoid premature comment termination
        code_snippet_display = code_snippet_display.replace('*/', '* /')

        educational_header = (
            "// ============================================================================\n"
            "// EDUCATIONAL TEMPLATE — Demonstrates the vulnerability pattern for this finding.\n"
            "// This is NOT a direct proof-of-concept against the flagged contract.\n"
            "// To create a targeted PoC, replace the synthetic contract below with imports\n"
            "// and interfaces matching the actual vulnerable contract.\n"
            "// ============================================================================\n"
            "\n"
            f"// TARGET CONTEXT:\n"
            f"// Contract: {finding_contract_name}\n"
            f"// File: {finding_file}\n"
            f"// Line: {finding.get('line_no', '?')}\n"
            f"// Function: {function_name}\n"
            f"// Code: {code_snippet_display}\n"
            f"// Rule: {finding.get('rule_id', 'UNKNOWN')} ({finding.get('severity', 'UNKNOWN')})\n"
            "\n"
        )

        # Legacy finding comment (kept for downstream consumers)
        finding_comment = (
            f"\n"
            f"/*\n"
            f" * Vulnerability: {finding.get('rule_id', 'UNKNOWN')}\n"
            f" * Severity: {finding.get('severity', 'UNKNOWN')}\n"
            f" * Location: {finding_file}:{finding.get('line_no', '?')}\n"
            f" * Description: {finding.get('message', finding.get('description', 'No description'))}\n"
            f" */\n"
        )
        filled = educational_header + finding_comment + filled

        return filled
    
    def _enhance_with_llm(
        self,
        exploit_code: str,
        finding: Dict[str, Any],
        contract_source: str
    ) -> str:
        """Enhance exploit code using LLM if configured.
        
        Args:
            exploit_code: The base exploit code.
            finding: The vulnerability finding.
            contract_source: The vulnerable contract source.
            
        Returns:
            Enhanced exploit code.
        """
        if self.llm_backend == "none":
            return exploit_code
        
        if self.llm_backend == "openai" and OPENAI_AVAILABLE and self.openai_key:
            return self._enhance_with_openai(exploit_code, finding, contract_source)
        
        if self.llm_backend == "anthropic" and ANTHROPIC_AVAILABLE and self.anthropic_key:
            return self._enhance_with_anthropic(exploit_code, finding, contract_source)
        
        logger.warning(f"LLM backend '{self.llm_backend}' not available, using template only")
        return exploit_code
    
    def _enhance_with_openai(
        self,
        exploit_code: str,
        finding: Dict[str, Any],
        contract_source: str
    ) -> str:
        """Enhance exploit using OpenAI GPT-4.
        
        Args:
            exploit_code: The base exploit code.
            finding: The vulnerability finding.
            contract_source: The vulnerable contract source.
            
        Returns:
            Enhanced exploit code.
        """
        try:
            client = openai.OpenAI(api_key=self.openai_key)
            
            prompt = f"""Enhance this Foundry test exploit for a smart contract vulnerability.

Vulnerability:
- Rule ID: {finding.get('rule_id')}
- Severity: {finding.get('severity')}
- Description: {finding.get('message', finding.get('description'))}
- Location: {finding.get('file')}:{finding.get('line_no')}

Current exploit code:
```solidity
{exploit_code}
```

Task:
1. Improve the exploit to be more realistic and effective
2. Add specific attack logic based on the vulnerability type
3. Ensure proper Foundry test patterns (vm.prank, vm.deal, etc.)
4. Add comprehensive assertions that prove the exploit works
5. Keep the same overall structure but enhance the implementation

Return ONLY the enhanced Solidity code, no explanations."""

            response = client.chat.completions.create(
                model="gpt-4-turbo-preview",
                messages=[
                    {
                        "role": "system",
                        "content": "You are an expert Solidity security researcher specializing in proof-of-concept exploits."
                    },
                    {"role": "user", "content": prompt}
                ],
                temperature=0.3,
                max_tokens=2000
            )
            
            enhanced: str = cast(str, response.choices[0].message.content)
            # Clean up markdown if present
            if "```solidity" in enhanced:
                enhanced = enhanced.split("```solidity")[1].split("```")[0]
            elif "```" in enhanced:
                enhanced = enhanced.split("```")[1].split("```")[0]
            
            logger.info("Exploit enhanced with OpenAI GPT-4")
            return enhanced.strip()
            
        except Exception as e:
            logger.error(f"OpenAI enhancement failed: {e}")
            return exploit_code
    
    def _enhance_with_anthropic(
        self,
        exploit_code: str,
        finding: Dict[str, Any],
        contract_source: str
    ) -> str:
        """Enhance exploit using Anthropic Claude.
        
        Args:
            exploit_code: The base exploit code.
            finding: The vulnerability finding.
            contract_source: The vulnerable contract source.
            
        Returns:
            Enhanced exploit code.
        """
        try:
            client = anthropic.Anthropic(api_key=self.anthropic_key)
            
            prompt = f"""Enhance this Foundry test exploit for a smart contract vulnerability.

Vulnerability:
- Rule ID: {finding.get('rule_id')}
- Severity: {finding.get('severity')}
- Description: {finding.get('message', finding.get('description'))}
- Location: {finding.get('file')}:{finding.get('line_no')}

Current exploit code:
```solidity
{exploit_code}
```

Task:
1. Improve the exploit to be more realistic and effective
2. Add specific attack logic based on the vulnerability type
3. Ensure proper Foundry test patterns (vm.prank, vm.deal, etc.)
4. Add comprehensive assertions that prove the exploit works
5. Keep the same overall structure but enhance the implementation

Return ONLY the enhanced Solidity code, no explanations."""

            response = client.messages.create(
                model="claude-3-opus-20240229",
                max_tokens=2000,
                temperature=0.3,
                messages=[{"role": "user", "content": prompt}]
            )
            
            first_block = response.content[0]
            enhanced: str = cast(str, first_block.text if hasattr(first_block, "text") else "")
            # Clean up markdown if present
            if "```solidity" in enhanced:
                enhanced = enhanced.split("```solidity")[1].split("```")[0]
            elif "```" in enhanced:
                enhanced = enhanced.split("```")[1].split("```")[0]
            
            logger.info("Exploit enhanced with Anthropic Claude")
            return enhanced.strip()
            
        except Exception as e:
            logger.error(f"Anthropic enhancement failed: {e}")
            return exploit_code
    
    def generate_exploit(
        self,
        finding: Dict[str, Any],
        contract_source: str = "",
        output_path: Optional[str] = None
    ) -> str:
        """Generate a single exploit PoC for a vulnerability finding.
        
        This is the main entry point for exploit generation. It maintains
        backward compatibility with the existing generate_exploit() function.
        
        Args:
            finding: The vulnerability finding dictionary.
            contract_source: Optional source code of the vulnerable contract.
            output_path: Optional custom output path for the exploit file.
            
        Returns:
            Path to the generated exploit file.
            
        Raises:
            CounterscarpValidationError: If finding data is invalid.
        """
        if not finding or not isinstance(finding, dict):
            raise CounterscarpValidationError(
                "Invalid finding data",
                details={"finding": finding}
            )
        
        rule_id = finding.get('rule_id', 'UNKNOWN')
        logger.info(f"Generating exploit for: {rule_id}")
        
        # Get template for rule
        template_path = self._get_template_for_rule(rule_id)
        if not template_path:
            logger.warning(f"No template found for rule: {rule_id}, using generic template")
            template_path = self._get_template_for_rule('REENTRANCY_PATTERN')
            if not template_path:
                raise CounterscarpValidationError(
                    f"No template available for rule: {rule_id}",
                    details={"rule_id": rule_id}
                )
        
        # Load template
        template = self._load_template(template_path)
        
        # Infer contract state
        state = self.infer_contract_state(contract_source, finding)
        
        # Fill template
        exploit_code = self._fill_template(template, finding, state, contract_source)
        
        # Enhance with LLM if configured
        if self.llm_backend != "none":
            exploit_code = self._enhance_with_llm(exploit_code, finding, contract_source)
        
        # Determine output path
        if output_path:
            safe_name = os.path.basename(output_path)
            output_path = str(self.output_dir / safe_name)
        else:
            filename = f"Exploit_{rule_id}.t.sol"
            output_path = str(self.output_dir / filename)
        
        # Write exploit file — use Path.parent so an empty dirname never raises
        out_parent = os.path.dirname(output_path)
        os.makedirs(out_parent if out_parent else ".", exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(exploit_code)
        
        logger.info(f"Exploit generated: {output_path}")
        return output_path
    
    def batch_generate(
        self,
        findings: List[Dict[str, Any]],
        contract_source: str = "",
        output_dir: Optional[str] = None,
        min_severity: str = "HIGH"
    ) -> List[ExploitResult]:
        """Generate exploits for multiple findings.
        
        Args:
            findings: List of vulnerability findings.
            contract_source: Source code of the vulnerable contract.
            output_dir: Directory to save exploits (uses instance default if None).
            min_severity: Minimum severity level to generate exploits for.
            
        Returns:
            List of ExploitResult objects.
        """
        if output_dir:
            safe_subdir = os.path.basename(output_dir.rstrip("/\\"))
            output_path = os.path.join(str(self.output_dir), safe_subdir)
            os.makedirs(output_path, exist_ok=True)
        else:
            output_path = str(self.output_dir)
        
        # Filter by minimum severity
        min_severity_level = SEVERITY_ORDER.get(min_severity.upper(), 1)
        filtered_findings = [
            f for f in findings
            if SEVERITY_ORDER.get(f.get('severity', 'INFO').upper(), 5) <= min_severity_level
        ]
        
        logger.info(f"Filtered {len(findings)} findings to {len(filtered_findings)} with severity >= {min_severity}")
        
        # Sort by severity (CRITICAL first)
        sorted_findings = sorted(
            filtered_findings,
            key=lambda f: SEVERITY_ORDER.get(f.get('severity', 'INFO').upper(), 5)
        )
        
        results = []
        rule_counts: Dict[str, int] = {}
        for finding in sorted_findings:
            rule_id = finding.get('rule_id', 'UNKNOWN')
            try:
                # Build unique filename to avoid overwriting when the same rule fires multiple times
                rule_counts[rule_id] = rule_counts.get(rule_id, 0) + 1
                count = rule_counts[rule_id]
                if count == 1:
                    filename = f"Exploit_{rule_id}.t.sol"
                else:
                    filename = f"Exploit_{rule_id}_{count}.t.sol"
                output_file = os.path.join(output_path, filename)
                exploit_path = self.generate_exploit(
                    finding,
                    contract_source,
                    output_file
                )
                results.append(ExploitResult(
                    finding=finding,
                    output_path=exploit_path,
                    status="success",
                    message="Exploit generated successfully"
                ))
            except Exception as e:
                logger.error(f"Failed to generate exploit for {rule_id}: {e}")
                results.append(ExploitResult(
                    finding=finding,
                    output_path=None,
                    status="failed",
                    message=str(e)
                ))
        
        logger.info(f"Batch generation complete: {len([r for r in results if r.status == 'success'])} succeeded")
        return results
    
    def validate_exploit(self, exploit_path: str) -> ValidationResult:
        """Validate that a generated exploit compiles.
        
        Args:
            exploit_path: Path to the exploit file to validate.
            
        Returns:
            ValidationResult with validation status and errors.
        """
        if not Path(exploit_path).exists():
            return ValidationResult(
                valid=False,
                errors=[f"Exploit file not found: {exploit_path}"]
            )
        
        # Security: validate exploit_path resolves within its own parent dir
        # (guards against path traversal via symlinks or .. sequences)
        try:
            _resolved_exploit = Path(exploit_path).resolve()
            _resolved_exploit.relative_to(_resolved_exploit.parent)
        except ValueError:
            logger.warning(
                f"[SECURITY] Path traversal rejected for exploit validation:"
                f" {exploit_path!r} escapes expected directory"
            )
            return ValidationResult(
                valid=False,
                errors=["Path traversal detected: exploit path escapes output directory"]
            )

        # Check if forge is available
        try:
            result = subprocess.run(
                ["forge", "build", "--contracts", exploit_path],
                capture_output=True,
                text=True,
                timeout=60
            )
            
            if result.returncode == 0:
                logger.info(f"Exploit validation passed: {exploit_path}")
                return ValidationResult(valid=True)
            else:
                errors = result.stderr.split('\n') if result.stderr else ["Compilation failed"]
                logger.warning(f"Exploit validation failed: {exploit_path}")
                return ValidationResult(valid=False, errors=errors)
                
        except FileNotFoundError:
            logger.warning("Forge not found in PATH, skipping compilation validation")
            return ValidationResult(
                valid=True,
                errors=["Forge not available - validation skipped"]
            )
        except subprocess.TimeoutExpired:
            return ValidationResult(
                valid=False,
                errors=["Compilation timed out after 60 seconds"]
            )
        except Exception as e:
            return ValidationResult(
                valid=False,
                errors=[f"Validation error: {str(e)}"]
            )


# ============================================================================
# Legacy Functions for Backward Compatibility
# ============================================================================

def generate_exploit_prompt(vulnerability: Dict[str, Any]) -> str:
    """Create a detailed prompt for GPT-4 to generate exploit code.
    
    Args:
        vulnerability: Dict with vulnerability details.
        
    Returns:
        Formatted prompt string for GPT-4.
    """
    return f"""You are a smart contract security researcher writing a Foundry test to exploit a vulnerability.

**Vulnerability Details:**
- ID: {vulnerability.get('rule_id', 'UNKNOWN')}
- Severity: {vulnerability.get('severity', 'UNKNOWN')}
- Description: {vulnerability.get('message', vulnerability.get('description', 'No description'))}
- Location: {vulnerability.get('file', 'Unknown')}:{vulnerability.get('line_no', '?')}
- Vulnerable Code:
```solidity
{vulnerability.get('line_text', 'Not provided')}
```

**Task:**
Generate a complete Foundry test file that demonstrates the exploit.

**Output Format:**
Return ONLY the Solidity code (no markdown formatting).
"""


def generate_exploit_from_finding(
    vulnerability: Dict[str, Any],
    output_dir: str = "test/exploits",
    model: str = "gpt-4-turbo-preview"
) -> Optional[str]:
    """Legacy function: Generate exploit test from vulnerability finding.
    
    Maintains backward compatibility with existing code.
    
    Args:
        vulnerability: Vulnerability data.
        output_dir: Where to save the test.
        model: GPT-4 model to use (legacy parameter, ignored).
        
    Returns:
        Path to generated test file, or None if failed.
    """
    generator = ExploitGenerator(output_dir=output_dir)
    try:
        return generator.generate_exploit(vulnerability, output_path=output_dir)
    except Exception as e:
        logger.error(f"Failed to generate exploit: {e}")
        return None


def generate_batch_exploits(
    findings: List[Dict[str, Any]],
    output_dir: str = "test/exploits",
    max_exploits: int = 5
) -> List[str]:
    """Legacy function: Generate exploits for multiple findings.
    
    Maintains backward compatibility with existing code.
    
    Args:
        findings: List of vulnerability findings.
        output_dir: Where to save tests.
        max_exploits: Maximum number to generate.
        
    Returns:
        List of generated file paths.
    """
    generator = ExploitGenerator(output_dir=output_dir)
    results = generator.batch_generate(
        findings[:max_exploits],
        output_dir=output_dir
    )
    return [r.output_path for r in results if r.output_path]


# ============================================================================
# CLI Entry Point
# ============================================================================

def main() -> None:
    """Main entry point for the exploit generator CLI."""
    parser = argparse.ArgumentParser(
        description="🤖 Enhanced Exploit PoC Auto-Generator for Counterscarp Engine"
    )
    
    parser.add_argument(
        "--finding-json",
        help="Path to JSON file with vulnerability finding"
    )
    parser.add_argument(
        "--rule-id",
        help="Vulnerability rule ID (e.g., UNCHECKED_EXTERNAL_CALL)"
    )
    parser.add_argument(
        "--description",
        help="Vulnerability description"
    )
    parser.add_argument(
        "--file",
        help="Vulnerable file path"
    )
    parser.add_argument(
        "--line",
        type=int,
        help="Vulnerable line number"
    )
    parser.add_argument(
        "--contract-source",
        type=argparse.FileType("r", encoding="utf-8"),
        help="Path to vulnerable contract source file"
    )
    parser.add_argument(
        "--output-dir",
        default="exploits/",
        help="Output directory for exploit files"
    )
    parser.add_argument(
        "--llm-backend",
        choices=["none", "openai", "anthropic"],
        default="none",
        help="LLM backend for enhancement (default: none)"
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        help="Validate generated exploit compiles"
    )
    parser.add_argument(
        "--batch",
        type=argparse.FileType("r", encoding="utf-8"),
        help="Path to JSON file with multiple findings for batch generation"
    )
    parser.add_argument(
        "--min-severity",
        default="HIGH",
        choices=["CRITICAL", "HIGH", "MEDIUM", "LOW"],
        help="Minimum severity for batch generation"
    )
    
    args = parser.parse_args()
    
    # Initialize generator
    generator = ExploitGenerator(
        output_dir=args.output_dir,
        llm_backend=args.llm_backend
    )
    
    # Batch generation mode
    if args.batch:
        findings = json.load(args.batch)
        
        results = generator.batch_generate(
            findings,
            output_dir=args.output_dir,
            min_severity=args.min_severity
        )
        
        print(f"\n{'='*60}")
        print(f" BATCH GENERATION RESULTS")
        print(f"{'='*60}")
        for result in results:
            status_icon = "✅" if result.status == "success" else "❌"
            print(f"{status_icon} {result.finding.get('rule_id')}: {result.status}")
            if result.output_path:
                print(f"   Output: {result.output_path}")
        return
    
    # Single finding mode
    if args.finding_json:
        safe_finding = sanitize_cli_path(args.finding_json, allowed_suffixes={".json"})
        vulnerability = json.loads(safe_finding.read_text(encoding="utf-8"))
    else:
        vulnerability = {
            "rule_id": args.rule_id or "UNKNOWN_VULNERABILITY",
            "severity": "HIGH",
            "description": args.description or "No description provided",
            "file": args.file or "Unknown.sol",
            "line_no": args.line or 0,
            "line_text": "// Code not provided"
        }
    
    # Load contract source if provided
    contract_source = ""
    if args.contract_source:
        sanitize_cli_path(args.contract_source.name)
        contract_source = args.contract_source.read()
    
    # Generate exploit
    try:
        exploit_path = generator.generate_exploit(
            vulnerability,
            contract_source=contract_source
        )
        print(f"\n[+] ✅ Exploit generated: {exploit_path}")
        
        # Validate if requested
        if args.validate:
            validation = generator.validate_exploit(exploit_path)
            if validation.valid:
                print("[+] ✅ Exploit validation passed")
            else:
                print("[!] ❌ Exploit validation failed:")
                for error in validation.errors:
                    print(f"    - {error}")
        
        print(f"[*] Run with: forge test --match-test testExploit")
        
    except Exception as e:
        print(f"[!] Failed to generate exploit: {e}")
        raise


if __name__ == "__main__":
    main()