betterbash

#!/usr/bin/env python3
"""
BetterBash CLI - Enhanced shell execution with safety controls
Direct CLI interface using the same logic as the BetterBash MCP server.

Usage:
  betterbash "ls -la"
  betterbash "git status"
  betterbash --help
"""

import sys
import subprocess
import argparse
import time
from pathlib import Path

def load_allowed_commands():
    """Load allowed commands from file"""
    allowed_file = Path(__file__).parent / "allowed_commands.txt"
    if allowed_file.exists():
        with open(allowed_file) as f:
            return set(line.strip() for line in f if line.strip() and not line.startswith('#'))
    return {
        'python3', 'pip', 'git', 'curl', 'ls', 'cat', 'head', 'tail',
        'mkdir', 'cp', 'mv', 'rm', 'find', 'grep', 'which', 'say',
        'docker', 'docker-compose', 'nvidia-smi'
    }

def is_command_allowed(command, allowed_commands):
    """Check if command is in allowlist"""
    base_cmd = command.split()[0] if command else ""
    return base_cmd in allowed_commands

def execute_command(command, timeout=30):
    """Execute command with safety controls and logging"""
    allowed_commands = load_allowed_commands()
    
    if not is_command_allowed(command, allowed_commands):
        base_cmd = command.split()[0] if command else ""
        print(f"Command '{base_cmd}' not in allowlist.", file=sys.stderr)
        print(f"Add it with: betterbash --add \"{base_cmd}\"", file=sys.stderr)
        print(f"Or list all allowed commands: betterbash --list", file=sys.stderr)
        return 1
    
    try:
        start_time = time.time()
        result = subprocess.run(
            command,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        
        duration = time.time() - start_time
        
        # Print stdout and stderr
        if result.stdout:
            print(result.stdout, end="")
        if result.stderr:
            print(result.stderr, file=sys.stderr, end="")
            
        # Log execution (basic version)
        log_entry = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {command} -> exit:{result.returncode} time:{duration:.2f}s"
        
        return result.returncode
        
    except subprocess.TimeoutExpired:
        print(f"Command timed out after {timeout} seconds", file=sys.stderr)
        return 124
    except Exception as e:
        print(f"Error executing command: {e}", file=sys.stderr)
        return 1

def add_command_to_allowlist(command):
    """Add a command to the allowlist file"""
    allowed_file = Path(__file__).parent / "allowed_commands.txt"
    
    try:
        # Read current commands
        current_commands = load_allowed_commands()
        
        if command in current_commands:
            print(f"Command '{command}' is already in the allowlist.")
            return 0
        
        # Append new command to file
        with open(allowed_file, 'a') as f:
            f.write(f"\n{command}")
        
        print(f"Added '{command}' to allowlist.")
        print(f"You can now run: betterbash \"{command} ...\"")
        return 0
        
    except Exception as e:
        print(f"Error adding command to allowlist: {e}", file=sys.stderr)
        return 1

def main():
    parser = argparse.ArgumentParser(
        description="BetterBash - Enhanced shell execution with safety controls",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  betterbash "ls -la"
  betterbash "git status --porcelain"
  betterbash "find . -name '*.py' | head -5"
        """
    )
    parser.add_argument("command", nargs="?", help="Shell command to execute")
    parser.add_argument("--timeout", "-t", type=int, default=30, 
                       help="Command timeout in seconds (default: 30)")
    parser.add_argument("--list", action="store_true",
                       help="List all allowed commands")
    parser.add_argument("--add", metavar="COMMAND",
                       help="Add a command to the allowlist")
    
    args = parser.parse_args()
    
    if args.list:
        allowed_commands = load_allowed_commands()
        print("Allowed commands:")
        for cmd in sorted(allowed_commands):
            print(f"  {cmd}")
        return 0
    
    if args.add:
        return add_command_to_allowlist(args.add)
    
    if not args.command:
        parser.print_help()
        return 1
        
    return execute_command(args.command, args.timeout)

if __name__ == "__main__":
    sys.exit(main())