// 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"
	"errors"
	"fmt"
	"os"
	"strings"
	"testing"
	"time"

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

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

func pentestRedirRun(t *testing.T, script, dir string) (string, string, int) {
	t.Helper()
	// Parse outside the timeout so that slow parsing under the race
	// detector does not eat into the 5-second execution budget.
	parser := syntax.NewParser()
	prog, err := parser.Parse(strings.NewReader(script), "")
	if err != nil {
		return "", err.Error(), 2
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	return pentestRedirRunProg(ctx, t, prog, dir)
}

func pentestRedirRunProg(ctx context.Context, t *testing.T, prog *syntax.File, dir string) (string, string, int) {
	t.Helper()

	var outBuf, errBuf bytes.Buffer
	opts := []interp.RunnerOption{
		interp.StdIO(nil, &outBuf, &errBuf),
		interpoption.AllowAllCommands().(interp.RunnerOption),
	}
	if dir != "" {
		opts = append(opts, interp.AllowedPaths([]string{dir}))
	}

	runner, err := interp.New(opts...)
	require.NoError(t, err)
	defer runner.Close()

	if dir != "" {
		runner.Dir = dir
	}

	err = runner.Run(ctx, prog)
	exitCode := 0
	if err != nil {
		var es interp.ExitStatus
		if errors.As(err, &es) {
			exitCode = int(es)
		} else if ctx.Err() != nil {
			return outBuf.String(), errBuf.String(), -1 // timeout
		} else {
			t.Fatalf("unexpected error: %v", err)
		}
	}
	return outBuf.String(), errBuf.String(), exitCode
}

// --- Path traversal attacks ---

func TestPentestRedirPathTraversal(t *testing.T) {
	dir := t.TempDir()
	tests := []struct {
		name   string
		script string
	}{
		{"dot-dot traversal", "echo hello > /dev/null/../../tmp/evil"},
		{"dot-dot from devnull", "echo hello > /dev/null/../passwd"},
		{"double slash", "echo hello > /dev//null"},
		{"dot in path", "echo hello > /dev/./null"},
		{"trailing slash", "echo hello > /dev/null/"},
		{"case variation", "echo hello > /Dev/Null"},
		{"relative devnull", "echo hello > dev/null"},
		{"bare null", "echo hello > null"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			stdout, stderr, code := pentestRedirRun(t, tt.script, dir)
			assert.Equal(t, "", stdout, "should produce no stdout")
			assert.NotEqual(t, 0, code, "should fail with non-zero exit")
			// Should either be validation error (exit 2) or runtime error
			assert.True(t, code == 2 || code == 1, "exit code should be 1 or 2, got %d", code)
			_ = stderr // error message varies
		})
	}
}

// --- Variable expansion attacks ---

func TestPentestRedirVariableExpansion(t *testing.T) {
	dir := t.TempDir()
	tests := []struct {
		name   string
		script string
	}{
		{"variable target", "TARGET=/dev/null; echo hello > $TARGET"},
		{"variable partial", "DEV=/dev; echo hello > $DEV/null"},
		{"variable with braces", "TARGET=/dev/null; echo hello > ${TARGET}"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, _, code := pentestRedirRun(t, tt.script, dir)
			assert.Equal(t, 2, code, "variable expansion in redirect target should be blocked at validation")
		})
	}
}

// --- Quoting attacks ---

func TestPentestRedirQuotedDevNull(t *testing.T) {
	dir := t.TempDir()
	tests := []struct {
		name   string
		script string
	}{
		{"single quoted", "echo hello > '/dev/null'"},
		{"double quoted", `echo hello > "/dev/null"`},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, _, code := pentestRedirRun(t, tt.script, dir)
			// Quoted paths have different AST structure (SglQuoted/DblQuoted vs Lit)
			// Our check requires a single Lit part, so quoted paths should be rejected
			assert.Equal(t, 2, code, "quoted /dev/null in redirect should be blocked at validation")
		})
	}
}

// --- Glob/wildcard attacks ---

func TestPentestRedirGlobInTarget(t *testing.T) {
	dir := t.TempDir()
	// Glob characters in redirect targets
	_, _, code := pentestRedirRun(t, "echo hello > /dev/nul?", dir)
	assert.Equal(t, 2, code, "glob in redirect target should be rejected")
}

// --- fd duplication attacks ---

func TestPentestRedirFdDupAttacks(t *testing.T) {
	dir := t.TempDir()
	tests := []struct {
		name   string
		script string
	}{
		{"fd 0 dup", "echo hello 0>&1"},
		{"fd 3 dup", "echo hello 3>&1"},
		{"fd 9 dup", "echo hello 9>&1"},
		{"fd to 0", "echo hello >&0"},
		{"fd to 3", "echo hello >&3"},
		{"close fd", "echo hello >&-"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, _, code := pentestRedirRun(t, tt.script, dir)
			assert.Equal(t, 2, code, "unsupported fd duplication should be blocked")
		})
	}
}

// --- Redirect-to-file attacks that must remain blocked ---

func TestPentestRedirToSensitiveFiles(t *testing.T) {
	dir := t.TempDir()
	tests := []struct {
		name   string
		script string
	}{
		{"etc passwd", "echo evil > /etc/passwd"},
		{"etc shadow", "echo evil > /etc/shadow"},
		{"tmp file", "echo evil > /tmp/evil.txt"},
		{"home file", "echo evil > ~/.bashrc"},
		{"proc self", "echo evil > /proc/self/mem"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, stderr, code := pentestRedirRun(t, tt.script, dir)
			assert.NotEqual(t, 0, code, "redirect to %s should be blocked", tt.name)
			_ = stderr
		})
	}
}

// --- DplIn still blocked ---

func TestPentestRedirDplInBlocked(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := pentestRedirRun(t, "echo hello <&0", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "fd duplication is not supported")
}

// --- Read-write redirect still blocked ---

func TestPentestRedirReadWriteBlocked(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := pentestRedirRun(t, "echo hello <> /dev/null", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "file redirection is not supported")
}

// --- Heredoc size limit ---

func TestPentestHeredocSizeLimitRegular(t *testing.T) {
	// A regular << heredoc that exceeds MaxHeredocBytes must be rejected with a stderr message.
	body := strings.Repeat("x", interp.MaxHeredocBytes+1)
	script := "cat <<EOF\n" + body + "\nEOF"
	_, stderr, code := pentestRedirRun(t, script, "")
	assert.NotEqual(t, 0, code, "<< heredoc exceeding MaxHeredocBytes should fail")
	assert.Contains(t, stderr, "exceeds maximum size")
}

func TestPentestHeredocSizeLimitDash(t *testing.T) {
	// A <<- heredoc that exceeds MaxHeredocBytes must be rejected with a stderr message.
	body := strings.Repeat("x", interp.MaxHeredocBytes+1)
	script := "cat <<-EOF\n" + body + "\nEOF"
	_, stderr, code := pentestRedirRun(t, script, "")
	assert.NotEqual(t, 0, code, "<<- heredoc exceeding MaxHeredocBytes should fail")
	assert.Contains(t, stderr, "exceeds maximum size")
}

func TestPentestHeredocSizeLimitUnderLimit(t *testing.T) {
	// A heredoc just under the limit should succeed.
	body := strings.Repeat("x", interp.MaxHeredocBytes-1)
	script := "cat <<EOF\n" + body + "\nEOF"
	stdout, _, code := pentestRedirRun(t, script, "")
	assert.Equal(t, 0, code, "heredoc under MaxHeredocBytes should succeed")
	assert.Equal(t, body+"\n", stdout)
}

// --- Herestring still blocked ---

func TestPentestRedirHerestringBlocked(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := pentestRedirRun(t, "cat <<< 'hello'", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "herestring")
}

// --- Multiple redirects mixing allowed and blocked ---

func TestPentestRedirMixedAllowedBlocked(t *testing.T) {
	dir := t.TempDir()
	// First redirect is allowed, second is not
	_, _, code := pentestRedirRun(t, "echo hello >/dev/null > /tmp/evil", dir)
	assert.Equal(t, 2, code, "mixed redirects with blocked target should fail at validation")
}

// --- Ensure /dev/null redirect doesn't create any files ---

func TestPentestRedirNoFileCreated(t *testing.T) {
	dir := t.TempDir()
	// Run a redirect to /dev/null
	pentestRedirRun(t, "echo hello >/dev/null", dir)

	// Inspect the temp dir directly from Go: running `ls $dir` through the shell
	// would mis-tokenize on Windows (backslash path separators) and would silently
	// pass if `ls` itself failed, since this test ignores stderr/exit code.
	entries, err := os.ReadDir(dir)
	require.NoError(t, err)
	assert.Empty(t, entries, "no files should have been created")
}

// --- vuln-hunt 2026-05-18-initial-audit ---

// A heredoc body can only carry as much data as the surrounding variable
// machinery allows. Three layered defenses prevent a > 10 MiB heredoc
// expansion from ever being assembled:
//
//  1. Per-variable cap: 1 MiB on any single assignment.
//  2. Per-runner variable storage cap: 2 MiB total across all variables.
//  3. Post-expansion heredoc check: 10 MiB on the final expanded body.
//
// Layers (1) and (2) make a multi-var concat attack unreachable; the
// post-expansion check at runner_redir.go is defense-in-depth that would
// fire if those caps were ever loosened. This test pins both reachable
// caps (1) and (2) — both emit a stderr error, and even though they don't
// fail the surrounding script (matching bash assignment semantics), the
// over-cap data never propagates into the heredoc body.
func TestVulnHuntHeredocPostExpansionExceedsCap(t *testing.T) {
	dir := t.TempDir()

	// Layer 1: over-cap single assignment. Stderr mentions the cap; the
	// expansion in the heredoc resolves to the empty string, so no data
	// is smuggled through.
	body := strings.Repeat("x", interp.MaxHeredocBytes+1)
	script1 := "BIG='" + body + "'\ncat <<EOF\nMARKER:$BIG:END\nEOF"
	stdout1, stderr1, _ := pentestRedirRun(t, script1, dir)
	assert.Contains(t, stderr1, "value too large",
		"over-cap assignment must emit the per-var cap message; stderr=%q", stderr1)
	assert.Equal(t, "MARKER::END\n", stdout1,
		"the over-cap value must not propagate into the heredoc; got %q", stdout1)

	// Layer 2: multi-variable accumulation over the per-runner storage
	// cap (2 MiB total). After 2 or 3 assignments succeed, subsequent ones
	// fail with "variable storage limit exceeded" and the storage cap
	// prevents the assignments accumulating to anywhere near the 10 MiB
	// heredoc cap.
	chunk := strings.Repeat("y", 1024*1024-1)
	var sb strings.Builder
	for i := 0; i < 11; i++ {
		sb.WriteString(fmt.Sprintf("V%d='%s'\n", i, chunk))
	}
	_, stderr2, _ := pentestRedirRun(t, sb.String(), dir)
	assert.Contains(t, stderr2, "storage limit",
		"multi-var accumulation past 2 MiB must hit the storage cap; stderr=%q", stderr2)
}

// Heredoc writer goroutine polls ctx.Err() between 32 KiB chunks. Cancelling
// the executor mid-heredoc must terminate the write rather than leak the
// goroutine. We can observe this by running a heredoc just under the cap and
// cancelling the context immediately — the runner must return promptly.
func TestVulnHuntHeredocWriterHonorsCtxCancel(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping ctx-cancel timing test in -short mode")
	}
	dir := t.TempDir()
	body := strings.Repeat("y", interp.MaxHeredocBytes-1)
	script := "cat <<EOF\n" + body + "\nEOF"

	parser := syntax.NewParser()
	prog, err := parser.Parse(strings.NewReader(script), "")
	require.NoError(t, err)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	done := make(chan struct{})
	go func() {
		pentestRedirRunProg(ctx, t, prog, dir)
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("heredoc writer did not honor ctx cancel within budget")
	}
}
