ZovOS_AI/internal/tools/filesystem/filesystem.go

package filesystem

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/config"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/tools"
)

// RegisterTools registers filesystem-related tools in the registry.
func RegisterTools(registry *tools.Registry, cfg config.FilesystemToolConfig, log *slog.Logger) {
	if !cfg.Enabled {
		return
	}

	registry.Register(tools.ToolDef{
		Tool: ai.Tool{
			Name:        "read_file",
			Description: "Read the contents of a file at the given path.",
			Parameters: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"path": map[string]interface{}{
						"type":        "string",
						"description": "Absolute path to the file to read",
					},
				},
				"required": []string{"path"},
			},
		},
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			var a struct{ Path string `json:"path"` }
			if err := json.Unmarshal(args, &a); err != nil {
				return "", err
			}
			if err := checkAllowed(a.Path, cfg.AllowedPaths); err != nil {
				return "", err
			}
			data, err := os.ReadFile(a.Path)
			if err != nil {
				return "", err
			}
			content := string(data)
			if len(content) > 50000 {
				content = content[:50000] + "\n... (file truncated)"
			}
			return content, nil
		},
	})

	registry.Register(tools.ToolDef{
		Tool: ai.Tool{
			Name:        "write_file",
			Description: "Write content to a file, creating it if it doesn't exist.",
			Parameters: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"path": map[string]interface{}{
						"type":        "string",
						"description": "Absolute path to the file",
					},
					"content": map[string]interface{}{
						"type":        "string",
						"description": "Content to write to the file",
					},
				},
				"required": []string{"path", "content"},
			},
		},
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			var a struct {
				Path    string `json:"path"`
				Content string `json:"content"`
			}
			if err := json.Unmarshal(args, &a); err != nil {
				return "", err
			}
			if err := checkAllowed(a.Path, cfg.AllowedPaths); err != nil {
				return "", err
			}
			dir := filepath.Dir(a.Path)
			if err := os.MkdirAll(dir, 0755); err != nil {
				return "", err
			}
			if err := os.WriteFile(a.Path, []byte(a.Content), 0644); err != nil {
				return "", err
			}
			return fmt.Sprintf("File written successfully: %s (%d bytes)", a.Path, len(a.Content)), nil
		},
	})

	registry.Register(tools.ToolDef{
		Tool: ai.Tool{
			Name:        "list_directory",
			Description: "List contents of a directory.",
			Parameters: map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"path": map[string]interface{}{
						"type":        "string",
						"description": "Path to the directory to list",
					},
				},
				"required": []string{"path"},
			},
		},
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			var a struct{ Path string `json:"path"` }
			if err := json.Unmarshal(args, &a); err != nil {
				return "", err
			}
			if err := checkAllowed(a.Path, cfg.AllowedPaths); err != nil {
				return "", err
			}
			entries, err := os.ReadDir(a.Path)
			if err != nil {
				return "", err
			}
			var sb strings.Builder
			for _, e := range entries {
				info, _ := e.Info()
				typeStr := "file"
				if e.IsDir() {
					typeStr = "dir "
				}
				size := int64(0)
				if info != nil {
					size = info.Size()
				}
				sb.WriteString(fmt.Sprintf("[%s] %s (%d bytes)\n", typeStr, e.Name(), size))
			}
			if sb.Len() == 0 {
				return "(empty directory)", nil
			}
			return sb.String(), nil
		},
	})

	log.Info("filesystem tools registered")
}

func checkAllowed(path string, allowedPaths []string) error {
	if len(allowedPaths) == 0 {
		return nil // no restrictions
	}
	absPath, err := filepath.Abs(path)
	if err != nil {
		return fmt.Errorf("resolving path: %w", err)
	}
	for _, allowed := range allowedPaths {
		if strings.HasPrefix(absPath, allowed) {
			return nil
		}
	}
	return fmt.Errorf("access denied: path %s is not in allowed directories", path)
}