LaodingBot/internal/tools/shelltool/shelltool.go

package shelltool

import (
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"laodingbot/internal/logger"
)

type Tool struct {
	allowedCommands map[string]struct{}
	workDir         string
	timeout         time.Duration
	log             *logger.Logger
}

func New(allowed []string, workDir string, timeout time.Duration, log *logger.Logger) *Tool {
	set := make(map[string]struct{}, len(allowed))
	for _, c := range allowed {
		cmd := strings.TrimSpace(c)
		if cmd != "" {
			set[cmd] = struct{}{}
		}
	}
	absDir, err := filepath.Abs(workDir)
	if err != nil {
		absDir = workDir
	}
	if timeout <= 0 {
		timeout = 15 * time.Second
	}
	if log != nil {
		log.Infof("shell tool initialized allowed_commands=%d work_dir=%s timeout=%s", len(set), absDir, timeout)
	}
	return &Tool{allowedCommands: set, workDir: absDir, timeout: timeout, log: log}
}

func (t *Tool) Name() string { return "shell" }

func (t *Tool) Description() string {
	return "Execute allowlisted shell commands in Linux"
}

func (t *Tool) Call(ctx context.Context, input string) (string, error) {
	trimmed := strings.TrimSpace(input)
	if trimmed == "" {
		if t.log != nil {
			t.log.Warnf("shell tool rejected empty command")
		}
		return "", fmt.Errorf("empty command")
	}

	parts := strings.Fields(trimmed)
	base := parts[0]
	if _, ok := t.allowedCommands[base]; !ok {
		if t.log != nil {
			t.log.Warnf("shell command denied command=%s full_command=%q", base, trimmed)
		}
		return "", fmt.Errorf("command not allowed: %s", base)
	}
	if t.log != nil {
		t.log.Infof("shell command start command=%s args=%d full_command=%q", base, len(parts)-1, trimmed)
	}

	runCtx, cancel := context.WithTimeout(ctx, t.timeout)
	defer cancel()

	cmd := exec.CommandContext(runCtx, base, parts[1:]...)
	cmd.Dir = t.workDir
	out, err := cmd.CombinedOutput()
	if err != nil {
		if t.log != nil {
			t.log.Errorf("shell command failed command=%s full_command=%q err=%v output_bytes=%d output=%q", base, trimmed, err, len(out), string(out))
		}
		return string(out), err
	}
	if t.log != nil {
		t.log.Infof("shell command success command=%s full_command=%q output_bytes=%d output=%q", base, trimmed, len(out), string(out))
	}
	return string(out), nil
}