// 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 grep_test

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/DataDog/rshell/builtins/grep"
	"github.com/DataDog/rshell/builtins/testutil"
	"github.com/DataDog/rshell/interp"
)

const pentestTimeout = 10 * time.Second

func grepRun(t *testing.T, script, dir string, extraPaths ...string) (string, string, int) {
	t.Helper()
	paths := append([]string{dir}, extraPaths...)
	return testutil.RunScript(t, script, dir, interp.AllowedPaths(paths))
}

func grepRunCtx(ctx context.Context, t *testing.T, script, dir string) (string, string, int) {
	t.Helper()
	return testutil.RunScriptCtx(ctx, t, script, dir, interp.AllowedPaths([]string{dir}))
}

func mustNotHang(t *testing.T, fn func()) {
	t.Helper()
	done := make(chan struct{})
	go func() {
		fn()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(pentestTimeout):
		t.Fatal("operation did not complete within timeout")
	}
}

func pentestWriteFile(t *testing.T, dir, name, content string) string {
	t.Helper()
	require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0644))
	return name
}

// --- Flag and argument injection ---

func TestGrepPentestUnknownLongFlag(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := grepRun(t, "grep --recursive foo", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestUnknownShortFlag(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := grepRun(t, "grep -P foo file.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestFlagViaExpansion(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "ok\n")
	_, stderr, code := grepRun(t, `for flag in -P; do grep $flag foo file.txt; done`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestDoubleDashThenFlagLikeArg(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "-v\nfoo\n")
	stdout, _, code := grepRun(t, "grep -- -v file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "-v\n", stdout)
}

// --- Path and filename edge cases ---

func TestGrepPentestNonexistentFile(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := grepRun(t, "grep foo does_not_exist", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestDirectoryAsFile(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.Mkdir(filepath.Join(dir, "subdir"), 0755))
	_, stderr, code := grepRun(t, "grep foo subdir", dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestEmptyStringFilename(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := grepRun(t, `grep foo ""`, dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "grep:")
}

func TestGrepPentestPathTraversal(t *testing.T) {
	dir := t.TempDir()
	outer := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(outer, "secret.txt"), []byte("secret"), 0644))
	outerPath := strings.ReplaceAll(filepath.Join(outer, "secret.txt"), `\`, `/`)
	_, stderr, code := grepRun(t, "grep secret "+outerPath, dir)
	assert.Equal(t, 2, code)
	assert.Contains(t, stderr, "grep:")
}

// --- Long lines ---

func TestGrepPentestLineAtCapMinus1(t *testing.T) {
	dir := t.TempDir()
	content := make([]byte, grep.MaxLineBytes-1)
	for i := range content {
		content[i] = 'a'
	}
	content = append(content, '\n')
	require.NoError(t, os.WriteFile(filepath.Join(dir, "file.txt"), content, 0644))
	stdout, _, code := grepRun(t, "grep a file.txt", dir)
	assert.Equal(t, 0, code)
	assert.NotEmpty(t, stdout)
}

func TestGrepPentestLineAtCapExact(t *testing.T) {
	dir := t.TempDir()
	content := make([]byte, grep.MaxLineBytes)
	for i := range content {
		content[i] = 'a'
	}
	require.NoError(t, os.WriteFile(filepath.Join(dir, "file.txt"), content, 0644))
	_, _, code := grepRun(t, "grep a file.txt", dir)
	assert.Equal(t, 2, code)
}

func TestGrepPentestLineAtCapPlus1(t *testing.T) {
	dir := t.TempDir()
	content := make([]byte, grep.MaxLineBytes+1)
	for i := range content {
		content[i] = 'a'
	}
	require.NoError(t, os.WriteFile(filepath.Join(dir, "file.txt"), content, 0644))
	_, _, code := grepRun(t, "grep a file.txt", dir)
	assert.Equal(t, 2, code)
}

// --- Memory / resource exhaustion ---

func TestGrepPentestManyFiles(t *testing.T) {
	dir := t.TempDir()
	var args []string
	for i := 0; i < 200; i++ {
		name := "f" + strings.Repeat("0", 3-len(string(rune('0'+i/100)))) + string(rune('0'+i/100)) + string(rune('0'+i%100/10)) + string(rune('0'+i%10)) + ".txt"
		pentestWriteFile(t, dir, name, "line\n")
		args = append(args, name)
	}
	mustNotHang(t, func() {
		_, _, code := grepRun(t, "grep line "+strings.Join(args, " "), dir)
		assert.Equal(t, 0, code)
	})
}

func TestGrepPentestLargeFile(t *testing.T) {
	dir := t.TempDir()
	var sb strings.Builder
	for i := 0; i < 10000; i++ {
		sb.WriteString("line\n")
	}
	pentestWriteFile(t, dir, "big.txt", sb.String())
	mustNotHang(t, func() {
		stdout, _, code := grepRun(t, "grep -c line big.txt", dir)
		assert.Equal(t, 0, code)
		assert.Equal(t, "10000\n", stdout)
	})
}

// --- ReDoS protection ---

func TestGrepPentestLinearTimeRegex(t *testing.T) {
	// Go's regexp uses RE2, which guarantees linear-time matching.
	// This test verifies that pathological patterns complete quickly.
	dir := t.TempDir()
	// Create a string that would cause exponential backtracking in PCRE
	pentestWriteFile(t, dir, "file.txt", strings.Repeat("a", 100)+"\n")
	mustNotHang(t, func() {
		_, _, _ = grepRun(t, "grep -E 'a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?a?aaaaaaaaaaaaaaaaaaaa' file.txt", dir)
	})
}

// --- Pattern edge cases ---

func TestGrepPentestEmptyPattern(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "foo\nbar\n")
	mustNotHang(t, func() {
		stdout, _, code := grepRun(t, `grep '' file.txt`, dir)
		assert.Equal(t, 0, code)
		assert.Equal(t, "foo\nbar\n", stdout)
	})
}

func TestGrepPentestNullByteInFile(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "foo\x00bar\nbaz\n")
	mustNotHang(t, func() {
		_, _, code := grepRun(t, "grep foo file.txt", dir)
		assert.Equal(t, 0, code)
	})
}

func TestGrepPentestFixedStringsRegexInjection(t *testing.T) {
	// Verify -F prevents regex metacharacter interpretation.
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "a.*b\naxb\n")
	stdout, _, code := grepRun(t, "grep -F 'a.*b' file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a.*b\n", stdout)
}

// --- Special files / infinite sources ---

func TestGrepPentestDevNull(t *testing.T) {
	if os.DevNull == "NUL" {
		dir := t.TempDir()
		mustNotHang(t, func() {
			_, stderr, code := grepRun(t, "grep foo NUL", dir)
			assert.Equal(t, 2, code)
			assert.Contains(t, stderr, "grep:")
		})
		return
	}
	dir := t.TempDir()
	mustNotHang(t, func() {
		_, _, code := grepRun(t, "grep foo "+os.DevNull, dir, filepath.Dir(os.DevNull))
		assert.Equal(t, 1, code) // no match in empty file
	})
}

func TestGrepPentestContextCancelledDuringSearch(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", strings.Repeat("x\n", 10000))
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	mustNotHang(t, func() {
		grepRunCtx(ctx, t, "grep x file.txt", dir)
	})
}

// --- No-messages flag edge cases ---

func TestGrepPentestNoMessagesSuppressesErrors(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := grepRun(t, "grep -s foo nonexistent.txt", dir)
	assert.Equal(t, 2, code)
	assert.Equal(t, "", stderr)
}

// --- Multiple -e patterns ---

func TestGrepPentestManyPatterns(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "target\n")
	// Build a command with many -e flags
	var args []string
	for i := 0; i < 50; i++ {
		args = append(args, "-e", "pattern"+strings.Repeat("x", i))
	}
	args = append(args, "-e", "target")
	mustNotHang(t, func() {
		stdout, _, code := grepRun(t, "grep "+strings.Join(args, " ")+" file.txt", dir)
		assert.Equal(t, 0, code)
		assert.Equal(t, "target\n", stdout)
	})
}

// --- Quiet mode with error ---

func TestGrepPentestQuietWithMatch(t *testing.T) {
	dir := t.TempDir()
	pentestWriteFile(t, dir, "file.txt", "foo\n")
	stdout, stderr, code := grepRun(t, "grep -q foo file.txt", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "", stdout)
	assert.Equal(t, "", stderr)
}

// --- Context byte cap ---

// TestGrepBeforeContextByteCapEvictsOldLines verifies that the aggregate byte
// cap on the before-context sliding window (MaxContextBytes) is enforced.
// The file contains 1000 non-matching lines of ~1 KiB each (total ~1 MiB of
// context) followed by a single matching line. With -B 1000 the uncapped
// window would hold all 1000 lines (~1 MiB); the cap limits the window to
// MaxContextBytes (512 KiB), so the oldest lines are evicted and the total
// before-context output stays within the cap.
func TestGrepBeforeContextByteCapEvictsOldLines(t *testing.T) {
	dir := t.TempDir()
	linePayload := strings.Repeat("x", 1023) // 1023 bytes + newline = 1024 bytes/line
	var sb strings.Builder
	for i := 0; i < 1000; i++ {
		sb.WriteString(linePayload)
		sb.WriteByte('\n')
	}
	sb.WriteString("MATCH\n")
	pentestWriteFile(t, dir, "ctx.txt", sb.String())

	stdout, _, code := grepRun(t, "grep -B 1000 MATCH ctx.txt", dir)
	require.Equal(t, 0, code)
	// Output = (evicted-before-context lines) + "MATCH\n".
	// Before-context bytes must not exceed the cap.
	assert.LessOrEqual(t, len(stdout)-len("MATCH\n"), grep.MaxContextBytes,
		"before-context output exceeded MaxContextBytes cap")
}

// TestGrepAfterContextByteCapTruncatesOutput verifies that the aggregate byte
// cap on after-context output (MaxContextBytes) is enforced per match group.
// The file has one matching line followed by 1000 non-matching lines of ~1 KiB
// each (total ~1 MiB of potential after-context). With -A 1000 the uncapped
// stream would emit all 1000 lines; the cap stops emission once MaxContextBytes
// have been written for the group, so total after-context output stays within
// the cap.
func TestGrepAfterContextByteCapTruncatesOutput(t *testing.T) {
	dir := t.TempDir()
	linePayload := strings.Repeat("x", 1023) // 1023 bytes + newline = 1024 bytes/line
	var sb strings.Builder
	sb.WriteString("MATCH\n")
	for i := 0; i < 1000; i++ {
		sb.WriteString(linePayload)
		sb.WriteByte('\n')
	}
	pentestWriteFile(t, dir, "ctx.txt", sb.String())

	stdout, _, code := grepRun(t, "grep -A 1000 MATCH ctx.txt", dir)
	require.Equal(t, 0, code)
	// Output = "MATCH\n" + (after-context lines, capped at MaxContextBytes).
	assert.LessOrEqual(t, len(stdout)-len("MATCH\n"), grep.MaxContextBytes,
		"after-context output exceeded MaxContextBytes cap")
}

// --- GTFOBins validation ---

// TestGrepGTFOBinsFileReadSandboxEscape verifies that the GTFOBins file-read
// technique for grep is blocked by the AllowedPaths sandbox.
//
// GTFOBins: https://gtfobins.org/gtfobins/grep/
// Technique: grep ” /path/to/file
func TestGrepGTFOBinsFileReadSandboxEscape(t *testing.T) {
	allowed := t.TempDir()
	secret := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(secret, "secret.txt"), []byte("secret data\n"), 0644))
	secretPath := strings.ReplaceAll(filepath.Join(secret, "secret.txt"), `\`, `/`)
	_, stderr, code := grepRun(t, "grep '' "+secretPath, allowed)
	assert.Equal(t, 2, code) // grep uses exit code 2 for errors
	assert.Contains(t, stderr, "grep:")
}
