Cyrene_v2/plugins/tools/io/file.go

package io

import (
	"context"
	"errors"
	"fmt"
	stdio "io"
	"io/fs"
	"log/slog"
	"os"
	"path/filepath"
	"strings"

	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// Compile-time assertion that fileTool implements sdk.Tool.
var _ sdk.Tool = (*fileTool)(nil)

// fileTool reads, writes, lists and searches files, confined to the configured
// data directory plus any explicit whitelist. Paths are cleaned and re-checked
// against those roots, so ".." traversal and symlink escapes are rejected.
type fileTool struct {
	roots    []string // absolute, cleaned allowed directories
	maxRead  int64
	maxWrite int64
	maxList  int
	log      *slog.Logger
}

func newFileTool(t *tools) *fileTool {
	roots := make([]string, 0, 1+len(t.fileWhitelist))
	if t.cfg != nil && t.cfg.DataDir != "" {
		roots = append(roots, absoluteClean(t.cfg.DataDir))
	}
	for _, w := range t.fileWhitelist {
		roots = append(roots, absoluteClean(w))
	}
	return &fileTool{
		roots:    roots,
		maxRead:  t.maxReadBytes,
		maxWrite: t.maxWriteBytes,
		maxList:  t.maxList,
		log:      t.log,
	}
}

// Name implements sdk.Tool.
func (f *fileTool) Name() string { return "file" }

// Description implements sdk.Tool.
func (f *fileTool) Description() string {
	return "Read, write, list or search files, restricted to the configured data directory and explicit whitelist."
}

// Parameters implements sdk.Tool.
func (f *fileTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"action":  map[string]any{"type": "string", "enum": []any{"read", "write", "list", "search"}, "description": "Operation to perform."},
			"path":    map[string]any{"type": "string", "description": "Absolute or data-directory-relative file path."},
			"content": map[string]any{"type": "string", "description": "Body to write."},
			"query":   map[string]any{"type": "string", "description": "Substring filter for search."},
		},
		"required": []any{"action", "path"},
	}, nil
}

// Execute implements sdk.Tool.
func (f *fileTool) Execute(ctx context.Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("io: file: args must be an object")
	}
	switch action := argString(m, "action"); action {
	case "read":
		return f.read(ctx, argString(m, "path"))
	case "write":
		return f.write(ctx, argString(m, "path"), argString(m, "content"))
	case "list":
		return f.list(ctx, argString(m, "path"))
	case "search":
		return f.search(ctx, argString(m, "path"), argString(m, "query"))
	default:
		return nil, fmt.Errorf("io: file: unknown action %q", action)
	}
}

// read returns a file's content, capped at maxRead bytes.
func (f *fileTool) read(ctx context.Context, p string) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	real, err := f.resolveExisting(p)
	if err != nil {
		return nil, err
	}
	info, err := os.Stat(real)
	if err != nil {
		return nil, fmt.Errorf("io: file: stat: %w", err)
	}
	if info.IsDir() {
		return nil, fmt.Errorf("io: file: %q is a directory", real)
	}
	if info.Size() > f.maxRead {
		return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
	}
	fh, err := os.Open(real)
	if err != nil {
		return nil, fmt.Errorf("io: file: open: %w", err)
	}
	defer fh.Close()
	data, err := stdio.ReadAll(stdio.LimitReader(fh, f.maxRead+1))
	if err != nil {
		return nil, fmt.Errorf("io: file: read: %w", err)
	}
	if int64(len(data)) > f.maxRead {
		return nil, fmt.Errorf("io: file: %q exceeds the %d-byte read limit", real, f.maxRead)
	}
	return map[string]any{"path": real, "size": len(data), "content": string(data)}, nil
}

// write creates or overwrites a file, capped at maxWrite bytes and confined to
// the allowed roots (parent dirs are created inside the root only).
func (f *fileTool) write(ctx context.Context, p, content string) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	clean, err := f.resolve(p)
	if err != nil {
		return nil, err
	}
	if int64(len(content)) > f.maxWrite {
		return nil, fmt.Errorf("io: file: content exceeds the %d-byte write limit", f.maxWrite)
	}
	parent := filepath.Dir(clean)
	if err := os.MkdirAll(parent, dirPerm); err != nil {
		return nil, fmt.Errorf("io: file: mkdir: %w", err)
	}
	realParent, err := filepath.EvalSymlinks(parent)
	if err != nil {
		return nil, fmt.Errorf("io: file: resolve parent: %w", err)
	}
	if !f.within(realParent) {
		return nil, fmt.Errorf("io: file: parent directory escapes the allowed directories")
	}
	if info, err := os.Lstat(clean); err == nil && info.Mode()&os.ModeSymlink != 0 {
		real, err := filepath.EvalSymlinks(clean)
		if err != nil || !f.within(real) {
			return nil, fmt.Errorf("io: file: %q escapes the allowed directories", p)
		}
	}
	if err := os.WriteFile(clean, []byte(content), filePerm); err != nil {
		return nil, fmt.Errorf("io: file: write: %w", err)
	}
	return map[string]any{"path": clean, "size": len(content)}, nil
}

// list returns the entries of a directory, capped at maxList.
func (f *fileTool) list(ctx context.Context, p string) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	real, err := f.resolveExisting(p)
	if err != nil {
		return nil, err
	}
	info, err := os.Stat(real)
	if err != nil {
		return nil, fmt.Errorf("io: file: stat: %w", err)
	}
	if !info.IsDir() {
		return nil, fmt.Errorf("io: file: %q is not a directory", real)
	}
	entries, err := os.ReadDir(real)
	if err != nil {
		return nil, fmt.Errorf("io: file: readdir: %w", err)
	}
	if len(entries) > f.maxList {
		entries = entries[:f.maxList]
	}
	out := make([]map[string]any, 0, len(entries))
	for _, e := range entries {
		ei, err := e.Info()
		if err != nil {
			continue
		}
		out = append(out, map[string]any{"name": e.Name(), "size": ei.Size(), "dir": ei.IsDir()})
	}
	return map[string]any{"path": real, "entries": out, "count": len(out)}, nil
}

// search walks a base directory (default the primary root) and returns files
// whose names contain query, capped at maxList results.
func (f *fileTool) search(ctx context.Context, base, query string) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if strings.TrimSpace(query) == "" {
		return nil, fmt.Errorf("io: file: search query is required")
	}
	if strings.TrimSpace(base) == "" {
		if len(f.roots) == 0 {
			return nil, fmt.Errorf("io: file: no base directory configured")
		}
		base = f.roots[0]
	}
	real, err := f.resolveExisting(base)
	if err != nil {
		return nil, err
	}
	info, err := os.Stat(real)
	if err != nil {
		return nil, fmt.Errorf("io: file: stat: %w", err)
	}
	if !info.IsDir() {
		return nil, fmt.Errorf("io: file: %q is not a directory", real)
	}
	q := strings.ToLower(query)
	hits := make([]map[string]any, 0)
	walkErr := filepath.WalkDir(real, func(pp string, d fs.DirEntry, err error) error {
		if err != nil || d.IsDir() {
			return nil
		}
		if !strings.Contains(strings.ToLower(d.Name()), q) {
			return nil
		}
		di, err := d.Info()
		if err != nil {
			return nil
		}
		rel, rerr := filepath.Rel(real, pp)
		if rerr != nil {
			rel = pp
		}
		hits = append(hits, map[string]any{"path": rel, "size": di.Size()})
		if len(hits) >= f.maxList {
			return fs.SkipAll
		}
		return nil
	})
	if walkErr != nil && !errors.Is(walkErr, fs.SkipAll) {
		return nil, fmt.Errorf("io: file: walk: %w", walkErr)
	}
	return map[string]any{"base": real, "query": query, "hits": hits, "count": len(hits)}, nil
}

// resolve cleans p and verifies it stays within one of the allowed roots. It
// rejects empty paths and any path that escapes via ".." or an absolute target
// outside the roots.
func (f *fileTool) resolve(p string) (string, error) {
	if strings.TrimSpace(p) == "" {
		return "", fmt.Errorf("io: file: path is required")
	}
	if len(f.roots) == 0 {
		return "", fmt.Errorf("io: file: no base directory configured")
	}
	clean := filepath.Clean(p)
	if !filepath.IsAbs(clean) {
		clean = filepath.Join(f.roots[0], clean)
	}
	clean = filepath.Clean(clean)
	if !f.within(clean) {
		return "", fmt.Errorf("io: file: %q is outside the allowed directories", p)
	}
	return clean, nil
}

// resolveExisting resolves p and verifies the symlink-resolved real path stays
// within the roots, guarding against escaping through a symlink. It errors when
// the target does not exist.
func (f *fileTool) resolveExisting(p string) (string, error) {
	clean, err := f.resolve(p)
	if err != nil {
		return "", err
	}
	real, err := filepath.EvalSymlinks(clean)
	if err != nil {
		return "", fmt.Errorf("io: file: resolve %q: %w", p, err)
	}
	if !f.within(real) {
		f.log.Warn("io file: rejected path escaping allowed directories", "path", p, "resolved", real)
		return "", fmt.Errorf("io: file: %q escapes the allowed directories", p)
	}
	return real, nil
}

// within reports whether p is inside (or equal to) one of the allowed roots.
func (f *fileTool) within(p string) bool {
	for _, root := range f.roots {
		rel, err := filepath.Rel(root, p)
		if err != nil {
			continue
		}
		if rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) {
			return true
		}
	}
	return false
}

// absoluteClean absolutizes and cleans p, returning the best-effort absolute
// form used as an allowed root.
func absoluteClean(p string) string {
	a, err := filepath.Abs(p)
	if err != nil {
		return filepath.Clean(p)
	}
	return filepath.Clean(a)
}