package winpwn

import (
	"os"
	"testing"
)

// testFixturePE is a real PE32+ checked into the repo (examples/bof_basic),
// used by the ROP-level tests below to validate parsing against an actual
// binary rather than only hand-built Gadget slices.
const testFixturePE = "examples/bof_basic/bof_win.c.exe"

func requireFixturePE(t *testing.T) {
	if _, err := os.Stat(testFixturePE); err != nil {
		t.Skipf("fixture %s not present: %v", testFixturePE, err)
	}
}

// requireRPWinTool skips a test if rp-win.exe isn't resolvable -- NewROP
// shells out to it, so integration tests that build a real ROP need it
// installed the same way a real target binary needs to exist.
func requireRPWinTool(t *testing.T) {
	toolPath := os.Getenv("RP_WIN_EXE")
	if toolPath == "" {
		toolPath = defaultRPWinTool
	}
	if _, err := os.Stat(toolPath); err != nil {
		t.Skipf("rp-win.exe not present at %s: %v", toolPath, err)
	}
}

// The Find/SearchRegex tests below construct a *ROP directly from a
// hand-built gadgets slice, bypassing NewROP entirely -- these two methods
// only ever read r.gadgets, so this exercises the actual filtering/sorting
// logic without needing rp-win.exe or a real PE on the test machine.

func TestROPFindSortsByAddressAscendingAsTiebreak(t *testing.T) {
	r := &ROP{gadgets: []Gadget{
		{Address: 0x3000, Instructions: "pop rcx ; ret"},
		{Address: 0x1000, Instructions: "pop rcx ; ret"},
		{Address: 0x2000, Instructions: "pop rcx ; ret"},
	}}

	found := r.Find("pop rcx ; ret")
	if len(found) != 3 {
		t.Fatalf("expected 3 matches, got %d", len(found))
	}
	for i := 1; i < len(found); i++ {
		if found[i-1].Address > found[i].Address {
			t.Fatalf("results not sorted ascending among equal-quality matches: %v", found)
		}
	}
	if found[0].Address != 0x1000 {
		t.Errorf("Find(...)[0] should be the lowest address among ties, got 0x%X", found[0].Address)
	}
}

func TestROPFindPrefersExactMatchOverDirtySubstring(t *testing.T) {
	// Regression test for a real bug caught against kernel32.dll: sorting
	// purely by address let a "dirty" longer gadget win [0] over the clean
	// one just because it happened to start a few bytes earlier in memory
	// (0x1800198B7 < 0x1800198BB numerically, even though only the latter
	// is a bare "pop rcx ; ret" with no side effects).
	r := &ROP{gadgets: []Gadget{
		{Address: 0x1800198B7, Instructions: "ror byte [rax-0x1], 0x15 ; pop rcx ; ret"},
		{Address: 0x1800198BB, Instructions: "pop rcx ; ret"},
	}}

	found := r.Find("pop rcx ; ret")
	if len(found) != 2 {
		t.Fatalf("expected 2 matches, got %d", len(found))
	}
	if found[0].Address != 0x1800198BB || found[0].Instructions != "pop rcx ; ret" {
		t.Errorf("Find(...)[0] should be the clean exact-match gadget, got %+v", found[0])
	}
}

func TestROPFindSubstringMatchIsCaseInsensitive(t *testing.T) {
	r := &ROP{gadgets: []Gadget{
		{Address: 0x1000, Instructions: "POP RCX ; RET"},
		{Address: 0x2000, Instructions: "pop rdx ; ret"},
	}}

	found := r.Find("pop rcx ; ret")
	if len(found) != 1 || found[0].Address != 0x1000 {
		t.Errorf("expected a case-insensitive match on 0x1000, got %v", found)
	}
}

func TestROPFindReturnsEmptyWhenNoMatch(t *testing.T) {
	r := &ROP{gadgets: []Gadget{{Address: 0x1000, Instructions: "pop rcx ; ret"}}}

	found := r.Find("pop rbp ; ret")
	if len(found) != 0 {
		t.Errorf("expected no matches, got %v", found)
	}
}

func TestROPFindEmptyIndexPanics(t *testing.T) {
	// Documented behavior: indexing an empty Find() result panics rather
	// than silently handing back a zero-value Gadget -- an exploit script
	// should fail loudly at the gadget lookup, not against a garbage
	// address three chain-steps later.
	defer func() {
		if recover() == nil {
			t.Error("expected indexing an empty Find() result to panic")
		}
	}()
	r := &ROP{gadgets: nil}
	_ = r.Find("nonexistent")[0]
}

func TestROPSearchRegexOnHandBuiltGadgets(t *testing.T) {
	r := &ROP{gadgets: []Gadget{
		{Address: 0x1000, Instructions: "pop rcx ; ret"},
		{Address: 0x2000, Instructions: "pop rdx ; ret"},
		{Address: 0x3000, Instructions: "mov [rcx], eax ; ret"},
	}}

	gadgets, err := r.SearchRegex(`^pop r\w+ ; ret$`)
	if err != nil {
		t.Fatal(err)
	}
	if len(gadgets) != 2 {
		t.Errorf("expected 2 pop-reg-then-ret gadgets, got %d: %v", len(gadgets), gadgets)
	}
}

func TestROPFindsRetGadget(t *testing.T) {
	requireFixturePE(t)
	requireRPWinTool(t)
	rop, err := NewROP(testFixturePE)
	if err != nil {
		t.Fatal(err)
	}
	defer rop.Close()

	gadgets, err := rop.Search("ret")
	if err != nil {
		t.Fatal(err)
	}
	if len(gadgets) == 0 {
		t.Error("expected at least one 'ret' gadget in a real x64 binary")
	}
}

func TestROPSearchRegex(t *testing.T) {
	requireFixturePE(t)
	requireRPWinTool(t)
	rop, err := NewROP(testFixturePE)
	if err != nil {
		t.Fatal(err)
	}
	defer rop.Close()

	gadgets, err := rop.SearchRegex(`^pop r\w+ ; ret$`)
	if err != nil {
		t.Fatal(err)
	}
	if len(gadgets) == 0 {
		t.Error("expected at least one pop-reg-then-ret gadget")
	}
}