drharness/tools/shell_tool.cpp

#include "shell_tool.hpp"
#include "json_types.hpp"
#include <array>
#include <cstdio>
#include <iostream>
#include <stdexcept>
#include <string>
#include <sys/wait.h>
#include <sys/poll.h>
#include <unistd.h>
#include <signal.h>
#include <cstring>

Tool make_shell_tool(bool safe) {
    std::string tool_name = safe ? "run_shell" : "run_dangerous_shell";

    return Tool{
        .name = tool_name,
        .description = "Run a shell command and return its stdout and stderr. "
                       "Use for inspecting files, running programs, or querying the system.",
        .parameters_schema = {
            {"type", "object"},
            {"properties", {
                {"command", {
                    {"type", "string"},
                    {"description", "The shell command to execute. Runs via /bin/sh."}
                }},
                {"timeout_seconds", {
                    {"type", "integer"},
                    {"description", "Maximum seconds to wait for the command to complete. Defaults to 30."}
                }}
            }},
            {"required", {"command"}}
        },
        .execute = create_tool_execute<ShellToolParams>([safe](const ShellToolParams& params) -> std::string {

            if (safe) {
                std::cerr << "Approve command: " << params.command << " ? [y/N]: ";
                char response = 0;
                std::cin.get(response);
                if (response != 'y' && response != 'Y') {
                    return "Execution cancelled by user.";
                }
            }

            // 1. Create a pipe for stdout
            int pipefd[2];
            if (pipe(pipefd) < 0) {
                return "Error: pipe creation failed";
            }

            pid_t pid = fork();
            if (pid < 0) {
                close(pipefd[0]); close(pipefd[1]);
                return "Error: fork failed";
            }

            if (pid == 0) {
                // Child process
                // Set process group so we can kill the whole tree
                setpgid(0, 0);
                
                // Close read end (child writes to the pipe)
                close(pipefd[0]);
                // Duplicate pipe write end to stdout
                if (dup2(pipefd[1], STDOUT_FILENO) < 0) _exit(1);
                close(pipefd[1]);
                
                // Merge stderr to stdout so it doesn't block
                dup2(STDOUT_FILENO, STDERR_FILENO);

                // Execute command
                std::string sh = "/bin/sh";
                char* argv[] = { const_cast<char*>(sh.c_str()), const_cast<char*>("-c"), 
                                 const_cast<char*>(params.command.c_str()), nullptr };
                execvp(argv[0], argv);
                _exit(127); // exec failed
            }

            // Parent process
            // Close write end (parent reads from pipe)
            close(pipefd[1]);

            std::string output;
            std::array<char, 4096> buf;
            struct pollfd pfd;
            pfd.fd = pipefd[0];
            pfd.events = POLLIN;

            bool timed_out = false;
            auto start_time = std::chrono::steady_clock::now();

            while (true) {
                // 1. Poll for stdout or process exit
                // Wait 1 sec before we check timeout
                int res = poll(&pfd, 1, 1000);

                if (res == 0) {
                    if (std::chrono::steady_clock::now() - start_time
                            < std::chrono::seconds(params.timeout_seconds))
                    {
                        continue; // Not timed out yet
                    }
                    // Timeout!
                    // Kill the entire process group to stop sleep/dd/etc.
                    killpg(pid, SIGKILL);
                    timed_out = true;
                    break;
                } else if (res < 0) {
                    if (errno == EINTR) continue; // Restart if interrupted by signal
                    return "Error: poll failed";
                }

                // 2. Check if process finished (POLLHUP)
                if (pfd.revents & POLLHUP) {
                    // Data might still be available in buffer!
                    if (pfd.revents & POLLIN) {
                        // Read ALL remaining data in a loop
                        while (true) {
                            res = read(pfd.fd, buf.data(), buf.size());
                            if (res <= 0) break; // EOF or error
                            output += std::string(buf.data(), res);
                        }
                    }
                    break; 
                }

                // 3. Read data if available - keep reading until no more data
                if (pfd.revents & POLLIN) {
                    while (true) {
                        res = read(pfd.fd, buf.data(), buf.size());
                        if (res < 0) {
                            if (errno == EINTR) continue; // Restart if interrupted by signal
                            break; // Error
                        } else if (res == 0) { // EOF
                            break;
                        }
                        output += std::string(buf.data(), res);
                    }
                    break;
                }
                
            }

            // Check exit code
            int status;
            waitpid(pid, &status, 0);

            if (timed_out) {
                return "[timed out after " + std::to_string(params.timeout_seconds) + "s]\n" + output;
            }

            int code = WIFEXITED(status) ? WEXITSTATUS(status) : -1;
            if (code != 0 && !timed_out) {
                output = "[exit " + std::to_string(code) + "]\n" + output;
            }
            
            close(pipefd[0]);
            return output.empty() ? "(no output)" : output;
        })
    };
}