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

// Exploratory pentest for the sort builtin.
//
// These tests probe rejected flags, memory safety, path edge cases,
// and flag injection scenarios. Tests that might hang are run in a
// goroutine with time.After to bound execution.

package sort_test

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

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

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

const pentestTimeout = 10 * time.Second

// sortRun is a shorthand for runScript with AllowedPaths=dir.
func sortRun(t *testing.T, script, dir string, extraPaths ...string) (stdout, stderr string, exitCode int) {
	t.Helper()
	paths := append([]string{dir}, extraPaths...)
	return runScript(t, script, dir, interp.AllowedPaths(paths))
}

// mustNotHang runs f in a goroutine and fails the test if it does not return
// within pentestTimeout.
func mustNotHang(t *testing.T, f func()) {
	t.Helper()
	done := make(chan struct{})
	go func() {
		defer close(done)
		f()
	}()
	select {
	case <-done:
	case <-time.After(pentestTimeout):
		t.Fatalf("operation did not complete within %s", pentestTimeout)
	}
}

// --- Rejected flags (GTFOBins vectors) ---

func TestCmdPentestOutputFlagRejected(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort -o out.txt f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
	// Verify no file was created.
	_, err := os.Stat(filepath.Join(dir, "out.txt"))
	assert.True(t, os.IsNotExist(err))
}

func TestCmdPentestOutputFlagLong(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort --output=out.txt f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

func TestCmdPentestCompressProgramRejected(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort --compress-program=sh f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

func TestCmdPentestTempDirRejected(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort --temporary-directory=/tmp f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

// --- Path traversal ---

func TestCmdPentestPathTraversal(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := sortRun(t, "sort ../../etc/passwd", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

func TestCmdPentestOutsideSandbox(t *testing.T) {
	allowed := t.TempDir()
	secret := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(secret, "s.txt"), []byte("secret"), 0644))
	secretPath := filepath.ToSlash(filepath.Join(secret, "s.txt"))
	_, stderr, code := runScript(t, "sort "+secretPath, allowed, interp.AllowedPaths([]string{allowed}))
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

// --- Nonexistent and empty files ---

func TestCmdPentestNonexistentFile(t *testing.T) {
	dir := t.TempDir()
	_, stderr, code := sortRun(t, "sort does_not_exist.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

func TestCmdPentestEmptyFilename(t *testing.T) {
	dir := t.TempDir()
	_, _, code := sortRun(t, `sort ""`, dir)
	assert.Equal(t, 1, code)
}

func TestCmdPentestDirectoryAsFile(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.Mkdir(filepath.Join(dir, "subdir"), 0755))
	_, stderr, code := sortRun(t, "sort subdir", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

// --- Memory safety ---

func TestCmdPentestLargeFile(t *testing.T) {
	// A file with 10000 lines should sort without hanging.
	dir := t.TempDir()
	var buf bytes.Buffer
	for i := 10000; i > 0; i-- {
		buf.WriteString(fmt.Sprintf("%d\n", i))
	}
	require.NoError(t, os.WriteFile(filepath.Join(dir, "big.txt"), buf.Bytes(), 0644))
	mustNotHang(t, func() {
		stdout, _, code := sortRun(t, "sort -n big.txt", dir)
		assert.Equal(t, 0, code)
		lines := strings.Split(strings.TrimSuffix(stdout, "\n"), "\n")
		assert.Equal(t, 10000, len(lines))
		assert.Equal(t, "1", lines[0])
		assert.Equal(t, "10000", lines[len(lines)-1])
	})
}

func TestCmdPentestLongLine(t *testing.T) {
	// A line just below the 1 MiB cap should succeed.
	dir := t.TempDir()
	line := bytes.Repeat([]byte("a"), 1<<20-2)
	line = append(line, '\n')
	require.NoError(t, os.WriteFile(filepath.Join(dir, "long.txt"), line, 0644))
	mustNotHang(t, func() {
		stdout, _, code := sortRun(t, "sort long.txt", dir)
		assert.Equal(t, 0, code)
		assert.Equal(t, string(line), stdout)
	})
}

func TestCmdPentestLongLineExceedsCap(t *testing.T) {
	// A line exceeding the 1 MiB cap should error, not crash.
	dir := t.TempDir()
	content := bytes.Repeat([]byte("a"), 1<<20+1)
	require.NoError(t, os.WriteFile(filepath.Join(dir, "huge.txt"), content, 0644))
	mustNotHang(t, func() {
		_, stderr, code := sortRun(t, "sort huge.txt", dir)
		assert.Equal(t, 1, code)
		assert.Contains(t, stderr, "sort:")
	})
}

// --- Flag injection ---

func TestCmdPentestFlagViaExpansion(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, `flag="--output=evil.txt"; sort $flag f.txt`, dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
	// Verify no file was created.
	_, err := os.Stat(filepath.Join(dir, "evil.txt"))
	assert.True(t, os.IsNotExist(err))
}

func TestCmdPentestUnknownLongFlag(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort --no-such-flag f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

func TestCmdPentestUnknownShortFlag(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "f.txt", "hello\n")
	_, stderr, code := sortRun(t, "sort -Z f.txt", dir)
	assert.Equal(t, 1, code)
	assert.Contains(t, stderr, "sort:")
}

// --- Input size cap ---

// TestSortInputExceedsMaxTotalBytes verifies that sort rejects input larger
// than MaxTotalBytes (5 MiB) with a descriptive error so callers know what
// limit was hit and can pre-filter their data.
//
// sort counts content bytes via sc.Text() which strips the line terminator,
// so we use lines of 1000 'a' chars (1001 bytes on disk, 1000 counted) to
// keep the content-to-file-size ratio high and ensure we reliably cross the
// limit.
func TestSortInputExceedsMaxTotalBytes(t *testing.T) {
	dir := t.TempDir()
	// Each line: 1000 content bytes + newline. Lines needed to exceed cap:
	// ceil(MaxTotalBytes/1000) + 1.
	line := []byte(strings.Repeat("a", 1000) + "\n")
	nLines := sort.MaxTotalBytes/1000 + 2
	content := bytes.Repeat(line, nLines)
	require.NoError(t, os.WriteFile(filepath.Join(dir, "big.txt"), content, 0644))

	mustNotHang(t, func() {
		_, stderr, code := sortRun(t, "sort big.txt", dir)
		assert.Equal(t, 1, code)
		assert.Contains(t, stderr, "sort:")
		assert.Contains(t, stderr, "exceeds maximum")
		assert.Contains(t, stderr, "5 MiB")
	})
}

// TestSortInputBelowMaxTotalBytes verifies that sort succeeds when input is
// just below the MaxTotalBytes cap.
func TestSortInputBelowMaxTotalBytes(t *testing.T) {
	dir := t.TempDir()
	// Each line: 1000 content bytes + newline; nLines × 1000 < MaxTotalBytes
	// and nLines << MaxLines (1,000,000) so neither cap is hit.
	line := []byte(strings.Repeat("a", 1000) + "\n")
	nLines := sort.MaxTotalBytes/1000 - 2
	content := bytes.Repeat(line, nLines)
	require.NoError(t, os.WriteFile(filepath.Join(dir, "ok.txt"), content, 0644))

	mustNotHang(t, func() {
		_, stderr, code := sortRun(t, "sort ok.txt", dir)
		assert.Equal(t, 0, code)
		assert.Empty(t, stderr)
	})
}

// --- Double dash ---

func TestCmdPentestFlagLikeName(t *testing.T) {
	dir := t.TempDir()
	require.NoError(t, os.WriteFile(filepath.Join(dir, "-r"), []byte("flag-file\n"), 0644))
	stdout, _, code := sortRun(t, "sort -- -r", dir)
	assert.Equal(t, 0, code)
	assert.Equal(t, "flag-file\n", stdout)
}

// --- Nil stdin ---

func TestCmdPentestNilStdin(t *testing.T) {
	dir := t.TempDir()
	stdout, stderr, code := runScript(t, "sort -", dir, interp.AllowedPaths([]string{dir}))
	assert.Equal(t, 0, code)
	assert.Equal(t, "", stdout)
	assert.Equal(t, "", stderr)
}
