Pwno MCP

"""
Subprocess execution tools for Pwno MCP

Provides tools for running system commands, particularly useful for:
- Compiling binaries with various flags (ASAN, debug symbols, etc.)
- Running helper scripts
- Managing background processes
"""

import logging
import subprocess
import shlex
from typing import Dict, Any, Optional, List
import psutil
import os
import tempfile
import time
import re

logger = logging.getLogger(__name__)


class SubprocessTools:
    """Tools for subprocess execution and management"""
    
    def __init__(self):
        """Initialize subprocess tools"""
        self.background_processes: Dict[int, Dict[str, Any]] = {}
        
    def run_command(self, command: str, cwd: Optional[str] = None, env: Optional[Dict[str, str]] = None, 
                   timeout: Optional[float] = 30.0) -> Dict[str, Any]:
        """
        Execute a command and wait for it to complete
        
        This is primarily intended for compilation commands like:
        - gcc -g -fsanitize=address program.c -o program
        - clang -O0 -g -fno-omit-frame-pointer vuln.c
        - make clean && make
        
        Args:
            command: Command to execute (will be shell-parsed)
            cwd: Working directory for the command
            env: Environment variables (merged with current env)
            timeout: Maximum execution time in seconds
            
        Returns:
            Dictionary with execution results
        """
        logger.info(f"Running command: {command}")
        
        try:
            # Parse command for safer execution
            cmd_parts = shlex.split(command)
            
            # Prepare environment
            cmd_env = os.environ.copy()
            if env:
                cmd_env.update(env)
                
            # Execute command
            result = subprocess.run(
                cmd_parts,
                cwd=cwd,
                env=cmd_env,
                capture_output=True,
                text=True,
                timeout=timeout
            )
            
            return {
                "success": result.returncode == 0,
                "command": command,
                "returncode": result.returncode,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "cwd": cwd or os.getcwd()
            }
            
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "command": command,
                "error": f"Command timed out after {timeout} seconds",
                "cwd": cwd or os.getcwd()
            }
        except Exception as e:
            logger.error(f"Failed to run command: {e}")
            return {
                "success": False,
                "command": command,
                "error": str(e),
                "cwd": cwd or os.getcwd()
            }
            
    def spawn_process(self, command: str, cwd: Optional[str] = None, env: Optional[Dict[str, str]] = None) -> Dict[str, Any]:
        """
        Spawn a background process and return immediately
        
        Useful for starting long-running processes like:
        - Web servers for exploitation
        - Network listeners
        - Monitoring scripts
        
        Args:
            command: Command to execute
            cwd: Working directory for the command
            env: Environment variables
            
        Returns:
            Dictionary with process information including PID
        """
        logger.info(f"Spawning process: {command}")
        # Create temp files for stdout and stderr
        stdout_fd, stdout_path = tempfile.mkstemp(prefix='pwno_stdout_', suffix='.log')
        stderr_fd, stderr_path = tempfile.mkstemp(prefix='pwno_stderr_', suffix='.log')
        stdout_file = os.fdopen(stdout_fd, 'w+')
        stderr_file = os.fdopen(stderr_fd, 'w+')
        
        try:
            # Parse command
            cmd_parts = shlex.split(command)
            
            # Prepare environment
            cmd_env = os.environ.copy()
            if env:
                cmd_env.update(env)
                
            # Spawn process with stdout and stderr redirected to files
            process = subprocess.Popen(
                cmd_parts,
                cwd=cwd,
                env=cmd_env,
                stdout=stdout_file,
                stderr=stderr_file,
                text=True
            )
            
            # Store process reference and file paths
            self.background_processes[process.pid] = {
                'process': process,
                'stdout_file': stdout_file,
                'stderr_file': stderr_file,
                'stdout_path': stdout_path,
                'stderr_path': stderr_path,
                'command': command,
                'cwd': cwd or os.getcwd()
            }
            
            # Give process a moment to potentially fail
            import time
            time.sleep(0.1)
            
            if process.poll() is not None:
                # Process already terminated, read outputs
                stdout_file.flush()
                stderr_file.flush()
                with open(stdout_path, 'r') as f:
                    stdout = f.read()
                with open(stderr_path, 'r') as f:
                    stderr = f.read()
                return {
                    'success': False,
                    'command': command,
                    'error': 'Process terminated immediately',
                    'returncode': process.returncode,
                    'stdout': stdout,
                    'stderr': stderr,
                    'stdout_path': stdout_path,
                    'stderr_path': stderr_path
                }
                
            return {
                'success': True,
                'command': command,
                'pid': process.pid,
                'cwd': cwd or os.getcwd(),
                'status': 'running',
                'stdout_path': stdout_path,
                'stderr_path': stderr_path
            }
            
        except Exception as e:
            logger.error(f"Failed to spawn process: {e}")
            return {
                "success": False,
                "command": command,
                "error": str(e)
            }
            
    def get_process(self, pid: int) -> Dict[str, Any]:
        """
        Get status of a spawned process
        
        Args:
            pid: Process ID to check
            
        Returns:
            Dictionary with process status. When available, includes live
            standard stream contents as strings:
            - ``stdout``: accumulated standard output so far
            - ``stderr``: accumulated standard error so far
            Paths to the backing log files (``stdout_path``, ``stderr_path``)
            are also returned for external tailing.
        """
        try:
            # Check if we're tracking this process
            if pid in self.background_processes:
                entry = self.background_processes[pid]
                process = entry['process']
                poll_result = process.poll()
                
                if poll_result is None:
                    # Still running, return status and output paths
                    # Flush current buffers and read accumulated outputs
                    try:
                        entry['stdout_file'].flush()
                    except Exception:
                        pass
                    try:
                        entry['stderr_file'].flush()
                    except Exception:
                        pass
                    try:
                        with open(entry['stdout_path'], 'r', encoding='utf-8', errors='ignore') as f:
                            live_stdout = f.read()
                    except Exception:
                        live_stdout = ''
                    try:
                        with open(entry['stderr_path'], 'r', encoding='utf-8', errors='ignore') as f:
                            live_stderr = f.read()
                    except Exception:
                        live_stderr = ''
                    return {
                        'success': True,
                        'pid': pid,
                        'status': 'running',
                        'stdout_path': entry['stdout_path'],
                        'stderr_path': entry['stderr_path'],
                        'stdout': live_stdout,
                        'stderr': live_stderr,
                        'cpu_percent': psutil.Process(pid).cpu_percent(),
                        'memory_info': psutil.Process(pid).memory_info()._asdict()
                    }
                else:
                    # Process finished, read outputs
                    entry['stdout_file'].close()
                    entry['stderr_file'].close()
                    with open(entry['stdout_path'], 'r') as f:
                        stdout = f.read()
                    with open(entry['stderr_path'], 'r') as f:
                        stderr = f.read()
                    del self.background_processes[pid]
                    return {
                        'success': True,
                        'pid': pid,
                        'status': 'terminated',
                        'returncode': poll_result,
                        'stdout': stdout,
                        'stderr': stderr,
                        'stdout_path': entry['stdout_path'],
                        'stderr_path': entry['stderr_path']
                    }
            else:
                # Try to check if process exists anyway
                if psutil.pid_exists(pid):
                    proc = psutil.Process(pid)
                    return {
                        "success": True,
                        "pid": pid,
                        "status": "running" if proc.is_running() else "unknown",
                        "name": proc.name(),
                        "cmdline": " ".join(proc.cmdline())
                    }
                else:
                    return {
                        "success": False,
                        "pid": pid,
                        "error": "Process not found"
                    }
                    
        except Exception as e:
            logger.error(f"Failed to get process status: {e}")
            return {
                "success": False,
                "pid": pid,
                "error": str(e)
            }
            
    def kill_process(self, pid: int, signal: int = 15) -> Dict[str, Any]:
        """
        Kill a process
        
        Args:
            pid: Process ID to kill
            signal: Signal to send (default: SIGTERM=15, use 9 for SIGKILL)
            
        Returns:
            Dictionary with kill result
        """
        try:
            if pid in self.background_processes:
                process = self.background_processes[pid]['process']
                process.terminate() if signal == 15 else process.kill()
                process.wait(timeout=5)
                # Keep the process entry cached even after kill, so get_process can retrieve outputs
                # del self.background_processes[pid]
            else:
                os.kill(pid, signal)
                
            return {
                "success": True,
                "pid": pid,
                "signal": signal,
                "status": "killed"
            }
            
        except ProcessLookupError:
            return {
                "success": False,
                "pid": pid,
                "error": "Process not found"
            }
        except Exception as e:
            logger.error(f"Failed to kill process: {e}")
            return {
                "success": False,
                "pid": pid,
                "error": str(e)
            }
            
    def list_processes(self) -> Dict[str, Any]:
        """
        List all tracked background processes
        
        Returns:
            Dictionary with process list
        """
        processes = []
        
        for pid, entry in list(self.background_processes.items()):
            process = entry['process']
            poll_result = process.poll()
            if poll_result is None:
                # Still running
                try:
                    proc_info = psutil.Process(pid)
                    processes.append({
                        "pid": pid,
                        "status": "running",
                        "name": proc_info.name(),
                        "cmdline": " ".join(proc_info.cmdline()),
                        "cpu_percent": proc_info.cpu_percent(),
                        "memory_mb": proc_info.memory_info().rss / 1024 / 1024
                    })
                except:
                    processes.append({
                        "pid": pid,
                        "status": "running",
                        "error": "Could not get process info"
                    })
            else:
                # Process terminated, clean up
                entry['stdout_file'].close()
                entry['stderr_file'].close()
                del self.background_processes[pid]
                
        return {
            "success": True,
            "processes": processes,
            "count": len(processes)
        } 

    def wait_for_pid_marker(self, stdout_path: str, timeout: float = 10.0, poll_interval: float = 0.05) -> Dict[str, Any]:
        """
        Wait for a PID marker of the form "<PID>{1234}</PID>" to appear in a stdout log file.

        Args:
            stdout_path: Path to the stdout log file produced by spawn_process
            timeout: Maximum time to wait in seconds
            poll_interval: How often to poll for new output

        Returns:
            Dictionary with keys:
            - success: whether a PID marker was found
            - pid: the captured PID if found
            - elapsed: time spent waiting
            - bytes_scanned: how many bytes of the file were scanned
            - error: error message if any
        """
        start_time = time.time()
        pattern = re.compile(r"<PID>\{(?P<pid>\d+)\}</PID>")
        bytes_scanned = 0

        try:
            if not os.path.exists(stdout_path):
                return {"success": False, "error": "stdout file not found", "elapsed": 0.0, "bytes_scanned": 0}

            with open(stdout_path, "r", encoding="utf-8", errors="ignore") as f:
                while True:
                    chunk = f.read()
                    if chunk:
                        bytes_scanned += len(chunk)
                        match = pattern.search(chunk)
                        if match:
                            elapsed = time.time() - start_time
                            return {
                                "success": True,
                                "pid": int(match.group("pid")),
                                "elapsed": elapsed,
                                "bytes_scanned": bytes_scanned
                            }

                    if time.time() - start_time >= timeout:
                        return {"success": False, "error": "timeout", "elapsed": timeout, "bytes_scanned": bytes_scanned}
                    time.sleep(poll_interval)
        except Exception as e:
            logger.error(f"Failed while waiting for PID marker: {e}")
            return {"success": False, "error": str(e), "elapsed": time.time() - start_time, "bytes_scanned": bytes_scanned}