98 lines
2.6 KiB
Go
98 lines
2.6 KiB
Go
package shelltool
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"time"
|
|
|
|
"laodingbot/internal/logger"
|
|
)
|
|
|
|
type Tool struct {
|
|
allowedCommands map[string]struct{}
|
|
workDir string
|
|
timeout time.Duration
|
|
maxOutputChars int
|
|
log *logger.Logger
|
|
}
|
|
|
|
func New(allowed []string, workDir string, timeout time.Duration, maxOutputChars int, 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 maxOutputChars <= 0 {
|
|
maxOutputChars = 4000
|
|
}
|
|
if log != nil {
|
|
log.Infof("shell tool initialized allowed_commands=%d work_dir=%s timeout=%s max_output_chars=%d", len(set), absDir, timeout, maxOutputChars)
|
|
}
|
|
return &Tool{allowedCommands: set, workDir: absDir, timeout: timeout, maxOutputChars: maxOutputChars, 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()
|
|
outText := string(out)
|
|
if len(outText) > t.maxOutputChars {
|
|
outText = outText[:t.maxOutputChars]
|
|
}
|
|
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), outText)
|
|
}
|
|
if runtime.GOOS == "windows" && strings.Contains(strings.ToLower(err.Error()), "executable file not found") {
|
|
return outText, fmt.Errorf("command not executable in current windows environment: %s", base)
|
|
}
|
|
return outText, 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), outText)
|
|
}
|
|
return outText, nil
|
|
}
|