//go:build windows

package winpwn

import (
	"debug/pe"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"unsafe"

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

// ProcessMemory is an io.ReaderAt/io.WriterAt over a remote process's
// address space, anchored at a base address -- the live-memory analogue of
// reading bytes off disk. Offset 0 in ReadAt/WriteAt means "Base itself",
// matching how a memory-backed PEFile's RVAToFileOffset treats RVAs as
// identical to read offsets: that's exactly what an RVA means once an
// image is actually loaded.
type ProcessMemory struct {
	Handle windows.Handle
	Base   uintptr
}

// OpenProcessMemory opens pid for VM read/write, anchored at base -- the
// building block behind OpenPEFromProcess, but also useful standalone for
// any arbitrary-read/write-shaped primitive once you have a target address
// (Phase 8's ARW interface is satisfied directly by *ProcessMemory).
func OpenProcessMemory(pid uint32, base uintptr) (*ProcessMemory, error) {
	h, err := windows.OpenProcess(
		windows.PROCESS_QUERY_INFORMATION|windows.PROCESS_VM_READ|windows.PROCESS_VM_WRITE|windows.PROCESS_VM_OPERATION,
		false, pid)
	if err != nil {
		return nil, err
	}
	return &ProcessMemory{Handle: h, Base: base}, nil
}

func (m *ProcessMemory) Close() error {
	return windows.CloseHandle(m.Handle)
}

func (m *ProcessMemory) ReadAt(p []byte, off int64) (int, error) {
	if off < 0 {
		return 0, errors.New("ProcessMemory.ReadAt: negative offset")
	}
	if len(p) == 0 {
		return 0, nil
	}
	var n uintptr
	err := windows.ReadProcessMemory(m.Handle, m.Base+uintptr(off), &p[0], uintptr(len(p)), &n)
	if err != nil {
		return int(n), err
	}
	if int(n) < len(p) {
		return int(n), io.ErrUnexpectedEOF
	}
	return int(n), nil
}

func (m *ProcessMemory) WriteAt(p []byte, off int64) (int, error) {
	if off < 0 {
		return 0, errors.New("ProcessMemory.WriteAt: negative offset")
	}
	if len(p) == 0 {
		return 0, nil
	}
	var n uintptr
	err := windows.WriteProcessMemory(m.Handle, m.Base+uintptr(off), &p[0], uintptr(len(p)), &n)
	return int(n), err
}

// OpenPEFromProcess opens the PE module loaded at base inside pid's address
// space, the live-process analogue of OpenPE. Every PEFile accessor
// (Checksec, ListExports/ListImports, the ROP gadget scanner, PatchBytes)
// works against it exactly as it does against a disk-backed PEFile -- the
// whole point of routing everything through PEFile.r/.w/.RVAToFileOffset.
// Find base with ResolveModuleBase, or read it straight from a leaked
// pointer once you have one.
func OpenPEFromProcess(pid uint32, base uintptr) (*PEFile, error) {
	mem, err := OpenProcessMemory(pid, base)
	if err != nil {
		return nil, err
	}
	f, err := pe.NewFile(mem)
	if err != nil {
		mem.Close()
		return nil, err
	}
	pf := &PEFile{File: f, r: mem, w: mem, closer: mem, live: true}
	if h, herr := pf.header(); herr == nil {
		pf.size = int64(h.sizeOfImage)
	}
	return pf, nil
}

// ResolveModuleBase walks pid's PEB -> Ldr -> InMemoryOrderModuleList to
// find moduleName's load address (e.g. "kernel32.dll", "ntdll.dll", or the
// process's own main-module file name), matched against the loader's
// FullDllName case-insensitively by base name.
//
// This is the live, cross-process twin of
// shellcode/asm/resolver.inc's get_kernel32_base: that NASM walks the exact
// same PEB/Ldr chain from *inside* the target process at IP-control time
// with no Windows API calls available; this does it from *outside*, with a
// debugger/exploit-tooling process's full WinAPI access, which is why it
// can resolve any module by name instead of relying on a fixed loader
// order. It's the real Windows analogue of how pwntools' DynELF defeats
// ASLR by walking ELF structures through a leak oracle -- here the "oracle"
// is ReadProcessMemory itself.
func ResolveModuleBase(pid uint32, moduleName string) (uintptr, error) {
	h, err := windows.OpenProcess(windows.PROCESS_QUERY_INFORMATION|windows.PROCESS_VM_READ, false, pid)
	if err != nil {
		return 0, err
	}
	defer windows.CloseHandle(h)

	var pbi windows.PROCESS_BASIC_INFORMATION
	var retLen uint32
	if err := windows.NtQueryInformationProcess(h, windows.ProcessBasicInformation,
		unsafe.Pointer(&pbi), uint32(unsafe.Sizeof(pbi)), &retLen); err != nil {
		return 0, fmt.Errorf("NtQueryInformationProcess(ProcessBasicInformation): %w", err)
	}
	pebAddr := uintptr(unsafe.Pointer(pbi.PebBaseAddress))
	if pebAddr == 0 {
		return 0, errors.New("PEB address is null")
	}

	peb, err := readRemoteStruct[windows.PEB](h, pebAddr)
	if err != nil {
		return 0, fmt.Errorf("read PEB: %w", err)
	}
	ldrAddr := uintptr(unsafe.Pointer(peb.Ldr))
	if ldrAddr == 0 {
		return 0, errors.New("PEB.Ldr is null -- ntdll's loader (LdrInitializeThunk) hasn't run in this process yet; " +
			"this is normal for a CREATE_SUSPENDED process before it's resumed, retry shortly after ResumeMainThread")
	}

	ldr, err := readRemoteStruct[windows.PEB_LDR_DATA](h, ldrAddr)
	if err != nil {
		return 0, fmt.Errorf("read PEB_LDR_DATA: %w", err)
	}

	// Flink/Blink point at the InMemoryOrderLinks *field* of each entry, not
	// at the start of its LDR_DATA_TABLE_ENTRY -- the field sits at a
	// nonzero offset (0x10 on x64: it's the second of three LIST_ENTRYs at
	// the head of the real struct), so every address read off the list has
	// to be corrected by that offset before it's used as an entry address.
	entryLinksOffset := unsafe.Offsetof(windows.LDR_DATA_TABLE_ENTRY{}.InMemoryOrderLinks)
	headAddr := ldrAddr + unsafe.Offsetof(windows.PEB_LDR_DATA{}.InMemoryOrderModuleList)

	cur := uintptr(unsafe.Pointer(ldr.InMemoryOrderModuleList.Flink))
	for cur != 0 && cur != headAddr {
		entryAddr := cur - entryLinksOffset

		entry, err := readRemoteStruct[windows.LDR_DATA_TABLE_ENTRY](h, entryAddr)
		if err != nil {
			return 0, fmt.Errorf("read LDR_DATA_TABLE_ENTRY: %w", err)
		}

		if name, err := readRemoteUTF16(h, entry.FullDllName); err == nil {
			if strings.EqualFold(moduleBaseName(name), moduleName) {
				return entry.DllBase, nil
			}
		}

		cur = uintptr(unsafe.Pointer(entry.InMemoryOrderLinks.Flink))
	}

	return 0, fmt.Errorf("module %q not found in process %d's loaded module list", moduleName, pid)
}

// readRemoteStruct copies sizeof(T) bytes from h's address space at addr
// into a T, by raw memcpy through ReadProcessMemory -- safe specifically
// because T's fields here are always fixed-width (uintptr/pointer-shaped)
// values being read as bit patterns, never dereferenced as if they were
// local pointers.
func readRemoteStruct[T any](h windows.Handle, addr uintptr) (T, error) {
	var v T
	size := int(unsafe.Sizeof(v))
	buf := make([]byte, size)
	var n uintptr
	if err := windows.ReadProcessMemory(h, addr, &buf[0], uintptr(size), &n); err != nil {
		return v, err
	}
	if int(n) < size {
		return v, io.ErrUnexpectedEOF
	}
	return *(*T)(unsafe.Pointer(&buf[0])), nil
}

func readRemoteUTF16(h windows.Handle, s windows.NTUnicodeString) (string, error) {
	if s.Buffer == nil || s.Length == 0 {
		return "", errors.New("empty NTUnicodeString")
	}
	buf := make([]uint16, s.Length/2)
	var n uintptr
	if err := windows.ReadProcessMemory(h, uintptr(unsafe.Pointer(s.Buffer)),
		(*byte)(unsafe.Pointer(&buf[0])), uintptr(s.Length), &n); err != nil {
		return "", err
	}
	return windows.UTF16ToString(buf), nil
}

func moduleBaseName(path string) string {
	if idx := strings.LastIndexAny(path, `/\`); idx >= 0 {
		path = path[idx+1:]
	}
	return path
}

// SpawnSuspended launches target with CREATE_SUSPENDED, returning the Tube
// (stdin/stdout wired exactly like Spawn) and the new process's PID, with
// the main thread parked before it executes a single instruction.
//
// Verified by testing, worth recording because the obvious assumption is
// wrong: you can NOT resolve module bases via ResolveModuleBase while the
// thread is still suspended. CREATE_SUSPENDED only pins the thread before
// its start routine runs, and that start routine *is*
// ntdll!LdrInitializeThunk -- the loader code that populates
// PEB->Ldr->InMemoryOrderModuleList in the first place. Until it runs,
// PEB.Ldr reads back as null (confirmed empirically: see the smoke test in
// this package's history). What this primitive is actually for is pausing
// *before the loader and entry point run*, e.g. to let Phase 4's debugger
// attach and plant a breakpoint before any application/TLS-callback code
// executes. Call ResumeMainThread(pid), then poll ResolveModuleBase for a
// few milliseconds (the loader runs fast, but it isn't instant) once you
// actually need module bases.
func SpawnSuspended(target string) (tube *Tube, pid uint32, err error) {
	if abs, aerr := filepath.Abs(target); aerr == nil {
		target = abs
	}

	cmd := exec.Command(target)
	cmd.SysProcAttr = &syscall.SysProcAttr{CreationFlags: windows.CREATE_SUSPENDED}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, 0, fmt.Errorf("StdinPipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, 0, fmt.Errorf("StdoutPipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return nil, 0, fmt.Errorf("Start: %w", err)
	}

	return newTube(cmd, nil, stdin, stdout), uint32(cmd.Process.Pid), nil
}

// ResumeMainThread resumes a process started with SpawnSuspended. A
// CREATE_SUSPENDED process has exactly one thread (its initial thread)
// until something resumes it, so finding "the" thread to resume is just
// finding the one thread Toolhelp reports for pid -- there's no race with
// the target spawning more threads first, because it hasn't run yet.
func ResumeMainThread(pid uint32) error {
	snap, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
	if err != nil {
		return fmt.Errorf("CreateToolhelp32Snapshot: %w", err)
	}
	defer windows.CloseHandle(snap)

	var entry windows.ThreadEntry32
	entry.Size = uint32(unsafe.Sizeof(entry))
	for err = windows.Thread32First(snap, &entry); err == nil; err = windows.Thread32Next(snap, &entry) {
		if entry.OwnerProcessID != pid {
			continue
		}
		th, err := windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, entry.ThreadID)
		if err != nil {
			return fmt.Errorf("OpenThread(%d): %w", entry.ThreadID, err)
		}
		_, err = windows.ResumeThread(th)
		windows.CloseHandle(th)
		return err
	}
	return fmt.Errorf("no thread found for process %d", pid)
}