byteclaw/internal/tools/shell.go

package tools

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/crypto/byteclaw/internal/llm"
)

// Shell runs allowlisted commands inside the workspace WITHOUT a shell
// interpreter: the command line is parsed here, every command of a chain
// (&&, ||, ;, |) is checked against the allowlist, and pipes are wired
// directly between processes. Expansions ($, `...`, <(...)), redirections,
// background jobs and escapes are rejected. The environment is replaced
// by a minimal one, so no API keys or tokens leak into child processes.
type Shell struct {
	Dir       string
	Allowed   map[string]bool
	Timeout   time.Duration
	MaxOutput int
}

// NeverAllowed commands would re-open an interpreter or escape the sandbox; they
// cannot be allowlisted.
var NeverAllowed = map[string]bool{
	"sh": true, "bash": true, "zsh": true, "dash": true, "ksh": true, "fish": true, "busybox": true,
	"env": true, "xargs": true, "sudo": true, "su": true, "doas": true, "nohup": true, "timeout": true,
	"eval": true, "exec": true, "command": true, "nice": true, "chroot": true, "setsid": true, "watch": true,
	"python": true, "python3": true, "perl": true, "ruby": true, "node": true, "php": true, "lua": true, "awk": true,
	"gawk": true, "sed": true, "ssh": true, "scp": true, "curl": true, "wget": true, "nc": true, "socat": true,
	"git": true, "make": true, "vim": true, "vi": true, "less": true, "more": true, "man": true,
}

const safePath = "/usr/local/bin:/usr/bin:/bin"

func NewShell(dir string, allowed []string, timeout time.Duration, maxOutput int) (*Shell, error) {
	m := map[string]bool{}
	for _, a := range allowed {
		a = strings.TrimSpace(a)
		if a == "" {
			continue
		}
		if strings.ContainsRune(a, '/') || NeverAllowed[a] {
			return nil, fmt.Errorf("command %q cannot be allowlisted", a)
		}
		m[a] = true
	}
	if len(m) == 0 {
		return nil, errors.New("shell tool enabled but allowed_commands is empty")
	}
	return &Shell{Dir: dir, Allowed: m, Timeout: timeout, MaxOutput: maxOutput}, nil
}

func (s *Shell) Spec() llm.ToolSpec {
	names := make([]string, 0, len(s.Allowed))
	for n := range s.Allowed {
		names = append(names, n)
	}
	return llm.ToolSpec{
		Name: "shell",
		Description: "Run a command in the workspace directory (no shell interpreter). Allowed commands: " +
			strings.Join(sortStrings(names), ", ") + ". Supports && || ; and |. No $, backticks, redirections or absolute paths.",
		Parameters: obj(map[string]any{"command": map[string]any{"type": "string"}}, "command"),
	}
}

func sortStrings(s []string) []string {
	for i := 1; i < len(s); i++ {
		for j := i; j > 0 && s[j] < s[j-1]; j-- {
			s[j], s[j-1] = s[j-1], s[j]
		}
	}
	return s
}

// ---- parsing

type step struct {
	op       string     // operator preceding this pipeline: "", "&&", "||", ";"
	pipeline [][]string // commands connected with |
}

var ErrRejected = errors.New("rejected")

func rejectf(format string, a ...any) error {
	return fmt.Errorf("%w: %s", ErrRejected, fmt.Sprintf(format, a...))
}

// Parse splits a command line into checked steps.
func (s *Shell) Parse(line string) ([]step, error) {
	if strings.TrimSpace(line) == "" {
		return nil, rejectf("empty command")
	}
	for _, bad := range []string{"$", "`", "<", ">", "\\", "\n", "\r", "\x00"} {
		if strings.Contains(line, bad) {
			return nil, rejectf("forbidden character %q (no expansion, substitution or redirection)", bad)
		}
	}
	var (
		steps  []step
		cur    []string // current command tokens
		pipe   [][]string
		tok    strings.Builder
		inTok  bool
		quote  rune
		nextOp = ""
		runes  = []rune(line)
	)
	flushTok := func() {
		if inTok {
			cur = append(cur, tok.String())
			tok.Reset()
			inTok = false
		}
	}
	endCmd := func() error {
		flushTok()
		if len(cur) == 0 {
			return rejectf("empty command in chain")
		}
		pipe = append(pipe, cur)
		cur = nil
		return nil
	}
	endPipeline := func(op string) error {
		if err := endCmd(); err != nil {
			return err
		}
		steps = append(steps, step{op: nextOp, pipeline: pipe})
		pipe, nextOp = nil, op
		return nil
	}
	for i := 0; i < len(runes); i++ {
		c := runes[i]
		if quote != 0 {
			if c == quote {
				quote = 0
			} else {
				tok.WriteRune(c)
			}
			continue
		}
		switch {
		case c == '\'' || c == '"':
			quote, inTok = c, true
		case c == ' ' || c == '\t':
			flushTok()
		case c == '&':
			if i+1 < len(runes) && runes[i+1] == '&' {
				if err := endPipeline("&&"); err != nil {
					return nil, err
				}
				i++
			} else {
				return nil, rejectf("background jobs (&) are not allowed")
			}
		case c == '|':
			if i+1 < len(runes) && runes[i+1] == '|' {
				if err := endPipeline("||"); err != nil {
					return nil, err
				}
				i++
			} else if err := endCmd(); err != nil {
				return nil, err
			}
		case c == ';':
			if err := endPipeline(";"); err != nil {
				return nil, err
			}
		case c == '(' || c == ')' || c == '{' || c == '}':
			return nil, rejectf("subshells and groups are not allowed")
		default:
			tok.WriteRune(c)
			inTok = true
		}
	}
	if quote != 0 {
		return nil, rejectf("unterminated quote")
	}
	if err := endPipeline(""); err != nil {
		return nil, err
	}
	for _, st := range steps {
		for _, argv := range st.pipeline {
			if err := s.check(argv); err != nil {
				return nil, err
			}
		}
	}
	return steps, nil
}

func (s *Shell) check(argv []string) error {
	name := argv[0]
	if strings.ContainsRune(name, '/') || !s.Allowed[name] || NeverAllowed[name] {
		return rejectf("command %q is not allowed", name)
	}
	for _, a := range argv[1:] {
		if strings.HasPrefix(a, "/") || strings.HasPrefix(a, "~") || strings.Contains(a, "=/") || strings.Contains(a, "=~") {
			return rejectf("absolute or home paths are not allowed: %q", a)
		}
		for _, part := range strings.FieldsFunc(a, func(r rune) bool { return r == '/' || r == '=' || r == ':' || r == ',' }) {
			if part == ".." {
				return rejectf("path traversal is not allowed: %q", a)
			}
		}
		if name == "find" && (strings.HasPrefix(a, "-exec") || strings.HasPrefix(a, "-ok") || a == "-delete" || strings.HasPrefix(a, "-fprint")) {
			return rejectf("find %s is not allowed", a)
		}
	}
	return nil
}

// ---- execution

type limitedBuf struct {
	mu        sync.Mutex
	buf       bytes.Buffer
	max       int
	truncated bool
}

func (l *limitedBuf) Write(p []byte) (int, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if room := l.max - l.buf.Len(); room < len(p) {
		if room > 0 {
			l.buf.Write(p[:room])
		}
		l.truncated = true
		return len(p), nil
	}
	return l.buf.Write(p)
}

func lookPath(name string) (string, error) {
	for _, dir := range filepath.SplitList(safePath) {
		p := filepath.Join(dir, name)
		if fi, err := os.Stat(p); err == nil && !fi.IsDir() && fi.Mode()&0o111 != 0 {
			return p, nil
		}
	}
	return "", fmt.Errorf("command %q not found", name)
}

func (s *Shell) env() []string {
	return []string{"PATH=" + safePath, "HOME=" + s.Dir, "LANG=C.UTF-8", "LC_ALL=C.UTF-8", "TERM=dumb"}
}

func (s *Shell) Call(ctx context.Context, raw json.RawMessage) (string, error) {
	var a struct {
		Command string `json:"command"`
	}
	if err := json.Unmarshal(raw, &a); err != nil {
		return "", err
	}
	steps, err := s.Parse(a.Command)
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(s.Dir, 0o750); err != nil {
		return "", err
	}
	ctx, cancel := context.WithTimeout(ctx, s.Timeout)
	defer cancel()
	out := &limitedBuf{max: s.MaxOutput}
	code := 0
	for _, st := range steps {
		if (st.op == "&&" && code != 0) || (st.op == "||" && code == 0) {
			continue
		}
		code, err = s.runPipeline(ctx, st.pipeline, out)
		if err != nil {
			return "", err
		}
	}
	res := out.buf.String()
	if out.truncated {
		res += "\n[output truncated]"
	}
	if ctx.Err() == context.DeadlineExceeded {
		return res + fmt.Sprintf("\n[timeout after %s]", s.Timeout), nil
	}
	return fmt.Sprintf("%s\n[exit %d]", strings.TrimRight(res, "\n"), code), nil
}

func (s *Shell) runPipeline(ctx context.Context, pipeline [][]string, out io.Writer) (int, error) {
	cmds := make([]*exec.Cmd, len(pipeline))
	for i, argv := range pipeline {
		bin, err := lookPath(argv[0])
		if err != nil {
			return 127, nil
		}
		c := exec.CommandContext(ctx, bin, argv[1:]...)
		c.Dir, c.Env, c.Stderr = s.Dir, s.env(), out
		c.WaitDelay = time.Second
		cmds[i] = c
	}
	for i := 0; i < len(cmds)-1; i++ {
		r, err := cmds[i].StdoutPipe()
		if err != nil {
			return 0, err
		}
		cmds[i+1].Stdin = r
	}
	cmds[len(cmds)-1].Stdout = out
	for _, c := range cmds {
		if err := c.Start(); err != nil {
			return 126, nil
		}
	}
	code := 0
	for i, c := range cmds {
		err := c.Wait()
		if i == len(cmds)-1 {
			var ee *exec.ExitError
			if errors.As(err, &ee) {
				code = ee.ExitCode()
			} else if err != nil {
				code = 1
			}
		}
	}
	return code, nil
}