// pe_multitool tours the static-analysis side of winpwn: checksec, section
// entropy, IAT/EAT navigation with forwarder resolution, and the native ROP
// gadget scanner. Run from: workspace/demos/pe_multitool -> go run .
package main

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

const target = "simple_rop.exe";

func main() {
	pe, err := winpwn.OpenPE(target)
	if err != nil {
		log.Fatalf("OpenPE: %v", err)
	}
	defer pe.Close()

	is64, _ := pe.Is64Bit()
	base, _ := pe.ImageBase()
	entry, _ := pe.EntryPoint()
	fmt.Printf("=== %s ===\n", target)
	fmt.Printf("64-bit: %v  ImageBase: 0x%X  EntryPoint: 0x%X\n\n", is64, base, entry)

	// --- checksec ---
	cs, err := pe.Checksec()
	if err != nil {
		log.Fatalf("Checksec: %v", err)
	}
	fmt.Println("--- checksec ---")
	fmt.Printf("ASLR (DYNAMIC_BASE):  %v\n", cs.ASLR)
	fmt.Printf("High-Entropy VA:      %v\n", cs.HighEntropyVA)
	fmt.Printf("DEP (NX_COMPAT):      %v\n", cs.DEP)
	fmt.Printf("CFG:                  %v\n", cs.CFG)
	if cs.SEHApplicable {
		fmt.Printf("SafeSEH:              %v\n", cs.SafeSEH)
	} else {
		fmt.Println("SafeSEH:              n/a (x64 uses table-based SEH)")
	}
	fmt.Printf("GS cookie (heuristic): %v\n", cs.GSHeuristic)
	fmt.Printf("Authenticode present:  %v\n", cs.AuthenticodeSigned)
	fmt.Println()

	// --- section entropy / packing ---
	fmt.Println("--- sections ---")
	for _, sec := range pe.Sections() {
		entropy, _ := sec.Entropy()
		fmt.Printf("%-10s R=%v W=%v X=%v entropy=%.2f\n",
			sec.Name, sec.IsReadable(), sec.IsWritable(), sec.IsExecutable(), entropy)
	}
	fmt.Println()

	// --- IAT: what does this binary already pull in? ---
	fmt.Println("--- interesting imports ---")
	for _, name := range []string{"VirtualProtect", "VirtualAlloc", "LoadLibraryA", "GetProcAddress", "CreateFileA"} {
		imp, err := pe.FindImport(name)
		if err != nil {
			fmt.Printf("%-16s not imported\n", name)
			continue
		}
		fmt.Printf("%-16s %s!%s  IAT RVA=0x%X\n", name, imp.DLL, imp.Name, imp.IATRVA)
	}
	fmt.Println()

	// --- EAT: does this binary export anything (e.g. a win() for ROP)? ---
	exports, err := pe.ListExports()
	if err == nil {
		fmt.Printf("--- exports (%d) ---\n", len(exports))
		for _, e := range exports {
			if e.ForwardTarget != "" {
				dll, fn := winpwn.ParseForwardTarget(e.ForwardTarget)
				fmt.Printf("%s -> forwards to %s!%s\n", e.Name, dll, fn)
			} else {
				fmt.Printf("%s RVA=0x%X\n", e.Name, e.RVA)
			}
		}
		fmt.Println()
	}

	rop, err := winpwn.NewROP(target)
	if err != nil {
		log.Fatalf("NewROP: %v", err)
	}
	defer rop.Close()

	fmt.Println("--- gadgets ---")
	if g, err := rop.Search("pop rcx ; ret"); err == nil {
		fmt.Printf("pop rcx ; ret: 0x%X (%d candidates)\n", g[0].Address, len(g))
		if lines, err := rop.Disassemble(g[0].Address, 2); err == nil {
			fmt.Printf("  verified: %v\n", lines)
		}
	}
	if g, err := rop.SearchRegex(`^pop r\w+ ; pop r\w+ ; ret$`); err == nil {
		fmt.Printf("pop r.. ; pop r.. ; ret: %d candidates, first at 0x%X\n", len(g), g[0].Address)
	}
}