package winpwn

import (
	"errors"
	"strings"
)

// Export describes one entry of a PE's export table (EAT), the analogue of
// pwntools' libc.symbols[...] entries.
type Export struct {
	// Name is empty when the function is exported by ordinal only.
	Name    string
	Ordinal uint16
	// RVA is the function's address. It is zero when ForwardTarget is set:
	// the export doesn't point at code in this module at all, it forwards
	// to a function in another DLL.
	RVA uint32
	// ForwardTarget is "DLLNAME.FuncName" when this export is a forwarder
	// (e.g. api-ms-win-core-*.dll entries that forward into kernelbase.dll).
	// Use ParseForwardTarget to split it.
	ForwardTarget string
}

// ParseForwardTarget splits a forwarder string ("KERNELBASE.CreateFileW")
// into the target DLL and function name.
func ParseForwardTarget(forward string) (dll string, fn string) {
	idx := strings.LastIndex(forward, ".")
	if idx == -1 {
		return "", forward
	}
	return forward[:idx], forward[idx+1:]
}

// exportDirectory mirrors winnt.h's IMAGE_EXPORT_DIRECTORY.
type exportDirectory struct {
	Characteristics       uint32
	TimeDateStamp         uint32
	MajorVersion          uint16
	MinorVersion          uint16
	Name                  uint32
	Base                  uint32
	NumberOfFunctions     uint32
	NumberOfNames         uint32
	AddressOfFunctions    uint32
	AddressOfNames        uint32
	AddressOfNameOrdinals uint32
}

// ListExports walks the full Export Address Table, the analogue of
// pwntools' libc.symbols when you need every entry rather than a single
// lookup. Functions exported by ordinal only (no name) come back with
// Name == "". Works against a live-process-backed PEFile exactly as well
// as a disk-backed one (see RVAToFileOffset).
func (p *PEFile) ListExports() ([]Export, error) {
	h, err := p.header()
	if err != nil {
		return nil, err
	}

	dir := h.dataDirectory[0]
	if dir.VirtualAddress == 0 {
		return nil, errors.New("export table not found")
	}
	exportOffset := p.RVAToFileOffset(dir.VirtualAddress)
	if exportOffset == 0 {
		return nil, errors.New("failed to map export directory RVA to file offset")
	}

	var expDir exportDirectory
	if err := p.readStructAt(exportOffset, &expDir); err != nil {
		return nil, err
	}

	funcsOffset := p.RVAToFileOffset(expDir.AddressOfFunctions)
	namesOffset := p.RVAToFileOffset(expDir.AddressOfNames)
	ordinalsOffset := p.RVAToFileOffset(expDir.AddressOfNameOrdinals)

	// Build ordinal-index -> name from the name table before walking
	// AddressOfFunctions, since not every function slot has a name.
	nameByOrdinalIndex := make(map[uint16]string, expDir.NumberOfNames)
	for i := uint32(0); i < expDir.NumberOfNames; i++ {
		var nameRVA uint32
		if err := p.readStructAt(namesOffset+int64(i*4), &nameRVA); err != nil {
			return nil, err
		}
		name, err := p.readCString(p.RVAToFileOffset(nameRVA))
		if err != nil {
			return nil, err
		}

		var ordinalIndex uint16
		if err := p.readStructAt(ordinalsOffset+int64(i*2), &ordinalIndex); err != nil {
			return nil, err
		}
		nameByOrdinalIndex[ordinalIndex] = name
	}

	// A forwarder's "RVA" doesn't point at code: it points back inside the
	// export directory itself, at an ASCII "DLL.Func" string.
	forwarderLo := dir.VirtualAddress
	forwarderHi := dir.VirtualAddress + dir.Size

	exports := make([]Export, 0, expDir.NumberOfFunctions)
	for i := uint32(0); i < expDir.NumberOfFunctions; i++ {
		var rva uint32
		if err := p.readStructAt(funcsOffset+int64(i*4), &rva); err != nil {
			return nil, err
		}
		if rva == 0 {
			continue // unused ordinal slot
		}

		e := Export{
			Name:    nameByOrdinalIndex[uint16(i)],
			Ordinal: uint16(expDir.Base + i),
		}
		if rva >= forwarderLo && rva < forwarderHi {
			fwd, err := p.readCString(p.RVAToFileOffset(rva))
			if err != nil {
				return nil, err
			}
			e.ForwardTarget = fwd
		} else {
			e.RVA = rva
		}
		exports = append(exports, e)
	}

	return exports, nil
}

// GetExport looks up a single export by name and reports whether it
// forwards to another DLL, unlike GetProcAddress which returns a bare RVA.
func (p *PEFile) GetExport(name string) (*Export, error) {
	exports, err := p.ListExports()
	if err != nil {
		return nil, err
	}
	for i := range exports {
		if exports[i].Name == name {
			return &exports[i], nil
		}
	}
	return nil, errors.New("function not found in export table: " + name)
}