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

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

	"github.com/DataDog/rshell/interp"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const pentestTimeout = 10 * time.Second

func pentestDir(t *testing.T, files map[string]string) string {
	t.Helper()
	dir := t.TempDir()
	for name, content := range files {
		require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0644))
	}
	return dir
}

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

func TestPentestUnknownFlags(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	for _, flag := range []string{"-f", "--follow", "--no-such-flag", "-z", "-s"} {
		t.Run(flag, func(t *testing.T) {
			_, stderr, code := cmdRun(t, "sed "+flag+" 's/a/b/' f.txt", dir)
			assert.Equal(t, 1, code)
			assert.Contains(t, stderr, "sed:")
		})
	}
}

func TestPentestRejectedInPlaceFlag(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "hello\n"})
	_, stderr, code := cmdRun(t, `sed -i 's/hello/bye/' f.txt`, dir)
	assert.NotEqual(t, 0, code)
	assert.Contains(t, stderr, "sed:")
	// Verify file was NOT modified.
	data, err := os.ReadFile(filepath.Join(dir, "f.txt"))
	require.NoError(t, err)
	assert.Equal(t, "hello\n", string(data))
}

func TestPentestDoubleDashEndOfFlags(t *testing.T) {
	dir := pentestDir(t, map[string]string{"-n": "hello\n"})
	// -- should allow flag-like filenames.
	stdout, _, code := runScript(t, `sed 's/hello/bye/' -- -n`, dir, interp.AllowedPaths([]string{dir}))
	assert.Equal(t, 0, code)
	assert.Equal(t, "bye\n", stdout)
}

func TestPentestFlagViaWordExpansion(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	// Flag injection via variable.
	_, stderr, code := cmdRun(t, `flag="-f"; sed $flag 's/a/b/' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

// --- Blocked commands ---

func TestPentestAllBlockedCommands(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	blocked := []struct {
		name, script string
	}{
		{"e-command", `sed 'e' f.txt`},
		{"w-command", `sed 'w /tmp/evil' f.txt`},
		{"W-command", `sed 'W /tmp/evil' f.txt`},
		{"r-command", `sed 'r /etc/passwd' f.txt`},
		{"R-command", `sed 'R /etc/passwd' f.txt`},
		{"s-w-flag", `sed 's/t/T/w /tmp/evil' f.txt`},
		{"s-e-flag", `sed 's/t/T/e' f.txt`},
	}
	for _, tc := range blocked {
		t.Run(tc.name, func(t *testing.T) {
			_, stderr, code := cmdRun(t, tc.script, dir)
			assert.Equal(t, 1, code)
			assert.Contains(t, stderr, "blocked")
		})
	}
}

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

func TestPentestLongLineNearLimit(t *testing.T) {
	// Line of MaxLineBytes - 1 bytes should succeed.
	const limit = 1 << 20 // 1 MiB
	dir := pentestDir(t, map[string]string{
		"near.txt": strings.Repeat("x", limit-1) + "\n",
	})
	stdout, _, code := cmdRun(t, `sed 's/x/y/' near.txt`, dir)
	assert.Equal(t, 0, code)
	assert.True(t, strings.HasPrefix(stdout, "y"))
}

func TestPentestPatternSpaceExhaustion(t *testing.T) {
	// Try to exceed MaxSpaceBytes via N command.
	var sb strings.Builder
	for i := 0; i < 2000; i++ {
		sb.WriteString(strings.Repeat("a", 600))
		sb.WriteByte('\n')
	}
	dir := pentestDir(t, map[string]string{"big.txt": sb.String()})
	_, stderr, code := cmdRun(t, `sed ':a;N;ba' big.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "pattern space exceeded size limit")
}

func TestPentestHoldSpaceExhaustion(t *testing.T) {
	var sb strings.Builder
	for i := 0; i < 2000; i++ {
		sb.WriteString(strings.Repeat("b", 600))
		sb.WriteByte('\n')
	}
	dir := pentestDir(t, map[string]string{"big.txt": sb.String()})
	_, stderr, code := cmdRun(t, `sed 'H' big.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "hold space exceeded size limit")
}

func TestPentestBranchLoopLimit(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	ctx, cancel := context.WithTimeout(context.Background(), pentestTimeout)
	defer cancel()
	_, stderr, code := runScriptCtx(ctx, t, `sed ':loop;b loop' f.txt`, dir, interp.AllowedPaths([]string{dir}))
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "branch loop limit exceeded")
}

func TestPentestSubstitutionGrowsPatternSpace(t *testing.T) {
	// s///g with a replacement much larger than the match can grow pattern space.
	dir := pentestDir(t, map[string]string{
		"f.txt": strings.Repeat("x", 100000) + "\n",
	})
	// Replace each 'x' with 20 chars — would grow to 2MB, exceeding MaxSpaceBytes (1MB).
	_, stderr, code := cmdRun(t, `sed 's/x/xxxxxxxxxxxxxxxxxxxx/g' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "pattern space exceeded size limit")
}

func TestPentestSmallFileWithLargeNthSub(t *testing.T) {
	// Substitution with very high N should not cause issues.
	dir := pentestDir(t, map[string]string{"f.txt": "aaa\n"})
	stdout, _, code := cmdRun(t, `sed 's/a/X/999' f.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "aaa\n", stdout) // No 999th occurrence, no change.
}

// --- Special files ---

func TestPentestDevNull(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("skipping /dev/null test on Windows")
	}
	stdout, _, code := runScript(t, `sed 's/a/b/' /dev/null`, "", interp.AllowedPaths([]string{"/dev"}))
	assert.Equal(t, 0, code)
	assert.Equal(t, "", stdout)
}

func TestPentestContextCancelledDuringProcessing(t *testing.T) {
	dir := pentestDir(t, map[string]string{
		"big.txt": strings.Repeat("line\n", 100000),
	})
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	_, _, code := runScriptCtx(ctx, t, `sed 's/line/LINE/g' big.txt`, dir, interp.AllowedPaths([]string{dir}))
	// Should either complete or be cancelled — both are acceptable.
	_ = code
}

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

func TestPentestNonExistentFile(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := cmdRun(t, `sed 's/a/b/' nonexistent.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "nonexistent.txt")
}

func TestPentestDirectoryAsFile(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.Mkdir(filepath.Join(dir, "subdir"), 0755))
	_, stderr, code := cmdRun(t, `sed 's/a/b/' subdir`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

func TestPentestEmptyStringFilename(t *testing.T) {
	_, stderr, code := cmdRun(t, `sed 's/a/b/' ""`, "")
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

func TestPentestMultipleStdinArgs(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "hello\n"})
	// Multiple - args: stdin should only be consumed once.
	stdout, _, code := cmdRun(t, `echo hello | sed 's/hello/bye/' - -`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "bye\n", stdout)
}

func TestPentestManyFiles(t *testing.T) {
	// Verify no FD leak with many files.
	files := make(map[string]string)
	var names []string
	for i := 0; i < 50; i++ {
		name := "f" + strings.Repeat("0", 3-len(strings.TrimLeft(string(rune('0'+i%10)), ""))) + ".txt"
		name = filepath.Base(name)
		actualName := "file_" + string(rune('a'+i%26)) + string(rune('a'+i/26)) + ".txt"
		files[actualName] = "line\n"
		names = append(names, actualName)
	}
	dir := pentestDir(t, files)
	script := "sed 's/line/LINE/' " + strings.Join(names, " ")
	stdout, _, code := cmdRun(t, script, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, strings.Repeat("LINE\n", len(names)), stdout)
}

// --- Input validation ---

func TestPentestInvalidRegexPatterns(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	patterns := []string{
		`sed 's/[/x/' f.txt`,    // Unterminated character class
		`sed '/[/d' f.txt`,      // Invalid address regex
		`sed -E 's/(/x/' f.txt`, // Unmatched paren in ERE
	}
	for _, script := range patterns {
		t.Run(script, func(t *testing.T) {
			_, stderr, code := cmdRun(t, script, dir)
			assert.Equal(t, 1, code)
			assert.Contains(t, stderr, "sed:")
		})
	}
}

func TestPentestEmptyScript(t *testing.T) {
	// Empty script is valid in GNU sed.
	_, _, code := cmdRun(t, `sed '' /dev/null`, "")
	_ = code // just don't crash
}

func TestPentestNoScriptNoFiles(t *testing.T) {
	_, stderr, code := cmdRun(t, `sed`, "")
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

func TestPentestUnterminatedSubstitution(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	_, _, code := cmdRun(t, `sed 's/foo' f.txt`, dir)
	// Just don't crash.
	_ = code
}

func TestPentestUnterminatedGroup(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	_, stderr, code := cmdRun(t, `sed '{p' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "unterminated")
}

func TestPentestDeeplyNestedGroups(t *testing.T) {
	// Many nested groups should not crash.
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	script := strings.Repeat("{", 50) + "p" + strings.Repeat("}", 50)
	stdout, _, code := cmdRun(t, `sed '`+script+`' f.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Contains(t, stdout, "test")
}

// --- Regex safety (ReDoS) ---

func TestPentestReDoSPattern(t *testing.T) {
	// RE2 should handle this in linear time.
	dir := pentestDir(t, map[string]string{
		"f.txt": strings.Repeat("a", 100) + "\n",
	})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _, code := runScriptCtx(ctx, t, `sed -E 's/(a+)+b/x/' f.txt`, dir, interp.AllowedPaths([]string{dir}))
	assert.Equal(t, 0, code)
}

// --- P and D commands ---

func TestPentestPCommandMultiline(t *testing.T) {
	dir := pentestDir(t, map[string]string{
		"f.txt": "line1\nline2\n",
	})
	stdout, _, code := cmdRun(t, `sed -n 'N;P' f.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "line1\n", stdout)
}

func TestPentestDCommandMultiline(t *testing.T) {
	dir := pentestDir(t, map[string]string{
		"f.txt": "a\nb\nc\n",
	})
	// D deletes first line of pattern space and restarts cycle.
	// N appends next line, D removes the first.
	stdout, _, code := cmdRun(t, `sed 'N;P;D' f.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\nb\nc\n", stdout)
}

// --- Quit commands with exit codes ---

func TestPentestQuitWithCode(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "line1\nline2\nline3\n"})
	_, _, code := cmdRun(t, `sed 'q 42' f.txt`, dir)
	assert.Equal(t, 42, code)
}

func TestPentestQuitNoPrintWithCode(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "line1\nline2\nline3\n"})
	stdout, _, code := cmdRun(t, `sed 'Q 7' f.txt`, dir)
	assert.Equal(t, 7, code)
	assert.Equal(t, "", stdout)
}

// --- Transliterate edge cases ---

func TestPentestTransliterateMismatch(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	_, stderr, code := cmdRun(t, `sed 'y/abc/de/' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "same length")
}

func TestPentestTransliterateEmpty(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	stdout, _, code := cmdRun(t, `sed 'y///' f.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "test\n", stdout) // No transliteration, identity.
}

// --- Symlinks ---

func TestPentestSymlinkToRegularFile(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(dir, "real.txt"), []byte("hello\n"), 0644))
	require.NoError(t, os.Symlink("real.txt", filepath.Join(dir, "link.txt")))
	stdout, _, code := cmdRun(t, `sed 's/hello/bye/' link.txt`, dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "bye\n", stdout)
}

func TestPentestDanglingSymlink(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.Symlink(filepath.Join(dir, "nonexistent"), filepath.Join(dir, "dangling.txt")))
	_, stderr, code := cmdRun(t, `sed 's/a/b/' dangling.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

// --- GTFOBins validation ---

// TestSedGTFOBinsFileReadSandboxEscape verifies that the GTFOBins file-read
// technique for sed is blocked by the AllowedPaths sandbox.
//
// GTFOBins: https://gtfobins.org/gtfobins/sed/
// Technique: sed ” /path/to/file
func TestSedGTFOBinsFileReadSandboxEscape(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 := runScript(t, "sed '' "+secretPath, allowed, interp.AllowedPaths([]string{allowed}))
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sed:")
}

// TestSedGTFOBinsShellEscapeBlocked verifies that the GTFOBins shell-escape
// technique for sed (the 'e' command) is blocked.
//
// GTFOBins: https://gtfobins.org/gtfobins/sed/
// Technique: sed -n '1e exec sh 1>&0' /dev/null
func TestSedGTFOBinsShellEscapeBlocked(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	_, stderr, code := cmdRun(t, `sed 'e' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "blocked")
}

// TestSedGTFOBinsFileWriteBlocked verifies that the GTFOBins file-write
// technique for sed (the 'w' command) is blocked.
//
// GTFOBins: https://gtfobins.org/gtfobins/sed/
// Technique: sed -n "s/.*/$data/w /path/to/output" /dev/null
func TestSedGTFOBinsFileWriteBlocked(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "test\n"})
	_, stderr, code := cmdRun(t, `sed 'w /tmp/evil' f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "blocked")
}

// TestSedGTFOBinsInPlaceBlocked verifies that the GTFOBins in-place edit
// technique for sed (-i flag) is blocked.
//
// GTFOBins: https://gtfobins.org/gtfobins/sed/
// Technique: sed -i ” /path/to/file
func TestSedGTFOBinsInPlaceBlocked(t *testing.T) {
	dir := pentestDir(t, map[string]string{"f.txt": "hello\n"})
	_, stderr, code := cmdRun(t, `sed -i 's/hello/bye/' f.txt`, dir)
	assert.NotEqual(t, 0, code)
	assert.Contains(t, stderr, "sed:")
	// Verify file was NOT modified.
	data, err := os.ReadFile(filepath.Join(dir, "f.txt"))
	require.NoError(t, err)
	assert.Equal(t, "hello\n", string(data))
}
