// Command winpwn is a thin CLI wrapper around the winpwn library, the
// analogue of pwntools' `pwn` command -- for the quick "just tell me the
// answer" cases (checksec, an offset, a gadget search) where spinning up a
// whole solve script is overkill. The library (package winpwn, used the
// way pwntools itself is: `import "winpwn"` in a real solve script) is
// still the primary interface; this is additive, not a replacement.
package main

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"

	"winpwn"
)

func main() {
	if len(os.Args) < 2 {
		usage()
		os.Exit(1)
	}

	cmd := os.Args[1]
	args := os.Args[2:]

	var err error
	switch cmd {
	case "checksec":
		err = cmdChecksec(args)
	case "cyclic":
		err = cmdCyclic(args)
	case "hex":
		err = cmdHex(args)
	case "unhex":
		err = cmdUnhex(args)
	case "hexdump":
		err = cmdHexdump(args)
	case "rop":
		err = cmdRop(args)
	case "bytes":
		err = cmdBytes(args)
	case "disasm":
		err = cmdDisasm(args)
	case "exports":
		err = cmdExports(args)
	case "imports":
		err = cmdImports(args)
	case "heap":
		err = cmdHeap(args)
	case "help", "-h", "--help":
		usage()
		return
	default:
		fmt.Fprintf(os.Stderr, "winpwn: unknown subcommand %q\n\n", cmd)
		usage()
		os.Exit(1)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "winpwn: %v\n", err)
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprint(os.Stderr, `winpwn -- quick-answer CLI for the winpwn library (PE/Windows pwn toolkit)

Usage:
  winpwn checksec <file>
  winpwn cyclic <length> [-n N]              generate a de Bruijn pattern
  winpwn cyclic -l <subseq> [-n N]           find subseq's offset (subseq may be "0x..." or literal bytes)
  winpwn hex                                 read raw bytes from stdin, print hex
  winpwn unhex                               read hex from stdin, print raw bytes
  winpwn hexdump <file>
  winpwn rop <file> -search "pop rcx ; ret"
  winpwn rop <file> -regex "^pop r.* ; ret$"
  winpwn bytes <file> <hex>                  find every VA where <hex> occurs (e.g. ebfe for jmp $)
  winpwn disasm <file> <hexaddr> <count>
  winpwn exports <file>
  winpwn imports <file>
  winpwn heap <pid>                          dump all heaps in a live process
  winpwn heap <pid> -walk                    also walk all NT Heap entries (slow on large heaps)
`)
}

func cmdChecksec(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: winpwn checksec <file>")
	}
	pf, err := winpwn.OpenPE(args[0])
	if err != nil {
		return err
	}
	defer pf.Close()

	r, err := pf.Checksec()
	if err != nil {
		return err
	}

	arch := "x86"
	if r.Is64Bit {
		arch = "x64"
	}
	fmt.Printf("Arch:           %s\n", arch)
	fmt.Printf("ASLR:           %s\n", yesNo(r.ASLR))
	fmt.Printf("HighEntropyVA:  %s\n", yesNo(r.HighEntropyVA))
	fmt.Printf("DEP/NX:         %s\n", yesNo(r.DEP))
	fmt.Printf("CFG:            %s\n", yesNo(r.CFG))
	if r.SEHApplicable {
		fmt.Printf("SafeSEH:        %s\n", yesNo(r.SafeSEH))
	} else {
		fmt.Printf("SafeSEH:        N/A (x64 uses table-based SEH)\n")
	}
	fmt.Printf("GS (heuristic): %s\n", yesNo(r.GSHeuristic))
	fmt.Printf("Authenticode:   %s\n", yesNo(r.AuthenticodeSigned))
	fmt.Printf(".NET (CLR):     %s\n", yesNo(r.DotNET))

	fmt.Println("\n--- sections ---")
	for _, sec := range pf.Sections() {
		perm := ""
		if sec.IsReadable() {
			perm += "R"
		} else {
			perm += "-"
		}
		if sec.IsWritable() {
			perm += "W"
		} else {
			perm += "-"
		}
		if sec.IsExecutable() {
			perm += "X"
		} else {
			perm += "-"
		}
		entropy, _ := sec.Entropy()
		fmt.Printf("  %-8s %s  VA=0x%-8X Offset=0x%-8X Size=0x%-8X entropy=%.2f\n",
			sec.Name, perm, sec.VirtualAddress, sec.Offset, sec.VirtualSize, entropy)
	}

	libs, err := pf.ImportedLibs()
	if err != nil {
		fmt.Printf("\n(imported libs unavailable: %v)\n", err)
		return nil
	}
	fmt.Println("\n--- imported libs (live image base -- stable until next reboot) ---")
	for _, lib := range libs {
		if lib.Err != nil {
			fmt.Printf("  %-24s (failed to load: %v)\n", lib.Name, lib.Err)
			continue
		}
		fmt.Printf("  %-24s 0x%016X\n", lib.Name, lib.Base)
	}
	return nil
}

func yesNo(b bool) string {
	if b {
		return "Yes"
	}
	return "No"
}

func cmdCyclic(args []string) error {
	n := 4
	var find string
	var rest []string

	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "-n":
			i++
			if i >= len(args) {
				return errors.New("-n requires a value")
			}
			v, err := strconv.Atoi(args[i])
			if err != nil {
				return fmt.Errorf("invalid -n value: %w", err)
			}
			n = v
		case "-l":
			i++
			if i >= len(args) {
				return errors.New("-l requires a value")
			}
			find = args[i]
		default:
			rest = append(rest, args[i])
		}
	}

	if find != "" {
		off := winpwn.CyclicFindN(parseSubseq(find, n), n)
		if off == -1 {
			return fmt.Errorf("subsequence %q not found in the n=%d cyclic pattern", find, n)
		}
		fmt.Println(off)
		return nil
	}

	if len(rest) != 1 {
		return errors.New("usage: winpwn cyclic <length> [-n N]  |  winpwn cyclic -l <subseq> [-n N]")
	}
	length, err := strconv.Atoi(rest[0])
	if err != nil {
		return fmt.Errorf("invalid length: %w", err)
	}
	fmt.Println(string(winpwn.CyclicN(length, n)))
	return nil
}

// parseSubseq accepts either a literal byte string ("aaab") or a hex-packed
// integer ("0x62616161", as read back from a crashed register/return
// address) and packs the latter little-endian at width n -- the CLI
// equivalent of pwntools' cyclic_find() accepting either bytes or an int.
func parseSubseq(s string, n int) []byte {
	if v, ok := strings.CutPrefix(s, "0x"); ok {
		if u, err := strconv.ParseUint(v, 16, 64); err == nil {
			if n == 8 {
				return winpwn.P64(u)
			}
			return winpwn.P32(uint32(u))
		}
	}
	return []byte(s)
}

func cmdHex(args []string) error {
	data, err := io.ReadAll(os.Stdin)
	if err != nil {
		return err
	}
	fmt.Println(winpwn.Enhex(data))
	return nil
}

func cmdUnhex(args []string) error {
	data, err := io.ReadAll(os.Stdin)
	if err != nil {
		return err
	}
	raw, err := winpwn.Unhex(strings.TrimSpace(string(data)))
	if err != nil {
		return err
	}
	_, err = os.Stdout.Write(raw)
	return err
}

func cmdHexdump(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: winpwn hexdump <file>")
	}
	data, err := os.ReadFile(args[0])
	if err != nil {
		return err
	}
	fmt.Print(winpwn.Hexdump(data))
	return nil
}

func cmdRop(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: winpwn rop <file> -search PATTERN | -regex PATTERN")
	}
	target := args[0]
	var search, pattern string
	for i := 1; i < len(args); i++ {
		switch args[i] {
		case "-search":
			i++
			if i >= len(args) {
				return errors.New("-search requires a value")
			}
			search = args[i]
		case "-regex":
			i++
			if i >= len(args) {
				return errors.New("-regex requires a value")
			}
			pattern = args[i]
		}
	}
	if search == "" && pattern == "" {
		return errors.New("specify -search or -regex (a full unfiltered gadget dump isn't supported from the CLI -- it can be tens of thousands of entries; use the library's NewROP+r.Search from a script instead)")
	}

	rop, err := winpwn.NewROP(target)
	if err != nil {
		return err
	}
	defer rop.Close()

	var gadgets []winpwn.Gadget
	if search != "" {
		gadgets, err = rop.Search(search)
	} else {
		gadgets, err = rop.SearchRegex(pattern)
	}
	if err != nil {
		return err
	}

	for _, g := range gadgets {
		fmt.Printf("0x%016X: %s\n", g.Address, g.Instructions)
	}
	return nil
}

func cmdBytes(args []string) error {
	if len(args) != 2 {
		return errors.New("usage: winpwn bytes <file> <hex>  (e.g. winpwn bytes kernel32.dll ebfe for jmp $)")
	}
	pattern, err := winpwn.Unhex(normalizeHex(args[1]))
	if err != nil {
		return fmt.Errorf("invalid hex pattern %q: %w", args[1], err)
	}

	pf, err := winpwn.OpenPE(args[0])
	if err != nil {
		return err
	}
	defer pf.Close()

	base, err := pf.ImageBase()
	if err != nil {
		return err
	}

	rvas, err := pf.SearchBytes(pattern)
	if err != nil {
		return err
	}
	for _, rva := range rvas {
		fmt.Printf("0x%016X  (RVA 0x%X)\n", base+rva, rva)
	}
	return nil
}

// normalizeHex strips the separators tools commonly paste hex bytes with
// ("\xeb\xfe", "EB FE", "eb-fe", "eb:fe") down to a bare hex.DecodeString-
// compatible string, so winpwn bytes accepts whatever got copied out of
// x64dbg/rp++/a disassembler without the user reformatting it by hand.
func normalizeHex(s string) string {
	replacer := strings.NewReplacer("\\x", "", " ", "", "-", "", ":", "", ",", "")
	return replacer.Replace(s)
}

func cmdDisasm(args []string) error {
	if len(args) != 3 {
		return errors.New("usage: winpwn disasm <file> <hexaddr> <count>")
	}
	target := args[0]
	addr, err := strconv.ParseUint(strings.TrimPrefix(args[1], "0x"), 16, 64)
	if err != nil {
		return fmt.Errorf("invalid address: %w", err)
	}
	count, err := strconv.Atoi(args[2])
	if err != nil {
		return fmt.Errorf("invalid count: %w", err)
	}

	rop, err := winpwn.NewROP(target)
	if err != nil {
		return err
	}
	defer rop.Close()

	lines, err := rop.Disassemble(addr, count)
	if err != nil {
		return err
	}
	for _, l := range lines {
		fmt.Println(l)
	}
	return nil
}

func cmdExports(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: winpwn exports <file>")
	}
	pf, err := winpwn.OpenPE(args[0])
	if err != nil {
		return err
	}
	defer pf.Close()

	exports, err := pf.ListExports()
	if err != nil {
		return err
	}
	base, err := pf.ImageBase()
	if err != nil {
		return err
	}

	for _, e := range exports {
		name := e.Name
		if name == "" {
			name = "(no name)"
		}
		if e.ForwardTarget != "" {
			fmt.Printf("%-40s ordinal=%-5d -> %s\n", name, e.Ordinal, e.ForwardTarget)
		} else {
			fmt.Printf("%-40s ordinal=%-5d 0x%016X\n", name, e.Ordinal, base+uint64(e.RVA))
		}
	}
	return nil
}

func cmdImports(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: winpwn imports <file>")
	}
	pf, err := winpwn.OpenPE(args[0])
	if err != nil {
		return err
	}
	defer pf.Close()

	imports, err := pf.ListImports()
	if err != nil {
		return err
	}

	for _, im := range imports {
		name := im.Name
		if name == "" {
			name = fmt.Sprintf("ordinal#%d", im.Ordinal)
		}
		fmt.Printf("%-20s %-40s IAT=0x%08X\n", im.DLL, name, im.IATRVA)
	}
	return nil
}

func cmdHeap(args []string) error {
	if len(args) < 1 {
		return errors.New("usage: winpwn heap <pid> [-walk]")
	}
	pidU, err := strconv.ParseUint(args[0], 10, 32)
	if err != nil {
		return fmt.Errorf("invalid pid %q: %w", args[0], err)
	}
	pid := uint32(pidU)

	walkEntries := len(args) >= 2 && args[1] == "-walk"

	heaps, err := winpwn.ListProcessHeaps(pid)
	if err != nil {
		return fmt.Errorf("ListProcessHeaps: %w", err)
	}
	fmt.Printf("pid %d: %d heap(s)\n", pid, len(heaps))

	mem, err := winpwn.OpenProcessMemory(pid, 0)
	if err != nil {
		return fmt.Errorf("OpenProcessMemory: %w", err)
	}
	defer mem.Close()

	for i, haddr := range heaps {
		kind, err := winpwn.DetectHeapKind(mem, haddr)
		if err != nil {
			fmt.Printf("\n[%d] 0x%016x  error: %v\n", i, haddr, err)
			continue
		}
		fmt.Printf("\n[%d] 0x%016x  %s\n", i, haddr, kind)

		switch kind {
		case winpwn.HeapKindNT:
			printNTHeap(mem, haddr, walkEntries)
		case winpwn.HeapKindSegment:
			printSegmentHeap(mem, haddr)
		}
	}
	return nil
}

func printNTHeap(r *winpwn.ProcessMemory, addr uint64, walkEntries bool) {
	h, err := winpwn.ReadHeap(r, addr)
	if err != nil {
		fmt.Printf("  ReadHeap error: %v\n", err)
		return
	}
	feType := "none"
	if h.FrontEndHeapType == winpwn.FrontEndHeapLFH {
		feType = fmt.Sprintf("LFH @ 0x%x", h.FrontEndHeap)
	} else if h.FrontEndHeapType == winpwn.FrontEndHeapLookaside {
		feType = "Lookaside"
	}
	fmt.Printf("  flags=0x%08x  encoding=%v  front-end=%s\n",
		h.Flags, h.EncodingActive(), feType)

	segs, _ := h.Segments(r)
	fmt.Printf("  segments: %d\n", len(segs)+1) // +1 for segment0 always present

	if !walkEntries {
		return
	}

	entries, err := h.WalkAllHeapEntries(r)
	if err != nil {
		fmt.Printf("  WalkAllHeapEntries: %v\n", err)
	}
	stats := winpwn.SummariseEntries(entries)
	fmt.Printf("  entries: total=%d busy=%d free=%d  busy_bytes=%d free_bytes=%d\n",
		stats.TotalEntries, stats.BusyEntries, stats.FreeEntries,
		stats.BusyBytes, stats.FreeBytes)

	pairs := winpwn.AdjacentBusyPairs(entries)
	if len(pairs) > 0 {
		fmt.Printf("  adjacent busy pairs: %d\n", len(pairs))
		shown := pairs
		if len(shown) > 5 {
			shown = shown[:5]
		}
		for _, p := range shown {
			fmt.Printf("    0x%x (size %d) <-> 0x%x (size %d)\n",
				p[0].UserData(), p[0].UserSize(),
				p[1].UserData(), p[1].UserSize())
		}
	}

	if h.FrontEndHeapType == winpwn.FrontEndHeapLFH {
		buckets, err := winpwn.ReadLFHBuckets(r, h.FrontEndHeap)
		if err != nil {
			fmt.Printf("  ReadLFHBuckets: %v\n", err)
			return
		}
		fmt.Printf("  LFH active buckets:")
		n := 0
		for _, b := range buckets {
			if b.BlockUnits == 0 {
				continue
			}
			fmt.Printf(" [%d]=%db", b.Index, b.BlockSize())
			n++
			if n >= 8 {
				fmt.Printf(" ...")
				break
			}
		}
		fmt.Println()
	}
}

func printSegmentHeap(r *winpwn.ProcessMemory, addr uint64) {
	h, err := winpwn.ReadSegmentHeap(r, addr)
	if err != nil {
		fmt.Printf("  ReadSegmentHeap error: %v\n", err)
		return
	}
	fmt.Printf("  GlobalFlags=0x%08x\n", h.GlobalFlags)
	fmt.Printf("  VS context @ 0x%x: committed=%d free=%d subsegments=%d\n",
		h.VS.Addr, h.VS.CommittedUnits, h.VS.FreeUnits, h.VS.SubsegmentCount)
	if len(h.LFH.ActiveBuckets) > 0 {
		fmt.Printf("  LFH active buckets (total-blocks):")
		for j, b := range h.LFH.ActiveBuckets {
			fmt.Printf(" [%d]=%d", b.Index, b.TotalBlockCount)
			if j >= 7 {
				fmt.Printf(" ...")
				break
			}
		}
		fmt.Println()
	} else {
		fmt.Printf("  LFH: no active buckets\n")
	}
}