Cyrene/backend/ai-core/internal/tools/file_tool.go

package tools

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

// FileTool provides sandboxed file system operations for the LLM.
// All paths are restricted to a DATA_DIR to prevent directory traversal attacks.
type FileTool struct {
	dataDir string
}

// NewFileTool creates a file operation tool with the given data directory.
func NewFileTool(dataDir string) *FileTool {
	if dataDir == "" {
		dataDir = "/tmp/cyrene_data"
	}
	return &FileTool{dataDir: dataDir}
}

// Definition returns the tool definition for LLM function calling.
func (t *FileTool) Definition() ToolDefinition {
	return ToolDefinition{
		Name:        "file_ops",
		Description: "文件操作工具。在服务端安全沙盒内读写文件、列出目录、检查文件是否存在、删除文件。所有操作限制在数据目录内，无法访问系统文件。",
		Parameters: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"action": map[string]interface{}{
					"type":        "string",
					"enum":        []string{"read", "write", "list", "exists", "delete"},
					"description": "操作类型。read: 读取文件；write: 写入文件（覆盖或创建）；list: 列出目录内容；exists: 检查路径是否存在；delete: 删除文件",
				},
				"path": map[string]interface{}{
					"type":        "string",
					"description": "文件或目录路径（相对于数据目录），如 \"notes/todo.txt\"",
				},
				"content": map[string]interface{}{
					"type":        "string",
					"description": "写入内容（write 操作时必需）",
				},
			},
			"required": []string{"action", "path"},
		},
	}
}

// Execute performs file operations.
func (t *FileTool) Execute(ctx context.Context, arguments map[string]interface{}) (*ToolResult, error) {
	action, ok := arguments["action"].(string)
	if !ok || action == "" {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    "缺少 action 参数",
		}, nil
	}

	relPath, ok := arguments["path"].(string)
	if !ok || relPath == "" {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    "缺少 path 参数",
		}, nil
	}

	safePath, err := t.resolveSafePath(relPath)
	if err != nil {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    err.Error(),
		}, nil
	}

	switch action {
	case "read":
		return t.handleRead(safePath, relPath)
	case "write":
		content, _ := arguments["content"].(string)
		return t.handleWrite(safePath, relPath, content)
	case "list":
		return t.handleList(safePath, relPath)
	case "exists":
		return t.handleExists(safePath, relPath)
	case "delete":
		return t.handleDelete(safePath, relPath)
	default:
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("未知操作: %s，支持: read, write, list, exists, delete", action),
		}, nil
	}
}

// resolveSafePath resolves a relative path and ensures it stays within dataDir.
func (t *FileTool) resolveSafePath(relPath string) (string, error) {
	// Clean the path first
	clean := filepath.Clean(relPath)

	// Ensure data directory exists
	if err := os.MkdirAll(t.dataDir, 0755); err != nil {
		return "", fmt.Errorf("创建数据目录失败: %v", err)
	}

	abs := filepath.Join(t.dataDir, clean)

	// Prevent directory traversal
	realPath, err := filepath.EvalSymlinks(abs)
	if err != nil {
		// If the path doesn't exist yet, we can still check the prefix
		if os.IsNotExist(err) {
			// Ensure the resolved path (without symlinks) is within dataDir
			if !strings.HasPrefix(filepath.Clean(abs), filepath.Clean(t.dataDir)+string(filepath.Separator)) &&
				filepath.Clean(abs) != filepath.Clean(t.dataDir) {
				return "", fmt.Errorf("路径穿越检测: %s 不在允许的数据目录内", relPath)
			}
			return abs, nil
		}
		return "", fmt.Errorf("路径解析失败: %v", err)
	}

	if !strings.HasPrefix(realPath, filepath.Clean(t.dataDir)+string(filepath.Separator)) &&
		realPath != filepath.Clean(t.dataDir) {
		return "", fmt.Errorf("路径穿越检测: %s 不在允许的数据目录内", relPath)
	}

	return realPath, nil
}

// handleRead reads a file, limited to 100KB.
func (t *FileTool) handleRead(absPath, relPath string) (*ToolResult, error) {
	const maxSize = 100 * 1024 // 100KB

	info, err := os.Stat(absPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &ToolResult{
				ToolName: "file_ops",
				Success:  false,
				Error:    fmt.Sprintf("文件不存在: %s", relPath),
			}, nil
		}
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("读取文件失败: %v", err),
		}, nil
	}

	if info.IsDir() {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("路径是目录，不能用 read 操作: %s", relPath),
		}, nil
	}

	if info.Size() > maxSize {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("文件过大 (%d bytes)，超过限制 (%d bytes)", info.Size(), maxSize),
		}, nil
	}

	data, err := os.ReadFile(absPath)
	if err != nil {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("读取文件失败: %v", err),
		}, nil
	}

	return &ToolResult{
		ToolName: "file_ops",
		Success:  true,
		Data:     fmt.Sprintf("文件: %s\n大小: %d bytes\n---\n%s", relPath, len(data), string(data)),
	}, nil
}

// handleWrite writes content to a file.
func (t *FileTool) handleWrite(absPath, relPath, content string) (*ToolResult, error) {
	// Ensure parent directory exists
	dir := filepath.Dir(absPath)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("创建目录失败: %v", err),
		}, nil
	}

	if err := os.WriteFile(absPath, []byte(content), 0644); err != nil {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("写入文件失败: %v", err),
		}, nil
	}

	return &ToolResult{
		ToolName: "file_ops",
		Success:  true,
		Data:     fmt.Sprintf("已写入文件: %s (%d bytes)", relPath, len(content)),
	}, nil
}

// handleList lists directory contents.
func (t *FileTool) handleList(absPath, relPath string) (*ToolResult, error) {
	entries, err := os.ReadDir(absPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &ToolResult{
				ToolName: "file_ops",
				Success:  false,
				Error:    fmt.Sprintf("目录不存在: %s", relPath),
			}, nil
		}
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("读取目录失败: %v", err),
		}, nil
	}

	if len(entries) == 0 {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  true,
			Data:     fmt.Sprintf("目录: %s\n(空目录)", relPath),
		}, nil
	}

	var result strings.Builder
	result.WriteString(fmt.Sprintf("目录: %s\n共 %d 项:\n", relPath, len(entries)))
	for _, entry := range entries {
		icon := "📄"
		if entry.IsDir() {
			icon = "📁"
		}
		info, _ := entry.Info()
		size := ""
		if info != nil && !entry.IsDir() {
			size = fmt.Sprintf(" (%d bytes)", info.Size())
		}
		result.WriteString(fmt.Sprintf("  %s %s%s\n", icon, entry.Name(), size))
	}

	return &ToolResult{
		ToolName: "file_ops",
		Success:  true,
		Data:     result.String(),
	}, nil
}

// handleExists checks whether a path exists.
func (t *FileTool) handleExists(absPath, relPath string) (*ToolResult, error) {
	info, err := os.Stat(absPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &ToolResult{
				ToolName: "file_ops",
				Success:  true,
				Data:     fmt.Sprintf("路径不存在: %s", relPath),
			}, nil
		}
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("检查路径失败: %v", err),
		}, nil
	}

	kind := "文件"
	if info.IsDir() {
		kind = "目录"
	}

	return &ToolResult{
		ToolName: "file_ops",
		Success:  true,
		Data:     fmt.Sprintf("路径存在: %s (%s, %d bytes)", relPath, kind, info.Size()),
	}, nil
}

// handleDelete deletes a file.
func (t *FileTool) handleDelete(absPath, relPath string) (*ToolResult, error) {
	info, err := os.Stat(absPath)
	if err != nil {
		if os.IsNotExist(err) {
			return &ToolResult{
				ToolName: "file_ops",
				Success:  false,
				Error:    fmt.Sprintf("文件不存在: %s", relPath),
			}, nil
		}
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("删除文件失败: %v", err),
		}, nil
	}

	if info.IsDir() {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("不能删除目录（安全限制）: %s", relPath),
		}, nil
	}

	if err := os.Remove(absPath); err != nil {
		return &ToolResult{
			ToolName: "file_ops",
			Success:  false,
			Error:    fmt.Sprintf("删除文件失败: %v", err),
		}, nil
	}

	return &ToolResult{
		ToolName: "file_ops",
		Success:  true,
		Data:     fmt.Sprintf("已删除文件: %s", relPath),
	}, nil
}