//go:build windows

package winpwn

import (
	"syscall"
	"testing"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
)

var (
	user32ForTest      = windows.NewLazySystemDLL("user32.dll")
	procEnumWindows    = user32ForTest.NewProc("EnumWindows")
	procGetWindowTextW = user32ForTest.NewProc("GetWindowTextW")
	procPostMessageW   = user32ForTest.NewProc("PostMessageW")
)

const wmClose = 0x0010

// TestShellcodeMessageBoxA runs the real shellcode (via ExecuteShellcode)
// and checks for a real window with the expected title -- not just "didn't
// crash". This proves resolve_export actually found LoadLibraryA in
// kernel32, loaded user32.dll (not guaranteed loaded in a test binary),
// resolved MessageBoxA inside it, and called it with the right calling
// convention. The window is dismissed programmatically (WM_CLOSE) so the
// test doesn't hang waiting for a human.
func TestShellcodeMessageBoxA(t *testing.T) {
	const wantTitle = "winpwn test"
	code, err := ShellcodeMessageBoxA("hello from winpwn", wantTitle)
	if err != nil {
		t.Fatal(err)
	}

	myPID := windows.GetCurrentProcessId()
	done := make(chan error, 1)
	go func() {
		done <- ExecuteShellcode(code)
	}()

	var found windows.HWND
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) && found == 0 {
		cb := syscall.NewCallback(func(hwnd windows.HWND, _ uintptr) uintptr {
			var pid uint32
			windows.GetWindowThreadProcessId(hwnd, &pid)
			if pid != myPID {
				return 1 // continue enumeration
			}
			buf := make([]uint16, 256)
			procGetWindowTextW.Call(uintptr(hwnd), uintptr(unsafe.Pointer(&buf[0])), uintptr(len(buf)))
			if windows.UTF16ToString(buf) == wantTitle {
				found = hwnd
				return 0 // stop enumeration
			}
			return 1
		})
		procEnumWindows.Call(cb, 0)
		if found == 0 {
			time.Sleep(100 * time.Millisecond)
		}
	}

	if found == 0 {
		t.Fatalf("never saw a real window titled %q -- MessageBoxA shellcode did not pop a window", wantTitle)
	}

	procPostMessageW.Call(uintptr(found), wmClose, 0, 0)

	select {
	case err := <-done:
		if err != nil {
			t.Errorf("ExecuteShellcode returned an error after dismissal: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("shellcode goroutine never returned after WM_CLOSE")
	}
}