package winpwn

import (
	"errors"
)

// Import describes one entry of a PE's Import Address Table (IAT), the
// analogue of pwntools reading a binary's .got/.dynsym to find which libc
// functions it pulls in.
type Import struct {
	DLL string
	// Name is empty when the function is imported by ordinal only.
	Name    string
	Ordinal uint16
	// IATRVA is the RVA of this import's slot in the IAT -- the address the
	// loader overwrites with the real function pointer at load time, and
	// the address you'd target if you wanted to overwrite the import to
	// redirect a call.
	IATRVA uint32
}

// ImportedLib describes one DLL a PE imports, plus that DLL's current live
// image base in this process (see (*PEFile).ImportedLibs).
type ImportedLib struct {
	Name string
	// Base is 0 if Err is set (the library failed to load in this process).
	Base uint64
	Err  error
}

const importOrdinalFlag64 = uint64(1) << 63
const importOrdinalFlag32 = uint32(1) << 31

// importDescriptor mirrors winnt.h's IMAGE_IMPORT_DESCRIPTOR.
type importDescriptor struct {
	OriginalFirstThunk uint32
	TimeDateStamp      uint32
	ForwarderChain     uint32
	Name               uint32
	FirstThunk         uint32
}

// ListImports walks the full Import Directory Table (IDT) and each DLL's
// thunk array, resolving every imported name/ordinal and its IAT slot
// address. The analogue of pwntools poking at a binary's dynamic symbol
// table to see what it links against.
func (p *PEFile) ListImports() ([]Import, error) {
	h, err := p.header()
	if err != nil {
		return nil, err
	}

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

	var imports []Import
	for i := 0; ; i++ {
		var desc importDescriptor
		if err := p.readStructAt(descOffset+int64(i*20), &desc); err != nil {
			return nil, err
		}
		if desc.OriginalFirstThunk == 0 && desc.Name == 0 && desc.FirstThunk == 0 {
			break // null terminator descriptor
		}

		dllName, err := p.readCString(p.RVAToFileOffset(desc.Name))
		if err != nil {
			return nil, err
		}

		thunkRVA := desc.OriginalFirstThunk
		if thunkRVA == 0 {
			thunkRVA = desc.FirstThunk // some linkers omit the ILT entirely
		}

		entries, err := p.walkThunks(h.is64, thunkRVA, desc.FirstThunk)
		if err != nil {
			return nil, err
		}
		for i := range entries {
			entries[i].DLL = dllName
		}
		imports = append(imports, entries...)
	}

	return imports, nil
}

func (p *PEFile) walkThunks(is64 bool, thunkRVA uint32, iatRVA uint32) ([]Import, error) {
	var out []Import
	thunkSize := uint32(4)
	if is64 {
		thunkSize = 8
	}

	for j := uint32(0); ; j++ {
		thunkOffset := p.RVAToFileOffset(thunkRVA + j*thunkSize)

		var imp Import
		imp.IATRVA = iatRVA + j*thunkSize

		if is64 {
			var thunk uint64
			if err := p.readStructAt(thunkOffset, &thunk); err != nil {
				return nil, err
			}
			if thunk == 0 {
				break
			}
			if thunk&importOrdinalFlag64 != 0 {
				imp.Ordinal = uint16(thunk & 0xFFFF)
			} else {
				name, err := p.readCString(p.RVAToFileOffset(uint32(thunk)) + 2) // skip Hint WORD
				if err != nil {
					return nil, err
				}
				imp.Name = name
			}
		} else {
			var thunk uint32
			if err := p.readStructAt(thunkOffset, &thunk); err != nil {
				return nil, err
			}
			if thunk == 0 {
				break
			}
			if thunk&importOrdinalFlag32 != 0 {
				imp.Ordinal = uint16(thunk & 0xFFFF)
			} else {
				name, err := p.readCString(p.RVAToFileOffset(thunk) + 2)
				if err != nil {
					return nil, err
				}
				imp.Name = name
			}
		}

		out = append(out, imp)
	}

	return out, nil
}

// FindImport looks for a specific imported function by name across every
// imported DLL -- the quick "does this binary already pull in
// VirtualProtect/LoadLibraryA/GetProcAddress" check.
func (p *PEFile) FindImport(funcName string) (*Import, error) {
	imports, err := p.ListImports()
	if err != nil {
		return nil, err
	}
	for i := range imports {
		if imports[i].Name == funcName {
			return &imports[i], nil
		}
	}
	return nil, errors.New("import not found: " + funcName)
}