"""
Heap overflow finder module for pwnelle.

Automatically detects heap overflows by analyzing adjacent allocations.
"""

import time
import os
import re
from typing import Dict, List, Any, Tuple, Optional, Union
from pwn import *

class HeapOverflowFinder:
    """Find heap overflow vulnerabilities across adjacent allocations."""
    
    def __init__(
        self,
        binary_path: str,
        analysis_results: Dict[str, Any],
        max_attempts: int = 20,
        max_length: int = 200,
        timeout: int = 5,
        verbose: bool = False
    ):
        """
        Initialize the heap overflow finder.
        
        Args:
            binary_path: Path to the binary to analyze
            analysis_results: Analysis results from the BinaryAnalyzer
            max_attempts: Maximum number of allocation patterns to try
            max_length: Maximum length to try for payloads
            timeout: Timeout for process in seconds
            verbose: Whether to show verbose output
        """
        self.binary_path = binary_path
        self.analysis_results = analysis_results
        self.max_attempts = max_attempts
        self.max_length = max_length
        self.timeout = timeout
        self.verbose = verbose
        self.context = context
        
        # Initialize pwntools ELF
        try:
            self.elf = ELF(binary_path)
            if self.verbose:
                self.context.log_level = 'debug'
            else:
                self.context.log_level = 'warning'
        except Exception as e:
            print(f"Error loading binary: {e}")
            raise
    
    def log(self, message: str) -> None:
        """Print message if verbose mode is enabled."""
        if self.verbose:
            print(message)
    
    def find_overflows(
        self, 
        interaction_patterns: List[Dict[str, Any]] = None
    ) -> Optional[Dict[str, Any]]:
        """
        Find potential heap overflow vulnerabilities.
        
        Args:
            interaction_patterns: List of interaction patterns to try
                                 Each pattern is a dict with keys:
                                 - 'name': Pattern name
                                 - 'setup': List of (expect, send) tuples for initial setup
                                 - 'overflow_action': Dict with keys:
                                   - 'expect': Bytes to expect before sending overflow
                                   - 'send_template': Payload template with '{length}' placeholder
                                 - 'verify_action': Dict with keys:
                                   - 'expect': Bytes to expect before sending verification
                                   - 'send': Bytes to send for verification
                                 - 'success_indicator': Bytes indicating successful overflow
                                 - 'fail_indicator': Bytes indicating failed overflow
        
        Returns:
            Dictionary with heap overflow info or None if not found
        """
        if interaction_patterns is None:
            # Default pattern based on common menu-driven applications
            interaction_patterns = [self._get_default_pattern()]
        
        self.log(f"Finding heap overflows in {self.binary_path}")
        
        for pattern in interaction_patterns:
            self.log(f"Trying pattern: {pattern['name']}")
            
            # Try different overflow lengths
            for length in range(8, self.max_length + 1, 4):
                try:
                    self.log(f"Trying length {length}...")
                    result = self._try_overflow(pattern, length)
                    if result:
                        return result
                except Exception as e:
                    self.log(f"Error in testing heap overflow: {e}")
                    continue
        
        return None
    
    def _get_default_pattern(self) -> Dict[str, Any]:
        """
        Generate a default interaction pattern for typical menu-driven heap applications.
        
        Returns:
            Dictionary with default pattern
        """
        return {
            'name': 'default_menu',
            'setup': [(b'choice:', b'1')],  # Typically option 1 is print heap
            'overflow_action': {
                'expect': b'buffer',
                'send_template': b'A' * '{length}'
            },
            'verify_action': {
                'expect': b'choice:',
                'send': b'1'  # Print heap again to verify
            },
            'success_indicator': b'AAA',  # Adjacent allocation shows our As
            'fail_indicator': None  # No specific failure indicator
        }
    
    def _try_overflow(self, pattern: Dict[str, Any], length: int) -> Optional[Dict[str, Any]]:
        """
        Try a heap overflow with the given interaction pattern and payload length.
        
        Args:
            pattern: Interaction pattern to use
            length: Length of overflow payload
            
        Returns:
            Dictionary with results if overflow found, None otherwise
        """
        try:
            self.log(f"Starting process with length {length}")
            p = process(self.binary_path, timeout=self.timeout)
            
            # Run setup interactions
            for expect, send in pattern['setup']:
                p.recvuntil(expect)
                p.sendline(send)
                time.sleep(0.1)
            
            # Run overflow action
            overflow_action = pattern['overflow_action']
            p.recvuntil(overflow_action['expect'])
            
            # Format the payload with the current length
            payload = overflow_action['send_template'].replace(b'{length}', str(length).encode())
            p.sendline(payload)
            time.sleep(0.1)
            
            # Run verify action
            verify_action = pattern['verify_action']
            p.recvuntil(verify_action['expect'])
            p.sendline(verify_action['send'])
            
            # Check for success indicator
            try:
                response = p.recvuntil(b'choice:', timeout=2)
                if pattern['success_indicator'] in response:
                    p.close()
                    return {
                        'found': True,
                        'type': 'heap_overflow',
                        'length': length,
                        'pattern': pattern['name'],
                        'confidence': 'high'
                    }
                
                # If fail indicator is specified, check for it
                if pattern['fail_indicator'] and pattern['fail_indicator'] in response:
                    self.log(f"Fail indicator found, trying next length")
                    p.close()
                    return None
            except EOFError:
                # Process crashed, might be good or bad depending on the binary
                self.log("Process crashed or exited")
                p.close()
                return {
                    'found': True,
                    'type': 'heap_crash',
                    'length': length,
                    'pattern': pattern['name'],
                    'confidence': 'medium'
                }
                
            # No clear success or failure indicated, continue to next length
            p.close()
            return None
            
        except Exception as e:
            self.log(f"Error during heap overflow test: {e}")
            if 'p' in locals() and p.poll() is None:
                p.close()
            return None
    
    def find_heap_functions(self) -> Dict[str, int]:
        """
        Find heap-related functions in the binary.
        
        Returns:
            Dictionary mapping function names to addresses
        """
        heap_funcs = {}
        try:
            for name, addr in self.elf.plt.items():
                if name in ['malloc', 'free', 'realloc', 'calloc']:
                    heap_funcs[name] = addr
        except Exception as e:
            self.log(f"Error finding heap functions: {e}")
        
        return heap_funcs
    
    def create_exploit(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """
        Create an exploit based on the heap overflow findings.
        
        Args:
            result: Result from find_overflows
            
        Returns:
            Dictionary with exploit information
        """
        if not result or not result.get('found'):
            return {'success': False, 'reason': 'No heap overflow found'}
        
        # Basic heap overflow exploit - just the correct length of data
        payload = b'A' * result['length']
        
        # Create interaction sequence based on the pattern that worked
        interaction_sequence = []
        
        # Create the exploit result
        exploit = {
            'success': True,
            'type': 'heap_overflow',
            'payload': payload.hex(),
            'interaction_sequence': interaction_sequence,
            'notes': f"Heap overflow with {result['length']} bytes. May need manual adjustment."
        }
        
        return exploit 