package windows

import (
	"bytes"
	"crypto/sha256"
	"errors"
	"strings"
	"testing"

	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
)

func TestResolveShellUsesOnlyConfiguredAbsoluteExecutable_BH_WINLAUNCH_01(t *testing.T) {
	t.Parallel()
	paths := ShellPaths{CMD: `C:\Trusted Tools\cmd.exe`, PowerShell: `C:\Trusted Tools\powershell.exe`}
	cmd, err := ResolveShell(rvboxv1.ShellType_SHELL_CMD, paths)
	if err != nil {
		t.Fatal(err)
	}
	if cmd.ApplicationName != paths.CMD || !equalStrings(cmd.Arguments, []string{"/D", "/S", "/C"}) || cmd.WrapperExtension != ".cmd" {
		t.Fatalf("cmd plan = %+v", cmd)
	}
	powershell, err := ResolveShell(rvboxv1.ShellType_SHELL_POWERSHELL, paths)
	if err != nil {
		t.Fatal(err)
	}
	if powershell.ApplicationName != paths.PowerShell || !equalStrings(powershell.Arguments, []string{"-NoLogo", "-NoProfile", "-NonInteractive", "-File"}) || powershell.WrapperExtension != ".ps1" || powershell.Encoding != WrapperEncodingUTF8BOM {
		t.Fatalf("PowerShell plan = %+v", powershell)
	}
	// A hostile request PATH or COMSPEC is not an input to resolution. The
	// configured absolute path remains the only application name.
	if strings.Contains(cmd.ApplicationName, "untrusted") || strings.Contains(powershell.ApplicationName, "untrusted") {
		t.Fatal("request environment influenced shell resolution")
	}
}

func TestResolveShellRejectsUnsupportedOrUnsafePaths_BH_WINLAUNCH_01(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name  string
		shell rvboxv1.ShellType
		paths ShellPaths
	}{
		{"unspecified", rvboxv1.ShellType_SHELL_TYPE_UNSPECIFIED, ShellPaths{CMD: `C:\Windows\System32\cmd.exe`}},
		{"unix shell", rvboxv1.ShellType_SHELL_BASH, ShellPaths{CMD: `C:\Windows\System32\cmd.exe`}},
		{"relative", rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: `cmd.exe`}},
		{"drive relative", rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: `C:cmd.exe`}},
		{"wrong extension", rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: `C:\Windows\System32\cmd.dll`}},
		{"trailing space", rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: "C:\\Windows\\System32\\cmd.exe "}},
		{"NUL", rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: "C:\\Windows\\System32\\cmd.exe\x00"}},
	}
	for _, test := range cases {
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			if _, err := ResolveShell(test.shell, test.paths); err == nil {
				t.Fatal("unsafe shell configuration accepted")
			}
		})
	}
}

func TestBuildWrapperValidatesAndEncodesSources_HP_WINLAUNCH_02(t *testing.T) {
	t.Parallel()
	command := &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo hello\r\n"}}
	wrapper, err := BuildWrapper(command, nil, 1024)
	if err != nil {
		t.Fatal(err)
	}
	if wrapper.Extension != ".cmd" || wrapper.Encoding != WrapperEncodingUTF8 || !bytes.Equal(wrapper.Bytes, []byte("echo hello\r\n")) {
		t.Fatalf("cmd wrapper = %+v", wrapper)
	}
	if wrapper.SHA256 != sha256.Sum256(wrapper.Bytes) {
		t.Fatal("wrapper digest does not cover materialized bytes")
	}

	body := []byte("Write-Output 'héllo'\n")
	digest := sha256.Sum256(body)
	powershell := &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_POWERSHELL, Source: &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: "display.ps1", SizeBytes: uint64(len(body)), Sha256: digest[:]}}}
	encoded, err := BuildWrapper(powershell, body, 1024)
	if err != nil {
		t.Fatal(err)
	}
	if encoded.Extension != ".ps1" || encoded.Encoding != WrapperEncodingUTF8BOM || !bytes.HasPrefix(encoded.Bytes, []byte{0xef, 0xbb, 0xbf}) || !bytes.HasSuffix(encoded.Bytes, body) {
		t.Fatalf("PowerShell wrapper bytes = %x", encoded.Bytes)
	}
}

func TestBuildWrapperRejectsMismatchNULAndLimit_BH_WINLAUNCH_02(t *testing.T) {
	t.Parallel()
	body := []byte("echo expected")
	digest := sha256.Sum256(body)
	script := func(size uint64, hash []byte, value []byte) *rvboxv1.ExecutionSpec {
		return &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_Script{Script: &rvboxv1.ScriptDescriptor{Filename: "x.cmd", SizeBytes: size, Sha256: hash}}}
	}
	cases := []struct {
		name string
		spec *rvboxv1.ExecutionSpec
		body []byte
		want error
	}{
		{"length", script(uint64(len(body)+1), digest[:], body), body, ErrWrapperLengthMismatch},
		{"digest", script(uint64(len(body)), make([]byte, sha256.Size), body), body, ErrWrapperDigestMismatch},
		{"NUL", &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "echo\x00bad"}}, nil, ErrInvalidWrapperSource},
		{"limit", &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_CMD, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "123456"}}, nil, ErrWrapperTooLarge},
	}
	for _, test := range cases {
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			_, err := BuildWrapper(test.spec, test.body, map[bool]uint64{true: 5, false: 1024}[test.name == "limit"])
			if !errors.Is(err, test.want) {
				t.Fatalf("error = %v, want %v", err, test.want)
			}
		})
	}
}

func TestBuildLaunchPlanQuotesWrapperAndKeepsWorkingDirectory(t *testing.T) {
	t.Parallel()
	plan, err := ResolveShell(rvboxv1.ShellType_SHELL_CMD, ShellPaths{CMD: `C:\Windows\System32\cmd.exe`})
	if err != nil {
		t.Fatal(err)
	}
	launch, err := plan.BuildLaunchPlan(`C:\ProgramData\RVBox\work\issue one\wrapper.cmd`, `C:\ProgramData\RVBox\work\issue one`, []uint16{0})
	if err != nil {
		t.Fatal(err)
	}
	if launch.ApplicationName != plan.ApplicationName || launch.WorkingDirectory == "" || !strings.Contains(launch.CommandLine, `"C:\ProgramData\RVBox\work\issue one\wrapper.cmd"`) {
		t.Fatalf("launch plan = %+v", launch)
	}
	if len(launch.Environment) != 1 || launch.Environment[0] != 0 {
		t.Fatalf("environment was not copied: %v", launch.Environment)
	}
	if _, err := plan.BuildLaunchPlan(`C:\ProgramData\RVBox\work\wrapper.ps1`, `C:\ProgramData\RVBox\work`, nil); !errors.Is(err, ErrInvalidWrapperPath) {
		t.Fatalf("extension mismatch error = %v", err)
	}
	if _, err := plan.BuildLaunchPlan(`relative\wrapper.cmd`, `C:\ProgramData\RVBox\work`, nil); !errors.Is(err, ErrInvalidWrapperPath) {
		t.Fatalf("relative wrapper error = %v", err)
	}
}

func equalStrings(left, right []string) bool {
	if len(left) != len(right) {
		return false
	}
	for index := range left {
		if left[index] != right[index] {
			return false
		}
	}
	return true
}