picoclaw/pkg/tools/shell_test.go

package tools

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/sipeed/picoclaw/pkg/agent/sandbox"
)

type stubSandbox struct {
	lastReq sandbox.ExecRequest
	err     error
	res     *sandbox.ExecResult
	fs      sandbox.FsBridge
}

func (s *stubSandbox) Start(ctx context.Context) error { return nil }
func (s *stubSandbox) Prune(ctx context.Context) error { return nil }

func (s *stubSandbox) Resolve(ctx context.Context) (sandbox.Sandbox, error) {
	return s, nil
}

func (s *stubSandbox) Fs() sandbox.FsBridge {
	if s.fs != nil {
		return s.fs
	}
	return sandbox.NewHostSandbox("", false).Fs()
}

func (s *stubSandbox) Exec(ctx context.Context, req sandbox.ExecRequest) (*sandbox.ExecResult, error) {
	return sandboxAggregateFromStub(ctx, req, s.ExecStream)
}

func (s *stubSandbox) ExecStream(
	ctx context.Context,
	req sandbox.ExecRequest,
	onEvent func(sandbox.ExecEvent) error,
) (*sandbox.ExecResult, error) {
	s.lastReq = req
	if s.err != nil {
		return nil, s.err
	}
	if s.res != nil {
		if onEvent != nil {
			if s.res.Stdout != "" {
				if err := onEvent(
					sandbox.ExecEvent{Type: sandbox.ExecEventStdout, Chunk: []byte(s.res.Stdout)},
				); err != nil {
					return nil, err
				}
			}
			if s.res.Stderr != "" {
				if err := onEvent(
					sandbox.ExecEvent{Type: sandbox.ExecEventStderr, Chunk: []byte(s.res.Stderr)},
				); err != nil {
					return nil, err
				}
			}
			if err := onEvent(sandbox.ExecEvent{Type: sandbox.ExecEventExit, ExitCode: s.res.ExitCode}); err != nil {
				return nil, err
			}
		}
		return s.res, nil
	}
	if onEvent != nil {
		if err := onEvent(sandbox.ExecEvent{Type: sandbox.ExecEventStdout, Chunk: []byte("ok")}); err != nil {
			return nil, err
		}
		if err := onEvent(sandbox.ExecEvent{Type: sandbox.ExecEventExit, ExitCode: 0}); err != nil {
			return nil, err
		}
	}
	return &sandbox.ExecResult{Stdout: "ok", ExitCode: 0}, nil
}

func sandboxAggregateFromStub(
	ctx context.Context,
	req sandbox.ExecRequest,
	streamFn func(context.Context, sandbox.ExecRequest, func(sandbox.ExecEvent) error) (*sandbox.ExecResult, error),
) (*sandbox.ExecResult, error) {
	var stdout strings.Builder
	var stderr strings.Builder
	exitCode := 0
	res, err := streamFn(ctx, req, func(event sandbox.ExecEvent) error {
		switch event.Type {
		case sandbox.ExecEventStdout:
			_, _ = stdout.Write(event.Chunk)
		case sandbox.ExecEventStderr:
			_, _ = stderr.Write(event.Chunk)
		case sandbox.ExecEventExit:
			exitCode = event.ExitCode
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	if res != nil {
		return res, nil
	}
	return &sandbox.ExecResult{Stdout: stdout.String(), Stderr: stderr.String(), ExitCode: exitCode}, nil
}

// TestShellTool_Success verifies successful command execution
func TestShellTool_Success(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "hello world", ExitCode: 0},
	})
	args := map[string]any{
		"command": "echo 'hello world'",
	}

	result := tool.Execute(ctx, args)

	// Success should not be an error
	if result.IsError {
		t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
	}

	// ForUser should contain command output
	if !strings.Contains(result.ForUser, "hello world") {
		t.Errorf("Expected ForUser to contain 'hello world', got: %s", result.ForUser)
	}

	// ForLLM should contain full output
	if !strings.Contains(result.ForLLM, "hello world") {
		t.Errorf("Expected ForLLM to contain 'hello world', got: %s", result.ForLLM)
	}
}

// TestShellTool_Failure verifies failed command execution
func TestShellTool_Failure(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{
			Stdout:   "",
			Stderr:   "ls: cannot access '/nonexistent_directory_12345': No such file or directory",
			ExitCode: 2,
		},
	})
	args := map[string]any{
		"command": "ls /nonexistent_directory_12345",
	}

	result := tool.Execute(ctx, args)

	// Failure should be marked as error
	if !result.IsError {
		t.Errorf("Expected error for failed command, got IsError=false")
	}

	// ForUser should contain error information
	if result.ForUser == "" {
		t.Errorf("Expected ForUser to contain error info, got empty string")
	}

	// ForLLM should contain exit code or error
	if !strings.Contains(result.ForLLM, "Exit code") && result.ForUser == "" {
		t.Errorf("Expected ForLLM to contain exit code or error, got: %s", result.ForLLM)
	}
}

// TestShellTool_Timeout verifies command timeout handling
func TestShellTool_Timeout(t *testing.T) {
	tool := NewExecTool("", false)
	tool.SetTimeout(100 * time.Millisecond)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		err: context.DeadlineExceeded,
	})
	args := map[string]any{
		"command": "sleep 10",
	}

	result := tool.Execute(ctx, args)

	// Timeout should be marked as error
	if !result.IsError {
		t.Errorf("Expected error for timeout, got IsError=false")
	}

	// Should mention timeout
	if !strings.Contains(result.ForLLM, "timed out") && !strings.Contains(result.ForUser, "timed out") {
		t.Errorf("Expected timeout message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
	}
}

// TestShellTool_WorkingDir verifies custom working directory
func TestShellTool_WorkingDir(t *testing.T) {
	// Create temp directory
	tmpDir := t.TempDir()
	testFile := filepath.Join(tmpDir, "test.txt")
	os.WriteFile(testFile, []byte("test content"), 0o644)

	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "test content\n", ExitCode: 0},
	})
	args := map[string]any{
		"command":     "cat test.txt",
		"working_dir": tmpDir,
	}

	result := tool.Execute(ctx, args)

	if result.IsError {
		t.Errorf("Expected success in custom working dir, got error: %s", result.ForLLM)
	}

	if !strings.Contains(result.ForUser, "test content") {
		t.Errorf("Expected output from custom dir, got: %s", result.ForUser)
	}
}

// TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands
func TestShellTool_DangerousCommand(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "", Stderr: "", ExitCode: 1},
	})
	args := map[string]any{
		"command": "rm -rf /",
	}

	result := tool.Execute(ctx, args)

	// Dangerous command should be blocked
	if !result.IsError {
		t.Errorf("Expected dangerous command to be blocked (IsError=true)")
	}

	if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") {
		t.Errorf("Expected 'blocked' message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
	}
}

// TestShellTool_MissingCommand verifies error handling for missing command
func TestShellTool_MissingCommand(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "", Stderr: "", ExitCode: 1},
	})
	args := map[string]any{}

	result := tool.Execute(ctx, args)

	// Should return error result
	if !result.IsError {
		t.Errorf("Expected error when command is missing")
	}
}

// TestShellTool_StderrCapture verifies stderr is captured and included
func TestShellTool_StderrCapture(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "stdout", Stderr: "stderr", ExitCode: 0},
	})
	args := map[string]any{
		"command": "sh -c 'echo stdout; echo stderr >&2'",
	}

	result := tool.Execute(ctx, args)

	// Both stdout and stderr should be in output
	if !strings.Contains(result.ForLLM, "stdout") {
		t.Errorf("Expected stdout in output, got: %s", result.ForLLM)
	}
	if !strings.Contains(result.ForLLM, "stderr") {
		t.Errorf("Expected stderr in output, got: %s", result.ForLLM)
	}
}

// TestShellTool_OutputTruncation verifies long output is truncated
func TestShellTool_OutputTruncation(t *testing.T) {
	tool := NewExecTool("", false)

	ctx := sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: strings.Repeat("x", 20000), ExitCode: 0},
	})
	// Generate long output (>10000 chars)
	args := map[string]any{
		"command": "python3 -c \"print('x' * 20000)\" || echo " + strings.Repeat("x", 20000),
	}

	result := tool.Execute(ctx, args)

	// Should have truncation message or be truncated
	if len(result.ForLLM) > 15000 {
		t.Errorf("Expected output to be truncated, got length: %d", len(result.ForLLM))
	}
}

// TestShellTool_WorkingDir_OutsideWorkspace verifies that working_dir cannot escape the workspace directly
func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) {
	root := t.TempDir()
	workspace := filepath.Join(root, "workspace")
	outsideDir := filepath.Join(root, "outside")
	if err := os.MkdirAll(workspace, 0o755); err != nil {
		t.Fatalf("failed to create workspace: %v", err)
	}
	if err := os.MkdirAll(outsideDir, 0o755); err != nil {
		t.Fatalf("failed to create outside dir: %v", err)
	}

	tool := NewExecTool(workspace, true)
	result := tool.Execute(sandbox.WithSandbox(context.Background(), &stubSandbox{
		res: &sandbox.ExecResult{Stdout: "", Stderr: "", ExitCode: 1},
	}), map[string]any{
		"command":     "pwd",
		"working_dir": outsideDir,
	})

	if !result.IsError {
		t.Fatalf("expected working_dir outside workspace to be blocked, got output: %s", result.ForLLM)
	}
	if !strings.Contains(result.ForLLM, "blocked") {
		t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM)
	}
}

// TestShellTool_WorkingDir_SymlinkEscape verifies that a symlink inside the workspace
// pointing outside cannot be used as working_dir to escape the sandbox.
func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) {
	root := t.TempDir()
	workspace := filepath.Join(root, "workspace")
	secretDir := filepath.Join(root, "secret")
	if err := os.MkdirAll(workspace, 0o755); err != nil {
		t.Fatalf("failed to create workspace: %v", err)
	}
	if err := os.MkdirAll(secretDir, 0o755); err != nil {
		t.Fatalf("failed to create secret dir: %v", err)
	}
	os.WriteFile(filepath.Join(secretDir, "secret.txt"), []byte("top secret"), 0o644)

	// symlink lives inside the workspace but resolves to secretDir outside it
	link := filepath.Join(workspace, "escape")
	if err := os.Symlink(secretDir, link); err != nil {
		t.Skipf("symlinks not supported in this environment: %v", err)
	}

	tool := NewExecTool(workspace, true)
	result := tool.Execute(context.Background(), map[string]any{
		"command":     "cat secret.txt",
		"working_dir": link,
	})

	if !result.IsError {
		t.Fatalf("expected symlink working_dir escape to be blocked, got output: %s", result.ForLLM)
	}
	if !strings.Contains(result.ForLLM, "blocked") {
		t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM)
	}
}

// TestShellTool_RestrictToWorkspace verifies workspace restriction
func TestShellTool_RestrictToWorkspace(t *testing.T) {
	tmpDir := t.TempDir()
	tool := NewExecTool(tmpDir, false)
	tool.SetRestrictToWorkspace(true)

	ctx := context.Background()
	args := map[string]any{
		"command": "cat ../../etc/passwd",
	}

	result := tool.Execute(ctx, args)

	// Path traversal should be blocked
	if !result.IsError {
		t.Errorf("Expected path traversal to be blocked with restrictToWorkspace=true")
	}

	if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") {
		t.Errorf(
			"Expected 'blocked' message for path traversal, got ForLLM: %s, ForUser: %s",
			result.ForLLM,
			result.ForUser,
		)
	}
}

func TestShellTool_SandboxMapsHostWorkingDirToRelative(t *testing.T) {
	workspace := t.TempDir()
	sb := &stubSandbox{}
	tool := NewExecTool(workspace, true)

	ctx := sandbox.WithSandbox(context.Background(), sb)
	args := map[string]any{
		"command":     "echo test",
		"working_dir": filepath.Join(workspace, "subdir"),
	}
	result := tool.Execute(ctx, args)
	if result.IsError {
		t.Fatalf("expected success, got error: %s", result.ForLLM)
	}
	if sb.lastReq.WorkingDir != "subdir" {
		t.Fatalf("sandbox working_dir = %q, want subdir", sb.lastReq.WorkingDir)
	}
}

func TestShellTool_SandboxAllowsAbsoluteWorkspaceWorkingDir(t *testing.T) {
	workspace := t.TempDir()
	sb := &stubSandbox{}
	tool := NewExecTool(workspace, true)

	ctx := sandbox.WithSandbox(context.Background(), sb)
	args := map[string]any{
		"command":     "echo test",
		"working_dir": "/workspace/subdir",
	}
	result := tool.Execute(ctx, args)
	if result.IsError {
		t.Fatalf("expected success, got error: %s", result.ForLLM)
	}
	if sb.lastReq.WorkingDir != "/workspace/subdir" {
		t.Fatalf("sandbox working_dir = %q, want /workspace/subdir", sb.lastReq.WorkingDir)
	}
}

func TestShellTool_SandboxBlocksAbsoluteNonWorkspaceWorkingDirWhenRestricted(t *testing.T) {
	workspace := t.TempDir()
	sb := &stubSandbox{}
	tool := NewExecTool(workspace, true)

	ctx := sandbox.WithSandbox(context.Background(), sb)
	args := map[string]any{
		"command":     "echo test",
		"working_dir": "/tmp/logs",
	}
	result := tool.Execute(ctx, args)
	if !result.IsError {
		t.Fatalf("expected error for /tmp/logs with restrict_to_workspace=true, got: %s", result.ForLLM)
	}
	if !strings.Contains(result.ForLLM, "blocked") {
		t.Fatalf("expected blocked error, got: %s", result.ForLLM)
	}
}

func TestShellTool_SandboxExecError(t *testing.T) {
	workspace := t.TempDir()
	sb := &stubSandbox{err: fmt.Errorf("sandbox down")}
	tool := NewExecTool(workspace, true)

	ctx := sandbox.WithSandbox(context.Background(), sb)
	result := tool.Execute(ctx, map[string]any{"command": "echo test"})
	if !result.IsError {
		t.Fatal("expected sandbox error result")
	}
	if !strings.Contains(result.ForLLM, "sandbox exec failed") {
		t.Fatalf("unexpected error message: %s", result.ForLLM)
	}
}