// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-present Datadog, Inc.

package tests_test

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"mvdan.cc/sh/v3/syntax"

	"github.com/DataDog/rshell/interp"
)

// --- Blocked features inside command substitution ---

func TestCmdSubstPentestProcessSubInSubst(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := cmdSubstRun(t, `echo $(<(echo hello))`, dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "process substitution is not supported")
}

// --- Command substitution with blocked features ---

func TestCmdSubstPentestFuncInSubst(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := cmdSubstRun(t, `x=$(f() { echo a; })`, dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "function declarations are not supported")
}

// --- Path traversal attempts ---

func TestCmdSubstPentestPathTraversal(t *testing.T) {
	dir := t.TempDir()
	// Attempt to read /etc/passwd via $(<file) - should fail due to sandbox
	_, stderr, code := cmdSubstRun(t, `x=$(<../../../../../../etc/passwd)`, dir)
	assert.NotEqual(t, 0, code)
	_ = stderr
}

func TestCmdSubstPentestAbsolutePath(t *testing.T) {
	dir := t.TempDir()
	// Absolute path outside sandbox should fail
	_, stderr, code := cmdSubstRun(t, `x=$(<"/etc/hosts")`, dir)
	assert.NotEqual(t, 0, code)
	_ = stderr
}

// --- No allowed paths at all ---

func TestCmdSubstPentestNoAllowedPaths(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(dir, "test.txt"), []byte("data"), 0644))
	// With no allowed paths, file access should be blocked
	_, stderr, code := cmdSubstRunWithOpts(t, `x=$(<test.txt)`, dir)
	assert.NotEqual(t, 0, code)
	_ = stderr
}

// --- Subshell pentest: exit doesn't kill parent ---

func TestSubshellPentestExitDoesntKillParent(t *testing.T) {
	dir := t.TempDir()
	stdout, _, code := subshellRun(t, `(exit 255); echo alive`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "alive\n", stdout)
}

// --- Large number of subshells ---

func TestSubshellPentestManySubshells(t *testing.T) {
	dir := t.TempDir()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	// Run 50 subshells in sequence
	var sb strings.Builder
	for i := 0; i < 50; i++ {
		sb.WriteString("(echo ok); ")
	}
	stdout, _, code := cmdSubstRunCtx(ctx, t, sb.String(), dir)
	assert.Equal(t, 0, code)
	lines := strings.Split(strings.TrimSuffix(stdout, "\n"), "\n")
	assert.Equal(t, 50, len(lines))
}

// --- Command substitution with stderr ---

func TestCmdSubstPentestStderrNotCaptured(t *testing.T) {
	dir := t.TempDir()
	// Stderr from command substitution should go to parent's stderr
	stdout, stderr, code := cmdSubstRun(t, `x=$(echo err >&2); echo "out=$x"`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "out=\n", stdout)
	assert.Equal(t, "err\n", stderr)
}

// --- Deeply nested command substitution ---

func TestCmdSubstPentestDeepNesting(t *testing.T) {
	dir := t.TempDir()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	// 5 levels deep
	stdout, _, code := cmdSubstRunCtx(ctx, t, `echo $(echo $(echo $(echo $(echo deep))))`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "deep\n", stdout)
}

// --- Command substitution as argument to builtin ---

func TestCmdSubstPentestAsGrepPattern(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(dir, "data.txt"), []byte("hello\nworld\n"), 0644))
	stdout, _, code := cmdSubstRun(t, `grep $(echo world) data.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "world\n", stdout)
}

// --- Command substitution with special characters ---

func TestCmdSubstPentestSpecialChars(t *testing.T) {
	dir := t.TempDir()
	stdout, _, code := cmdSubstRun(t, `echo "$(printf 'tab\there')"`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "tab\there\n", stdout)
}

func TestCmdSubstPentestNullBytes(t *testing.T) {
	dir := t.TempDir()
	// Null bytes in output — printf with \0
	stdout, _, code := cmdSubstRun(t, `x=$(printf "a\x00b"); echo "$x"`, dir)
	// Should not crash; behavior may vary (null bytes often truncate)
	assert.Contains(t, []int{0, 1}, code)
	_ = stdout
}

// --- Subshell with blocked features ---

func TestSubshellPentestBackgroundInSubshell(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := subshellRun(t, "(echo hi &\n)", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "background execution (&) is not supported")
}

// --- Cat shortcut edge cases ---

func TestCmdSubstPentestCatShortcutDirectory(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.MkdirAll(filepath.Join(dir, "subdir"), 0755))
	stdout, stderr, code := cmdSubstRun(t, `x=$(<subdir); echo "[$x]"`, dir)
	// Bash silently produces empty output with exit 0 for $(<dir).
	assert.Equal(t, 0, code)
	assert.Equal(t, "[]\n", stdout)
	assert.Equal(t, "", stderr)
}

func TestCmdSubstPentestCatShortcutEmptyFile(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(dir, "empty.txt"), []byte{}, 0644))
	stdout, _, code := cmdSubstRun(t, `x=$(<empty.txt); echo "[$x]"`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "[]\n", stdout)
}

// --- Subshell with redirect ---

func TestSubshellPentestRedirectOutBlocked(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := subshellRun(t, `(echo data) > /tmp/evil.txt`, dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "file redirection is not supported")
}

// --- Context cancellation ---

func TestCmdSubstPentestTimeout(t *testing.T) {
	dir := t.TempDir()
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()
	// Run a subshell that would loop forever using for loop with many items
	// Since while loops are blocked, this tests with a large for loop
	var items strings.Builder
	for i := 0; i < 100000; i++ {
		items.WriteString("x ")
	}
	script := `for x in ` + items.String() + `; do echo "$x"; done`
	_, _, code := cmdSubstRunCtx(ctx, t, `y=$(`+script+`); echo done`, dir)
	// Should either timeout or complete
	_ = code
}

func TestCmdSubstFatalExitPropagated(t *testing.T) {
	dir := t.TempDir()
	// Use a very short timeout so the long-running command substitution
	// triggers context deadline exceeded inside the subshell.
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	var items strings.Builder
	for i := 0; i < 100000; i++ {
		items.WriteString("x ")
	}
	script := `for x in ` + items.String() + `; do echo "$x"; done`

	parser := syntax.NewParser()
	prog, err := parser.Parse(strings.NewReader(`for y in $(`+script+`); do echo "$y"; done`), "")
	require.NoError(t, err)

	var outBuf, errBuf bytes.Buffer
	runner, err := interp.New(
		interp.StdIO(nil, &outBuf, &errBuf),
		interp.AllowedPaths([]string{dir}),
	)
	require.NoError(t, err)
	defer runner.Close()
	runner.Dir = dir

	// The fatal error from the subshell must propagate so that Run()
	// returns a non-nil error; prior to the fix, callers could see
	// a nil error (exit 0) despite the context being cancelled.
	runErr := runner.Run(ctx, prog)
	assert.Error(t, runErr, "Run should return non-nil error after context deadline in cmdsubst")
}

// --- Verify process substitution is still blocked ---

func TestPentestProcessSubstitutionStillBlocked(t *testing.T) {
	dir := t.TempDir()
	tests := []string{
		`cat <(echo hello)`,
		`echo >(cat)`,
		`grep pattern <(echo hello)`,
	}
	for _, tt := range tests {
		t.Run(tt, func(t *testing.T) {
			_, stderr, code := cmdSubstRun(t, tt, dir)
			assert.Equal(t, 2, code)
			assert.Contains(t, stderr, "process substitution is not supported")
		})
	}
}

// --- Subshell inherits parent variables ---

func TestSubshellPentestInheritsVars(t *testing.T) {
	dir := t.TempDir()
	stdout, _, code := subshellRun(t, `x=parent; (echo "$x")`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "parent\n", stdout)
}

// --- Command substitution in different contexts ---

func TestCmdSubstPentestInHeredocUnquoted(t *testing.T) {
	dir := t.TempDir()
	stdout, _, code := cmdSubstRun(t, "cat <<EOF\nhello $(echo world)\nEOF", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "hello world\n", stdout)
}

func TestCmdSubstPentestInHeredocQuoted(t *testing.T) {
	dir := t.TempDir()
	// Quoted heredoc should NOT expand command substitution
	stdout, _, code := cmdSubstRun(t, "cat <<'EOF'\nhello $(echo world)\nEOF", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "hello $(echo world)\n", stdout)
}
