package winpwn

import (
	"debug/pe"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"

	"golang.org/x/arch/x86/x86asm"
)

// Gadget describes a found ROP gadget.
type Gadget struct {
	Address      uint64
	Instructions string
}

// ROP finds ROP gadgets in a binary, the winpwn analogue of pwntools' ROP.
// NewROP shells out to rp-win.exe (a Windows build of rp++) — a real,
// battle-tested gadget finder instead of a hand-rolled scanner. Disassemble
// still works natively (via golang.org/x/arch/x86/x86asm) for verifying a
// chain in-script.
type ROP struct {
	binaryPath string
	toolPath   string
	gadgets    []Gadget
	pe         *PEFile // held open for Disassemble() and arch detection
	mode       int     // x86asm.Mode equivalent (32/64)
}

// defaultRPWinTool is where rp-win.exe lives on this machine. Override with
// the RP_WIN_EXE environment variable if it's installed somewhere else.
const defaultRPWinTool = `C:\tools\rp-win\rp-win.exe`

// NewROP scans binaryPath for ROP/JOP gadgets using rp-win.exe, resolved
// from the RP_WIN_EXE environment variable or defaultRPWinTool. Use
// NewROPExternal to point at a specific tool binary instead (a different
// rp++ build, or a copy kept somewhere else).
func NewROP(binaryPath string) (*ROP, error) {
	toolPath := os.Getenv("RP_WIN_EXE")
	if toolPath == "" {
		toolPath = defaultRPWinTool
	}
	if _, err := os.Stat(toolPath); err != nil {
		return nil, fmt.Errorf("can't find rp-win.exe at %q (%w) -- set RP_WIN_EXE to override, or use NewROPExternal(path, toolPath)", toolPath, err)
	}
	return newROP(binaryPath, toolPath)
}

func newROP(binaryPath, toolPath string) (*ROP, error) {
	peFile, err := OpenPE(binaryPath)
	if err != nil {
		return nil, err
	}

	mode, err := archMode(peFile)
	if err != nil {
		peFile.Close()
		return nil, err
	}

	r := &ROP{binaryPath: binaryPath, toolPath: toolPath, pe: peFile, mode: mode}
	if err := r.findGadgetsExternal(); err != nil {
		peFile.Close()
		return nil, err
	}
	return r, nil
}

// Close releases the PE handle held for Disassemble/arch detection.
func (r *ROP) Close() {
	if r.pe != nil {
		r.pe.Close()
	}
}

func archMode(p *PEFile) (int, error) {
	switch p.File.Machine {
	case pe.IMAGE_FILE_MACHINE_AMD64:
		return 64, nil
	case pe.IMAGE_FILE_MACHINE_I386:
		return 32, nil
	default:
		return 0, fmt.Errorf("unsupported machine type for gadget scanning: 0x%X", p.File.Machine)
	}
}

// Disassemble decodes up to count instructions starting at the absolute
// address addr, the verification step pwntools leaves to objdump/Capstone:
// "did the gadget chain I built actually decode the way I think it did".
func (r *ROP) Disassemble(addr uint64, count int) ([]string, error) {
	imageBase, err := r.pe.ImageBase()
	if err != nil {
		return nil, err
	}
	if addr < imageBase {
		return nil, fmt.Errorf("address is below ImageBase")
	}
	rva := uint32(addr - imageBase)
	offset := r.pe.RVAToFileOffset(rva)
	if offset == 0 {
		return nil, fmt.Errorf("address does not map to any section")
	}

	buf := make([]byte, 16*count)
	n, _ := r.pe.r.ReadAt(buf, offset)
	buf = buf[:n]

	var lines []string
	pos := 0
	for i := 0; i < count && pos < len(buf); i++ {
		inst, err := x86asm.Decode(buf[pos:], r.mode)
		if err != nil {
			return lines, fmt.Errorf("decode failed at +%d: %w", pos, err)
		}
		lines = append(lines, strings.ToLower(x86asm.IntelSyntax(inst, addr+uint64(pos), nil)))
		pos += inst.Len
	}
	return lines, nil
}

// Find returns every gadget whose formatted instruction text contains
// pattern as a substring, ranked so index 0 is the best candidate to use --
// the one-expression version of Search for the common case:
//
//	rop.Find("pop rcx ; ret")[0].Address
//
// Ranking, in order: an exact match to pattern beats a mere substring match
// (a bare "pop rcx ; ret" outranks "ror byte [rax-0x1], 0x15 ; pop rcx ;
// ret" even though both contain the pattern and the latter may sit at a
// numerically lower address); shorter instruction text (fewer side-effect
// instructions riding along) beats longer; address ascending breaks
// remaining ties for determinism. Without this, sorting by raw address
// alone can hand back a "dirty" multi-instruction gadget at [0] purely
// because it happens to start a few bytes earlier in memory.
//
// Indexing an empty result panics -- deliberate for exploit scripts: fail
// loudly at the gadget lookup itself, not three chain-steps later against a
// garbage address.
func (r *ROP) Find(pattern string) []Gadget {
	needle := strings.ToLower(pattern)
	var out []Gadget
	for _, g := range r.gadgets {
		if strings.Contains(strings.ToLower(g.Instructions), needle) {
			out = append(out, g)
		}
	}
	sort.Slice(out, func(i, j int) bool {
		ei, ej := strings.EqualFold(out[i].Instructions, pattern), strings.EqualFold(out[j].Instructions, pattern)
		if ei != ej {
			return ei
		}
		if len(out[i].Instructions) != len(out[j].Instructions) {
			return len(out[i].Instructions) < len(out[j].Instructions)
		}
		return out[i].Address < out[j].Address
	})
	return out
}

// SearchRegex finds gadgets whose instruction text matches the given
// regular expression (e.g. `^pop r[a-z]+ ; ret$`), for when a plain
// substring (Search/Find) isn't precise enough.
func (r *ROP) SearchRegex(pattern string) ([]Gadget, error) {
	re, err := regexp.Compile(pattern)
	if err != nil {
		return nil, err
	}

	var results []Gadget
	for _, g := range r.gadgets {
		if re.MatchString(strings.ToLower(g.Instructions)) {
			results = append(results, g)
		}
	}
	if len(results) == 0 {
		return nil, fmt.Errorf("no gadget matched regex %q", pattern)
	}
	return results, nil
}