tools.go

package main

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

type ToolResult struct {
	Success bool
	Output  string
}

var allowedCommands = []string{"go", "ls", "cat", "pwd", "echo", "grep"}

// Commands that will never be allowed, even if the base executable is in the allowed list.
var blockedCommands = []string{"go run", "./server", "go run ."}

func isAllowed(cmd string) bool {
	cmd = strings.TrimSpace(cmd)
	// Check blocked patterns first
	for _, blocked := range blockedCommands {
		if strings.HasPrefix(cmd, blocked) || strings.Contains(cmd, blocked) {
			return false
		}
	}
	parts := strings.Fields(cmd)
	if len(parts) == 0 {
		return false
	}
	base := parts[0]
	for _, allowed := range allowedCommands {
		if base == allowed {
			return true
		}
	}
	return false
}

func ExecuteToolFromCall(tc ToolCall) ToolResult {
	switch tc.Name {
	case "write_file":
		path := tc.Args["path"]
		content := tc.Args["content"]
		err := os.WriteFile(path, []byte(content), 0644)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}

		return ToolResult{Success: true, Output: fmt.Sprintf("wrote %d bytes to %s", len(content), path)}

	case "read_file":
		path := tc.Args["path"]
		data, err := os.ReadFile(path)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		return ToolResult{Success: true, Output: string(data)}

	case "list_dir":
		dir := tc.Args["dir"]
		if dir == "" {
			dir = "."
		}
		entries, err := os.ReadDir(dir)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		var names []string
		for _, e := range entries {
			if e.IsDir() {
				names = append(names, e.Name()+"/")
			} else {
				names = append(names, e.Name())
			}
		}
		return ToolResult{Success: true, Output: strings.Join(names, "\n")}

	case "search_content":
		pattern := tc.Args["pattern"]
		searchPath := tc.Args["path"]
		if searchPath == "" {
			searchPath = "."
		}
		matches, err := grep(searchPath, pattern)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		if len(matches) == 0 {
			return ToolResult{Success: true, Output: "no matches found"}
		}
		return ToolResult{Success: true, Output: strings.Join(matches, "\n")}

	case "edit_file":
		path := tc.Args["path"]
		old := tc.Args["old"]
		new := tc.Args["new"]
		data, err := os.ReadFile(path)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		content := string(data)
		if !strings.Contains(content, old) {
			return ToolResult{Success: false, Output: fmt.Sprintf("old string not found in %s", path)}
		}
		updated := strings.Replace(content, old, new, 1) // replace first occurrence
		err = os.WriteFile(path, []byte(updated), 0644)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		return ToolResult{Success: true, Output: fmt.Sprintf("edited %s: replaced 1 occurrence", path)}

	case "replace_all":
		path := tc.Args["path"]
		old := tc.Args["old"]
		new := tc.Args["new"]
		data, err := os.ReadFile(path)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		content := string(data)
		count := strings.Count(content, old)
		if count == 0 {
			return ToolResult{Success: false, Output: "old string not found"}
		}
		updated := strings.ReplaceAll(content, old, new)
		err = os.WriteFile(path, []byte(updated), 0644)
		if err != nil {
			return ToolResult{Success: false, Output: err.Error()}
		}
		return ToolResult{Success: true, Output: fmt.Sprintf("replaced %d occurrences in %s", count, path)}

	case "run_command":
		cmdStr := tc.Args["command"]
		if !isAllowed(cmdStr) {
			return ToolResult{Success: false, Output: fmt.Sprintf("command not allowed: %s", cmdStr)}
		}
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		cmd := exec.CommandContext(ctx, "sh", "-c", cmdStr)
		out, err := cmd.CombinedOutput()
		if err != nil {
			if ctx.Err() == context.DeadlineExceeded {
				return ToolResult{Success: false, Output: "command timed out after 5s"}
			}
			return ToolResult{Success: false, Output: fmt.Sprintf("command failed: %s\n%s", err, string(out))}
		}
		return ToolResult{Success: true, Output: string(out)}

	case "finish":
		return ToolResult{Success: true, Output: "task completed"}

	default:
		return ToolResult{Success: false, Output: fmt.Sprintf("unknown tool: %s", tc.Name)}
	}
}

func grep(root, pattern string) ([]string, error) {
	var matches []string
	err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return nil
		}
		if d.IsDir() {
			return nil
		}
		data, err := os.ReadFile(path)
		if err != nil {
			return nil
		}
		for i, line := range strings.Split(string(data), "\n") {
			if strings.Contains(line, pattern) {
				matches = append(matches, fmt.Sprintf("%s:%d: %s", path, i+1, line))
			}
		}
		return nil
	})
	return matches, err
}