"""
Exploit tester module for pwnelle.

Tests and validates automatically generated exploits.
"""

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

class ExploitTester:
    """Run and test exploits, detecting successful exploitation."""
    
    def __init__(
        self,
        binary_path: str,
        analysis_results: Dict[str, Any],
        output_dir: str,
        timeout: int = 10,
        flag_regex: str = r"(flag\{[^}]*\}|CTF\{[^}]*\})",
        verbose: bool = False,
        remote_options: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize the exploit tester.
        
        Args:
            binary_path: Path to the binary to analyze
            analysis_results: Analysis results from the BinaryAnalyzer
            output_dir: Directory to write output files
            timeout: Timeout for process in seconds
            flag_regex: Regex pattern to identify flags
            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.output_dir = output_dir
        self.timeout = timeout
        self.flag_regex = flag_regex
        self.verbose = verbose
        self.remote_options = remote_options
        self.context = context
        
        # Initialize pwntools context
        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 test_exploit(self, payload: bytes, interaction_sequence: Optional[List[Tuple[bytes, bytes]]] = None) -> Dict[str, Any]:
        """
        Test an exploit payload against the binary.
        
        Args:
            payload: The exploit payload to send
            interaction_sequence: Optional sequence of (expect, send) tuples for initial interaction
                                  (e.g., menu navigation before sending main payload)
        
        Returns:
            Dictionary with test results
        """
        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"Testing exploit against {target_desc}")
        
        result = {
            "success": False,
            "flag": None,
            "shell_obtained": False,
            "crash": False,
            "output": "",
            "error": None,
            "is_remote": bool(self.remote_options)
        }
        
        try:
            # Launch the process or connect to remote
            p = self.create_process()
            
            # Perform initial interactions if needed
            if interaction_sequence:
                for expect_bytes, send_bytes in interaction_sequence:
                    p.sendafter(expect_bytes, send_bytes)
            
            # Send the payload
            p.sendline(payload)
            
            # Try to check for shell
            shell_check_result = self._check_for_shell(p)
            if shell_check_result["shell_obtained"]:
                result.update(shell_check_result)
                p.close()
                return result
            
            # Try to find flag in output
            flag_check_result = self._check_for_flag(p)
            if flag_check_result["flag"]:
                result.update(flag_check_result)
                p.close()
                return result
            
            # If no immediate shell or flag, collect output
            try:
                output = p.recvall(timeout=5).decode('utf-8', errors='ignore')
                result["output"] = output
                
                # Check for flag in full output
                flag_match = re.search(self.flag_regex, output)
                if flag_match:
                    result["success"] = True
                    result["flag"] = flag_match.group(0)
            except:
                # If recvall times out, the process might be running but not producing output
                # We'll check if it's a shell below
                pass
            
            # Final shell check
            if hasattr(p, 'poll') and p.poll() is None:  # Process is still running (local only)
                shell_check_result = self._check_for_shell(p)
                result.update(shell_check_result)
            elif hasattr(p, 'poll'):  # Local process that has exited
                if p.poll() != 0:
                    result["crash"] = True
                    self.log(f"Process exited with status {p.poll()}")
            
            p.close()
        
        except Exception as e:
            result["error"] = str(e)
            self.log(f"Error testing exploit: {e}")
            if 'p' in locals():
                try:
                    p.close()
                except:
                    pass
        
        return result
    
    def _check_for_shell(self, p) -> Dict[str, Any]:
        """
        Check if we've obtained a shell.
        
        Args:
            p: The process to check
        
        Returns:
            Dictionary with shell check results
        """
        result = {
            "success": False,
            "shell_obtained": False,
            "shell_check_output": "",
            "flag": None
        }
        
        try:
            # Check if we can run commands in the shell
            p.sendline(b"id")
            p.sendline(b"echo SHELLCHECK")
            
            # Give it a moment
            time.sleep(0.5)
            
            # Receive output
            try:
                output = p.recvuntil(b"SHELLCHECK", timeout=2)
                output = output.decode('utf-8', errors='ignore')
                result["shell_check_output"] = output
                
                # Check for common shell indicators
                if any(indicator in output for indicator in ["uid=", "gid=", "SHELLCHECK"]):
                    result["shell_obtained"] = True
                    result["success"] = True
                    self.log("Shell obtained successfully!")
                    
                    # Try to read flag files
                    p.sendline(b"cat flag.txt || cat flag || find / -name flag.txt 2>/dev/null")
                    time.sleep(0.5)
                    flag_output = p.recv(timeout=2).decode('utf-8', errors='ignore')
                    
                    # Look for a flag in the output
                    flag_match = re.search(self.flag_regex, flag_output)
                    if flag_match:
                        result["flag"] = flag_match.group(0)
                        self.log(f"Flag found: {result['flag']}")
            except:
                # No shell indicators found
                pass
        
        except Exception as e:
            self.log(f"Error checking for shell: {e}")
        
        return result
    
    def _check_for_flag(self, p) -> Dict[str, Any]:
        """
        Check if output contains a flag.
        
        Args:
            p: The process to check
        
        Returns:
            Dictionary with flag check results
        """
        result = {
            "success": False,
            "flag": None,
            "output": ""
        }
        
        try:
            # Try to receive output that might contain a flag
            output = p.recv(timeout=2).decode('utf-8', errors='ignore')
            result["output"] = output
            
            # Look for flag pattern
            flag_match = re.search(self.flag_regex, output)
            if flag_match:
                result["flag"] = flag_match.group(0)
                result["success"] = True
                self.log(f"Flag found: {result['flag']}")
        
        except Exception as e:
            self.log(f"Error checking for flag: {e}")
        
        return result
    
    def generate_exploit_report(self, test_result: Dict[str, Any]) -> str:
        """
        Generate a human-readable report from test results.
        
        Args:
            test_result: The test results from test_exploit
        
        Returns:
            A formatted string with the report
        """
        report = []
        report.append("## Exploit Test Report")
        report.append(f"Target: {self.binary_path}")
        report.append(f"Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        report.append("")
        
        if test_result["success"]:
            report.append("### 🟢 SUCCESS")
            if test_result.get("shell_obtained"):
                report.append("Shell successfully obtained!")
            if test_result.get("flag"):
                report.append(f"Flag found: `{test_result['flag']}`")
        else:
            report.append("### 🔴 FAILURE")
            if test_result.get("crash"):
                report.append("Binary crashed during exploitation.")
            if test_result.get("error"):
                report.append(f"Error: {test_result['error']}")
        
        report.append("")
        report.append("### Output")
        report.append("```")
        report.append(test_result.get("output", "No output captured"))
        report.append("```")
        
        if test_result.get("shell_check_output"):
            report.append("")
            report.append("### Shell Check Output")
            report.append("```")
            report.append(test_result["shell_check_output"])
            report.append("```")
        
        return "\n".join(report)
    
    def save_exploit_report(self, test_result: Dict[str, Any]) -> str:
        """
        Save the exploit report to a file.
        
        Args:
            test_result: The test results from test_exploit
        
        Returns:
            Path to the saved report file
        """
        report = self.generate_exploit_report(test_result)
        
        report_path = os.path.join(self.output_dir, "exploit_report.md")
        with open(report_path, "w", encoding="utf-8") as f:
            f.write(report)
        
        return report_path 