async_shell.py

#!/usr/bin/env python3
"""
Async Shell Runner: A utility for executing shell commands asynchronously with timeout control.

This module provides a robust way to execute shell commands in an asynchronous manner,
with support for timeouts, output truncation, and comprehensive error handling.

Features:
- Run shell commands asynchronously
- Set timeouts to prevent hanging on long-running commands
- Capture and truncate stdout/stderr if needed
- Handle various error conditions gracefully
- Support for command cancellation
"""

import asyncio
import logging
import os
import shlex
import signal
import sys
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('async_shell')

# Constants
DEFAULT_TIMEOUT = 120.0  # Default timeout in seconds
MAX_OUTPUT_LENGTH = 16000  # Maximum length of captured output
TRUNCATION_MESSAGE = "<output truncated...>"


@dataclass
class CommandResult:
    """Container for the results of a command execution."""
    command: str
    returncode: int
    stdout: str
    stderr: str
    timeout_occurred: bool = False
    cancelled: bool = False
    error: Optional[str] = None
    
    @property
    def success(self) -> bool:
        """Returns True if the command completed successfully with return code 0."""
        return self.returncode == 0 and not self.timeout_occurred and not self.cancelled
    
    def __str__(self) -> str:
        """Pretty string representation of the command result."""
        status = "SUCCESS" if self.success else "FAILED"
        if self.timeout_occurred:
            status = "TIMEOUT"
        elif self.cancelled:
            status = "CANCELLED"
            
        result = [
            f"Command: {self.command}",
            f"Status: {status}",
            f"Return code: {self.returncode}",
        ]
        
        if self.error:
            result.append(f"Error: {self.error}")
        
        if self.stdout:
            stdout_excerpt = self.stdout[:500] + "..." if len(self.stdout) > 500 else self.stdout
            result.append(f"Stdout: {stdout_excerpt}")
        
        if self.stderr:
            stderr_excerpt = self.stderr[:500] + "..." if len(self.stderr) > 500 else self.stderr
            result.append(f"Stderr: {stderr_excerpt}")
            
        return "\n".join(result)


def maybe_truncate(content: str, max_length: Optional[int] = MAX_OUTPUT_LENGTH) -> str:
    """Truncate content if it exceeds the specified maximum length.
    
    Args:
        content: The string content to truncate if necessary
        max_length: Maximum length allowed, or None for no limit
        
    Returns:
        The original string if it's within length limits, or a truncated version
    """
    if not max_length or len(content) <= max_length:
        return content
    return content[:max_length] + TRUNCATION_MESSAGE


async def run_shell_command(
    cmd: str,
    timeout: Optional[float] = DEFAULT_TIMEOUT,
    truncate_length: Optional[int] = MAX_OUTPUT_LENGTH,
    env: Optional[Dict[str, str]] = None,
    cwd: Optional[str] = None,
    shell: bool = True,
    raise_on_error: bool = False,
) -> CommandResult:
    """Run a shell command asynchronously with timeout.
    
    Args:
        cmd: The shell command to run
        timeout: Maximum execution time in seconds (None for no timeout)
        truncate_length: Maximum length for stdout/stderr (None for no truncation)
        env: Optional environment variables for the subprocess
        cwd: Optional working directory for the subprocess
        shell: Whether to run the command through the shell
        raise_on_error: Whether to raise exceptions on non-zero return codes
        
    Returns:
        CommandResult object containing execution results
        
    Raises:
        TimeoutError: If the command times out and raise_on_error is True
        RuntimeError: If the command returns non-zero and raise_on_error is True
    """
    logger.debug(f"Executing command: {cmd}")
    
    # Prepare process environment
    process_env = os.environ.copy()
    if env:
        process_env.update(env)
    
    # Start the subprocess
    try:
        if shell:
            process = await asyncio.create_subprocess_shell(
                cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=process_env,
                cwd=cwd,
            )
        else:
            process = await asyncio.create_subprocess_exec(
                *shlex.split(cmd),
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=process_env,
                cwd=cwd,
            )
    except Exception as e:
        logger.error(f"Failed to start command: {cmd}", exc_info=True)
        return CommandResult(
            command=cmd,
            returncode=-1,
            stdout="",
            stderr="",
            error=str(e)
        )
    
    # Execute with timeout handling
    try:
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout)
        stdout_text = maybe_truncate(stdout.decode(), truncate_length)
        stderr_text = maybe_truncate(stderr.decode(), truncate_length)
        
        result = CommandResult(
            command=cmd,
            returncode=process.returncode,
            stdout=stdout_text,
            stderr=stderr_text
        )
        
        if raise_on_error and process.returncode != 0:
            raise RuntimeError(f"Command failed with return code {process.returncode}: {cmd}")
            
        return result
        
    except asyncio.TimeoutError:
        logger.warning(f"Command timed out after {timeout} seconds: {cmd}")
        
        # Try to kill the process
        try:
            process.kill()
            await process.wait()
        except ProcessLookupError:
            logger.debug("Process already terminated")
        except Exception as e:
            logger.error(f"Error killing process: {e}")
        
        result = CommandResult(
            command=cmd,
            returncode=-1,
            stdout="",
            stderr="",
            timeout_occurred=True,
            error=f"Command timed out after {timeout} seconds"
        )
        
        if raise_on_error:
            raise TimeoutError(f"Command timed out after {timeout} seconds: {cmd}")
            
        return result
    
    except asyncio.CancelledError:
        logger.info(f"Command cancelled: {cmd}")
        
        # Try to kill the process when cancelled
        try:
            process.kill()
            await process.wait()
        except Exception:
            pass
            
        return CommandResult(
            command=cmd,
            returncode=-1,
            stdout="",
            stderr="",
            cancelled=True,
            error="Command was cancelled"
        )
        
    except Exception as e:
        logger.error(f"Error executing command: {cmd}", exc_info=True)
        
        # Try to kill the process if it's still running
        try:
            process.kill()
            await process.wait()
        except Exception:
            pass
            
        return CommandResult(
            command=cmd,
            returncode=-1,
            stdout="",
            stderr="",
            error=str(e)
        )


async def run_multiple_commands(
    commands: List[str],
    timeout: Optional[float] = DEFAULT_TIMEOUT,
    truncate_length: Optional[int] = MAX_OUTPUT_LENGTH,
    parallel: bool = False,
    **kwargs
) -> List[CommandResult]:
    """Run multiple shell commands, either sequentially or in parallel.
    
    Args:
        commands: List of shell commands to execute
        timeout: Maximum execution time per command
        truncate_length: Maximum output length per command
        parallel: If True, run commands in parallel; otherwise run sequentially
        **kwargs: Additional arguments to pass to run_shell_command
        
    Returns:
        List of CommandResult objects
    """
    if not commands:
        return []
        
    if parallel:
        # Run all commands in parallel
        tasks = [
            run_shell_command(cmd, timeout, truncate_length, **kwargs)
            for cmd in commands
        ]
        return await asyncio.gather(*tasks)
    else:
        # Run commands sequentially
        results = []
        for cmd in commands:
            result = await run_shell_command(cmd, timeout, truncate_length, **kwargs)
            results.append(result)
        return results


# Command-line interface
if __name__ == "__main__":
    async def main():
        import argparse
        
        parser = argparse.ArgumentParser(description="Run shell commands asynchronously with timeout")
        parser.add_argument("command", nargs="?", help="Command to execute")
        parser.add_argument("--commands-file", "-f", help="File with commands to run (one per line)")
        parser.add_argument("--timeout", "-t", type=float, default=DEFAULT_TIMEOUT, help="Command timeout in seconds")
        parser.add_argument("--parallel", "-p", action="store_true", help="Run multiple commands in parallel")
        parser.add_argument("--no-truncate", action="store_true", help="Don't truncate command output")
        parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging")
        
        args = parser.parse_args()
        
        if args.verbose:
            logger.setLevel(logging.DEBUG)
            
        truncate_length = None if args.no_truncate else MAX_OUTPUT_LENGTH
        
        if args.command:
            # Run a single command
            result = await run_shell_command(args.command, args.timeout, truncate_length)
            print(result)
            sys.exit(0 if result.success else 1)
            
        elif args.commands_file:
            # Run commands from file
            try:
                with open(args.commands_file, 'r') as f:
                    commands = [line.strip() for line in f if line.strip()]
            except Exception as e:
                logger.error(f"Failed to read commands file: {e}")
                sys.exit(1)
                
            if not commands:
                logger.error("No commands found in the file")
                sys.exit(1)
                
            results = await run_multiple_commands(
                commands, 
                args.timeout,
                truncate_length,
                parallel=args.parallel
            )
            
            # Print all results
            for i, result in enumerate(results):
                print(f"\n=== Command {i+1}/{len(results)} ===")
                print(result)
                
            # Check if all commands succeeded
            success = all(result.success for result in results)
            sys.exit(0 if success else 1)
        
        else:
            parser.print_help()
            sys.exit(1)
    
    asyncio.run(main())