bantam-cpp/tools.cpp

/* tools.cpp — bash/read/write tools and oversized result offloading. */
#include "bantam.hpp"
#include <algorithm>
#include <array>
#include <cstring>
#include <fstream>
#include <cerrno>
#include <dirent.h>
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>

/* ----------------------------- shell ---------------------------------- */
std::string shell_with_workdir(std::string_view cmd, std::string_view workdir,
                               int timeout, volatile int *interrupted) {
  std::string fcmd = filter_text(cmd);
  if (timeout <= 0) timeout = 120;

  int pipefd[2];
  if (pipe(pipefd) != 0) return "[tool error: pipe failed]";
  pid_t pid = fork();
  if (pid < 0) {
    close(pipefd[0]);
    close(pipefd[1]);
    return "[tool error: fork failed]";
  }
  if (pid == 0) {
    /* child */
    dup2(pipefd[1], STDOUT_FILENO);
    dup2(pipefd[1], STDERR_FILENO);
    close(pipefd[0]);
    close(pipefd[1]);
    if (!workdir.empty()) chdir(std::string(workdir).c_str());
    execl("/bin/bash", "bash", "-c", fcmd.c_str(), (char *)nullptr);
    _exit(127);
  }
  close(pipefd[1]);

  int flags = fcntl(pipefd[0], F_GETFL, 0);
  fcntl(pipefd[0], F_SETFL, flags | O_NONBLOCK);

  std::string out;
  struct pollfd pfd; pfd.fd = pipefd[0]; pfd.events = POLLIN;
  time_t deadline = time(nullptr) + timeout;
  bool killed = false;
  while (true) {
    if (*interrupted) break;
    int remain = (int)(deadline - time(nullptr));
    if (remain <= 0) { killed = true; break; }
    int pr = poll(&pfd, 1, remain * 1000);
    if (pr < 0) { if (errno == EINTR) continue; break; }
    if (pr == 0) { killed = true; break; }
    char buf[4096];
    ssize_t n = read(pipefd[0], buf, sizeof buf);
    if (n > 0) out.append(buf, (size_t)n);
    else if (n == 0) break; /* EOF: child exited */
    else if (errno == EAGAIN || errno == EWOULDBLOCK) continue;
    else break;
  }
  if (*interrupted || killed) {
    kill(pid, SIGKILL);
    /* drain remaining output briefly */
    struct pollfd dpfd; dpfd.fd = pipefd[0]; dpfd.events = POLLIN;
    for (int i = 0; i < 20; i++) {
      int pr = poll(&dpfd, 1, 50);
      if (pr <= 0) break;
      char buf[4096];
      ssize_t n = read(pipefd[0], buf, sizeof buf);
      if (n > 0) out.append(buf, (size_t)n);
      else break;
    }
  }
  int status = 0;
  while (waitpid(pid, &status, 0) < 0 && errno == EINTR) {}
  close(pipefd[0]);

  std::string filtered = filter_text(out);
  /* trim trailing whitespace */
  while (!filtered.empty() &&
         (filtered.back() == '\n' || filtered.back() == '\r' ||
          filtered.back() == ' '  || filtered.back() == '\t'))
    filtered.pop_back();

  if (*interrupted) {
    return filtered + "\n\nexit: -1";
  }
  if (killed) {
    return sfmt("%s\n\n[shell timeout after %ds]\nexit: -1", filtered.c_str(), timeout);
  }
  int code = 0;
  if (WIFEXITED(status)) code = WEXITSTATUS(status);
  else if (WIFSIGNALED(status)) code = -1;
  return sfmt("%s\n\nexit: %d", filtered.c_str(), code);
}

/* ----------------------------- read ----------------------------------- */
std::string read_file_or_dir(std::string_view path, int offset, int limit, int *style) {
  std::string p(path);
  /* trim */
  size_t s = 0;
  while (s < p.size() && (p[s] == ' ' || p[s] == '\t')) s++;
  p = p.substr(s);
  if (p.empty()) { *style = 31; return "[tool error: read requires 'filePath' parameter]"; }
  struct stat st;
  if (stat(p.c_str(), &st) != 0) {
    *style = 31;
    return sfmt("[tool error: read %s: %s]", p.c_str(), strerror(errno));
  }
  if (S_ISDIR(st.st_mode)) {
    std::vector<std::string> raw_names;
    {
      DIR *dir = opendir(p.c_str());
      if (!dir) { *style = 31; return "[tool error: read (directory listing failed)]"; }
      struct dirent *de;
      while ((de = readdir(dir))) {
        std::string nm = de->d_name;
        if (nm == "." || nm == "..") continue;
        raw_names.push_back(std::move(nm));
      }
      closedir(dir);
    }
    std::sort(raw_names.begin(), raw_names.end());
    std::string b;
    for (auto &nm : raw_names) {
      std::string full = p + "/" + nm;
      struct stat ds;
      bool isdir = (stat(full.c_str(), &ds) == 0 && S_ISDIR(ds.st_mode));
      b += nm;
      if (isdir) b += '/';
      b += '\n';
    }
    if (!b.empty()) b.pop_back(); /* trailing newline */
    *style = 2;
    return b;
  }
  /* regular file */
  std::ifstream f(p, std::ios::binary);
  if (!f) { *style = 31; return sfmt("[tool error: read %s: %s]", p.c_str(), strerror(errno)); }
  if (offset < 1) offset = 1;
  if (limit <= 0) limit = 2000;
  std::string b;
  std::string line;
  int line_num = 0, read_count = 0;
  while (std::getline(f, line)) {
    line_num++;
    if (line_num < offset) continue;
    read_count++;
    while (!line.empty() && (line.back() == '\n' || line.back() == '\r')) line.pop_back();
    std::string_view lview = utf8_truncate(line, 2000);
    str_appendf(b, "%d: ", line_num);
    b.append(lview);
    b += '\n';
    if (read_count >= limit) break;
  }
  if (!b.empty()) b.pop_back(); /* trailing newline */
  *style = 2;
  return b;
}

/* ----------------------------- write ---------------------------------- */
std::string write_file(std::string_view path, int offset, int del_bytes,
                       std::string_view content, int *err) {
  *err = 0;
  std::string p(path);
  size_t s = 0;
  while (s < p.size() && (p[s] == ' ' || p[s] == '\t')) s++;
  p = p.substr(s);
  if (p.empty()) { *err = 1; return "[tool error: path is required]"; }
  if (offset < 0) offset = 0;
  if (del_bytes < 0) del_bytes = 0;

  std::string data;
  struct stat st;
  if (stat(p.c_str(), &st) == 0) {
    if (S_ISDIR(st.st_mode)) {
      *err = 1;
      return sfmt("[tool error: write %s: Is a directory]", p.c_str());
    }
    std::ifstream f(p, std::ios::binary);
    if (!f) { *err = 1; return sfmt("[tool error: write %s: %s]", p.c_str(), strerror(errno)); }
    data.assign((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
  } else {
    auto slash = p.find_last_of('/');
    if (slash != std::string::npos) mkdir_p(std::string_view(p).substr(0, slash));
    else mkdir_p(".");
  }
  if ((size_t)offset > data.size()) {
    data.resize((size_t)offset, '\0');
  }
  size_t end_del = (size_t)offset + (size_t)del_bytes;
  size_t prefix_len = (size_t)offset;
  size_t suffix_len = (end_del < data.size()) ? (data.size() - end_del) : 0;
  size_t clen = content.size();
  size_t newlen = prefix_len + clen + suffix_len;
  std::string nd;
  nd.reserve(newlen);
  nd.append(data, 0, prefix_len);
  nd.append(content);
  if (suffix_len) nd.append(data, end_del, suffix_len);
  {
    std::ofstream out(p, std::ios::binary | std::ios::trunc);
    if (!out) { *err = 1; return sfmt("[tool error: write %s: %s]", p.c_str(), strerror(errno)); }
    out.write(nd.data(), (std::streamsize)nd.size());
  }
  return sfmt("Successfully wrote %zu bytes to %s", clen, p.c_str());
}

/* ----------------------------- offload -------------------------------- */
std::string offload_tool_result(std::string_view res, int max_res,
                                std::vector<std::string> &tmps) {
  if (max_res <= 0) max_res = 65536;
  if (res.size() <= (size_t)max_res) return std::string(res);
  std::string dir = tool_res_dir();
  std::string tpl = dir + "/toolres-XXXXXX";
  int fd = mkstemp(tpl.data());
  if (fd < 0) return std::string(res);
  size_t written = 0;
  bool write_ok = true;
  while (written < res.size()) {
    ssize_t w = write(fd, res.data() + written, res.size() - written);
    if (w < 0) {
      if (errno == EINTR) continue;
      write_ok = false;
      break;
    }
    written += (size_t)w;
  }
  close(fd);
  if (!write_ok) { unlink(tpl.c_str()); return std::string(res); }
  tmps.push_back(tpl);
  return sfmt("[tool result too large for context: %zu bytes (limit %d). The full output was saved to %s. Read it with the read tool; to read it partially, use offset and limit. The file is %zu bytes long.]",
              res.size(), max_res, tpl.c_str(), res.size());
}
void cleanup_temps(std::vector<std::string> &tmps) {
  for (auto &t : tmps) unlink(t.c_str());
  tmps.clear();
}