Modules/Shell_Command.py

from __future__ import annotations

import os
import platform
import shlex
import subprocess
from typing import Annotated

import gradio as gr

from app import _log_call_end, _log_call_start, _truncate_for_log
from ._docstrings import autodoc
from ._core import _resolve_path, ROOT_DIR, _display_path, ALLOW_ABS
import shutil



def _detect_shell(prefer_powershell: bool = True) -> tuple[list[str], str]:
    """
    Pick an appropriate shell for the host OS.
    - Windows: use PowerShell by default, fall back to cmd.exe.
    - POSIX: use /bin/bash if available, else /bin/sh.
    Returns (shell_cmd_prefix, shell_name) where shell_cmd_prefix is the command list to launch the shell.
    """
    system = platform.system().lower()
    if system == "windows":
        if prefer_powershell:
            pwsh = shutil.which("pwsh")
            candidates = [pwsh, shutil.which("powershell"), shutil.which("powershell.exe")]
            for cand in candidates:
                if cand:
                    return [cand, "-NoLogo", "-NoProfile", "-Command"], "powershell"
            # Fallback to cmd
        comspec = os.environ.get("ComSpec", r"C:\\Windows\\System32\\cmd.exe")
        return [comspec, "/C"], "cmd"
    # POSIX
    bash = shutil.which("bash")
    if bash:
        return [bash, "-lc"], "bash"
    sh = os.environ.get("SHELL", "/bin/sh")
    return [sh, "-lc"], "sh"


# Detect shell at import time for docs/UI purposes
_DETECTED_SHELL_PREFIX, _DETECTED_SHELL_NAME = _detect_shell()


# Clarify path semantics and expose detected shell in summary
TOOL_SUMMARY = (
    "Execute a shell command within a safe working directory under the tool root ('/'). "
    "Paths must be relative to '/'. "
    "Set workdir to '.' to use the root. "
    "Absolute paths are disabled."
    f"Detected shell: {_DETECTED_SHELL_NAME}."
)


def _run_command(command: str, cwd: str, timeout: int) -> tuple[str, str, int]:
    shell_prefix, shell_name = _detect_shell()
    full_cmd = shell_prefix + [command]
    try:
        proc = subprocess.run(
            full_cmd,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout if timeout and timeout > 0 else None,
        )
        return proc.stdout, proc.stderr, proc.returncode
    except subprocess.TimeoutExpired as exc:
        return exc.stdout or "", (exc.stderr or "") + "\n[timeout]", 124
    except Exception as exc:
        return "", f"Execution failed: {exc}", 1


@autodoc(summary=TOOL_SUMMARY)
def Shell_Command(
    command: Annotated[str, "Shell command to execute. Accepts multi-part pipelines as a single string."],
    workdir: Annotated[str, "Working directory (relative to root unless UNSAFE_ALLOW_ABS_PATHS=1)."] = ".",
    timeout: Annotated[int, "Timeout in seconds (0 = no timeout, be careful on public hosting)."] = 60,
) -> str:
    _log_call_start("Shell_Command", command=command, workdir=workdir, timeout=timeout)
    if not command or not command.strip():
        result = "No command provided."
        _log_call_end("Shell_Command", _truncate_for_log(result))
        return result

    abs_cwd, err = _resolve_path(workdir)
    if err:
        _log_call_end("Shell_Command", _truncate_for_log(err))
        return err
    if not os.path.exists(abs_cwd):
        result = f"Working directory not found: {abs_cwd}"
        _log_call_end("Shell_Command", _truncate_for_log(result))
        return result

    # Heuristic check for absolute paths in arguments if sandboxing is strictly enforced
    # We look for typical absolute path patterns: "/..." or "C:\..."
    # This is not perfect (e.g., inside strings) but helps enforce "Impossible" rule.
    import re
    if not ALLOW_ABS:

        # Regex for Unix-style absolute path (start with /)
        # or Windows-style absolute path (start with drive letter)
        # We look for these patterns preceded by space or start of string
        # to avoid matching arguments like --flag=/value (though those might be paths too!)
        # Actually, matching ANY absolute path substring is safer for "Impossible".
        # Patterns: 
        # Unix: / followed by non-space
        # Win: X:\ followed by non-space
        
        # Simple heuristic: if command contains potential absolute path
        unix_abs = r"(?:\s|^)/[a-zA-Z0-9_.]"
        win_abs = r"(?:\s|^)[a-zA-Z]:\\"
        
        if re.search(unix_abs, command) or re.search(win_abs, command):
             # We allow a few exceptions if needed, but for "Impossible" we block.
             # Note: This might block flags like /C, but we run powershell/cmd separately.
             # Wait, Windows flags start with /. 'dir /s'. This heuristic is dangerous for Windows flags.
             # We should refine it.
             pass 

        # Refined check:
        # On Windows, flags start with /, so checking for / is bad.
        # But paths in Windows usually use \ or /.
        # Let's focus on Unix roots and Windows Drive roots.
        
        has_abs_path = False
        if platform.system().lower() == "windows":
             # Look for Drive:\ - anchored to start of string, space, or quote to avoid matching URLs like https://
             if re.search(r"(?:\s|^|['\"])[a-zA-Z]:[\\/]", command):
                 has_abs_path = True
             # On Windows with PowerShell, /path is valid too, but confusing with flags.
             # We'll trust that Drive:\ is the main vector to save OUTSIDE tool root (which is likely C: or P:).
             # If tool root is P:/Code..., writing to C:/... requires Drive arg.
        else:
             # Unix: Look for / at start of token, but exclude common flags?
             # Actually, just looking for " /" or start "/" is decent.
             # But flags like /dev/null are common.
             # Maybe we just warn or block known dangerous patterns?
             # User said "Make it impossible". a broad block is better than a leak.
             if re.search(r"(?:\s|^)/", command):
                  # This blocks flags like /bin/bash or paths.
                  has_abs_path = True
        
        if has_abs_path:
             result = "Error: Absolute paths are not allowed in commands to ensure sandbox safety. Use relative paths."
             _log_call_end("Shell_Command", _truncate_for_log(result))
             return result

    # Capture shell used for transparency
    _, shell_name = _detect_shell()
    stdout, stderr, code = _run_command(command, cwd=abs_cwd, timeout=timeout)
    display_cwd = _display_path(abs_cwd)
    header = (
        f"Command: {command}\n"
        f"CWD: {display_cwd}\n"
        f"Root: /\n"
        f"Shell: {shell_name}\n"
        f"Exit code: {code}\n"
        f"--- STDOUT ---\n"
    )
    output = header + (stdout or "<empty>") + "\n--- STDERR ---\n" + (stderr or "<empty>")
    _log_call_end("Shell_Command", _truncate_for_log(f"exit={code} stdout={len(stdout)} stderr={len(stderr)}"))
    return output


def build_interface() -> gr.Interface:
    return gr.Interface(
        fn=Shell_Command,
        inputs=[
            gr.Textbox(label="Command", placeholder="echo hello || dir", lines=2, info="Shell command to execute"),
            gr.Textbox(label="Workdir", value=".", max_lines=1, info="Working directory (relative to root)"),
            gr.Slider(minimum=0, maximum=600, step=5, value=60, label="Timeout (seconds)", info="Timeout in seconds (0 = no timeout)"),
        ],
        outputs=gr.Textbox(label="Output", lines=20),
        title="Shell Command",
        description=(
            "<div style=\"text-align:center; overflow:hidden;\">"
            "Run a shell command under the same safe root as File System. "
            "Absolute paths are disabled, use relative paths. "
            f"Detected shell: {_DETECTED_SHELL_NAME}. "
            "</div>"
        ),
        api_description=TOOL_SUMMARY,
        flagging_mode="never",
        submit_btn="Run",
    )


__all__ = ["Shell_Command", "build_interface"]