"""
Buffer overflow finder module for pwnelle.

Automatically detects the offset needed to control the return address.
"""

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

class BufferOverflowFinder:
    """Find the exact offset for a buffer overflow vulnerability."""
    
    def __init__(
        self,
        binary_path: str,
        analysis_results: Dict[str, Any],
        max_length: int = 2000,
        step_size: int = 8,
        timeout: int = 3,
        verbose: bool = False,
        remote_options: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize the buffer overflow finder.
        
        Args:
            binary_path: Path to the binary to analyze
            analysis_results: Analysis results from the BinaryAnalyzer
            max_length: Maximum length to try for the payload
            step_size: Step size for the payload length
            timeout: Timeout for process in seconds
            verbose: Whether to show verbose output
            remote_options: Optional dict with remote server connection details
                           Format: {"host": str, "port": int, "protocol": str, 
                                   "delay": float, "retries": int}
        """
        self.binary_path = binary_path
        self.analysis_results = analysis_results
        self.max_length = max_length
        self.step_size = step_size
        self.timeout = timeout
        self.verbose = verbose
        self.remote_options = remote_options
        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 create_process(self) -> tube:
        """
        Create either a local process or remote connection based on remote_options.
        
        Returns:
            A pwntools tube (process or remote)
        """
        if self.remote_options:
            self.log(f"Connecting to remote server: {self.remote_options['host']}:{self.remote_options['port']}")
            
            # Perform retry logic for remote connections
            retries = self.remote_options.get('retries', 3)
            for attempt in range(retries):
                try:
                    # Create remote connection
                    if self.remote_options.get('protocol', 'tcp').lower() == 'tcp':
                        p = remote(self.remote_options['host'], self.remote_options['port'], timeout=self.timeout)
                    else:
                        # UDP connection
                        p = remote(self.remote_options['host'], self.remote_options['port'], typ='udp', timeout=self.timeout)
                    
                    # Apply delay if specified
                    if self.remote_options.get('delay', 0) > 0:
                        time.sleep(self.remote_options['delay'])
                    
                    return p
                except Exception as e:
                    self.log(f"Connection attempt {attempt+1}/{retries} failed: {e}")
                    if attempt < retries - 1:
                        time.sleep(1)  # Wait between retries
                    else:
                        raise
            
            # If we get here, all retries failed
            raise Exception("All connection attempts to remote server failed")
        else:
            # Create local process
            return process(self.binary_path, timeout=self.timeout)
    
    def find_offset(self, interaction_sequence: Optional[List[bytes]] = None) -> Optional[Dict[str, Any]]:
        """
        Find the buffer overflow offset using pattern payload.
        
        Args:
            interaction_sequence: Optional sequence of interactions before sending payload
                                  (e.g., menus to navigate through)
        
        Returns:
            Dictionary with offset info or None if not found
        """
        if self.remote_options:
            target_desc = f"remote {self.remote_options['host']}:{self.remote_options['port']}"
        else:
            target_desc = self.binary_path
            
        self.log(f"Finding buffer overflow offset in {target_desc}")
        
        # Try cyclic pattern with different lengths
        for length in range(8, self.max_length + 1, self.step_size):
            try:
                self.log(f"Trying length {length}...")
                result = self._try_pattern(length, interaction_sequence)
                if result:
                    return result
            except Exception as e:
                self.log(f"Error in finding offset: {e}")
                continue
        
        return None
    
    def _try_pattern(self, length: int, interaction_sequence: Optional[List[bytes]]) -> Optional[Dict[str, Any]]:
        """
        Try a cyclic pattern of the given length.
        
        Args:
            length: Length of the pattern
            interaction_sequence: Optional sequence of interactions before sending payload
        
        Returns:
            Dictionary with offset info or None if not found
        """
        # Generate cyclic pattern
        pattern = cyclic(length)
        
        # Start the process or connect to remote
        try:
            p = self.create_process()
            
            # Perform required interactions if any
            if interaction_sequence:
                for interaction in interaction_sequence:
                    p.sendafter(interaction, b"Y\n")
            
            # Send the pattern
            p.sendline(pattern)
            
            # Check for crash
            try:
                # Wait a bit for crash to happen
                time.sleep(0.5)
                # Try to receive anything (this will timeout if process crashed)
                p.recv(timeout=1)
                # If we get here, no crash yet
                self.log("No crash detected, trying next length")
                p.close()
                return None
            except EOFError:
                # Process crashed or exited
                self.log("Process crashed or exited")
                crash_result = self._analyze_crash(pattern)
                p.close()
                return crash_result
            finally:
                # Only attempt to kill/close if it's a local process
                if not self.remote_options and p.poll() is None:
                    p.kill()
                    p.close()
                elif self.remote_options:
                    try:
                        p.close()
                    except:
                        pass
        
        except Exception as e:
            self.log(f"Error during pattern test: {e}")
            return None
    
    def _analyze_crash(self, pattern: bytes) -> Optional[Dict[str, Any]]:
        """
        Analyze the crash to find the exact offset.
        
        Args:
            pattern: The pattern that caused the crash
        
        Returns:
            Dictionary with offset info or None if not found
        """
        # For remote targets, we can't analyze the crash directly
        # We need to use binary analysis and educated guesses
        if self.remote_options:
            self.log("Remote target detected, using heuristic analysis to estimate offset")
            return self._estimate_offset(pattern)
            
        # For local targets, we can use GDB
        try:
            # Try to run with GDB to get exact register values
            self.log("Analyzing crash with GDB...")
            
            # Create a GDB script to run and get register state
            script_content = "run < <(python3 -c 'print(\"" + pattern.decode() + "\")')\n"
            script_content += "info registers\n"
            
            # Create a temporary gdb script
            script_file = "gdb_crash_script.txt"
            with open(script_file, "w") as f:
                f.write(script_content)
            
            # Run GDB with the script
            try:
                gdb_cmd = ["gdb", "-batch", "-x", script_file, self.binary_path]
                gdb_output = subprocess.check_output(gdb_cmd, stderr=subprocess.STDOUT, universal_newlines=True)
                
                # Extract register values from output
                rip_match = re.search(r"rip\s+0x([0-9a-f]+)", gdb_output, re.IGNORECASE)
                rsp_match = re.search(r"rsp\s+0x([0-9a-f]+)", gdb_output, re.IGNORECASE)
                rbp_match = re.search(r"rbp\s+0x([0-9a-f]+)", gdb_output, re.IGNORECASE)
                
                os.remove(script_file)
                
                # If RIP was overwritten with our pattern
                if rip_match:
                    rip_value = int(rip_match.group(1), 16)
                    try:
                        # Try to decode the pattern
                        offset = cyclic_find(p64(rip_value)[:4])
                        self.log(f"Found offset: {offset} (RIP)")
                        return {
                            "offset": offset,
                            "register": "rip",
                            "value": rip_value,
                            "type": "direct_ret_overwrite"
                        }
                    except:
                        # Could not find pattern in RIP
                        pass
                
                # If RSP or RBP was overwritten with something interesting
                if rsp_match or rbp_match:
                    # If we have stack pointer, we might be able to find an offset
                    if rsp_match:
                        rsp_value = int(rsp_match.group(1), 16)
                        # Not implemented: use rsp value to derive offset
                        # This is more complex and requires reading memory
                    
                    # If we have base pointer, check for a pattern in it
                    if rbp_match:
                        rbp_value = int(rbp_match.group(1), 16)
                        try:
                            offset = cyclic_find(p64(rbp_value)[:4])
                            self.log(f"Found offset: {offset} (RBP)")
                            # RBP overwrite usually means RIP is at offset+8
                            return {
                                "offset": offset + 8,  # RIP is usually 8 bytes after RBP
                                "register": "rbp",
                                "value": rbp_value,
                                "type": "frame_pointer_overwrite"
                            }
                        except:
                            # Could not find pattern in RBP
                            pass
            
            except Exception as e:
                self.log(f"Error running GDB: {e}")
                if os.path.exists(script_file):
                    os.remove(script_file)
            
            return self._estimate_offset(pattern)
                
        except Exception as e:
            self.log(f"Error analyzing crash: {e}")
            return None
    
    def _estimate_offset(self, pattern: bytes) -> Dict[str, Any]:
        """
        Estimate buffer overflow offset using binary analysis.
        
        Args:
            pattern: The pattern that caused the crash
            
        Returns:
            Dictionary with estimated offset info
        """
        # Fallback: Try to use heuristics based on binary size
        if "basic_info" in self.analysis_results:
            binary_size = self.analysis_results["basic_info"].get("file_size", 0)
            # Very basic heuristic - small binaries often have small buffers
            if binary_size < 20000:  # Less than 20KB
                self.log("Small binary detected, trying common offset of 40")
                return {
                    "offset": 40,
                    "register": "unknown",
                    "value": 0,
                    "type": "estimated_small_binary",
                    "confidence": "low"
                }
            else:
                self.log("Medium/large binary detected, trying common offset of 72")
                return {
                    "offset": 72,
                    "register": "unknown",
                    "value": 0,
                    "type": "estimated_medium_binary",
                    "confidence": "low"
                }
        
        # Last resort: return a default offset with very low confidence
        return {
            "offset": 64,
            "register": "unknown",
            "value": 0, 
            "type": "default_guess",
            "confidence": "very_low"
        } 