package winpwn

import (
	"bytes"
	"debug/pe"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"
)

// PEFile is a read-only-by-default handle on a PE32/PE32+ image, the winpwn
// analogue of pwntools' ELF. Unlike ELF, a PE's bytes can come from two
// meaningfully different places: a file on disk (OpenPE) or a loaded
// module's live address space inside a running process (OpenPEFromProcess).
// Every accessor in this package (checksec, IAT/EAT, gadget scanning,
// patching) is written against the r/w fields below so both backings get
// every feature for free.
type PEFile struct {
	File *pe.File

	r      io.ReaderAt
	w      io.WriterAt // nil when opened read-only
	closer io.Closer   // nil if there is nothing to close
	size   int64       // total backing size, for whole-image operations (checksum recompute)

	// live is true when r/w address a process's memory (OpenPEFromProcess)
	// rather than a file's bytes (OpenPE/OpenPEForWrite). See
	// RVAToFileOffset for why this changes the RVA translation.
	live bool
}

func OpenPE(path string) (*PEFile, error) {
	fd, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	f, err := pe.NewFile(fd)
	if err != nil {
		fd.Close()
		return nil, err
	}
	info, err := fd.Stat()
	if err != nil {
		fd.Close()
		return nil, err
	}
	return &PEFile{File: f, r: fd, closer: fd, size: info.Size()}, nil
}

func (p *PEFile) Close() {
	if p.closer != nil {
		p.closer.Close()
	}
}

// IsLive reports whether this PEFile is backed by a running process's
// address space (OpenPEFromProcess) rather than a file on disk.
func (p *PEFile) IsLive() bool {
	return p.live
}

// peHeader normalizes the PE32 (32-bit) vs PE32+ (64-bit) optional header
// split into the fields callers actually need, so the rest of the package
// doesn't have to type-switch on pe.OptionalHeader32/64 everywhere.
type peHeader struct {
	is64                bool
	imageBase           uint64
	addressOfEntryPoint uint32
	sizeOfImage         uint32
	dllCharacteristics  uint16
	dataDirectory       [16]pe.DataDirectory
}

func (p *PEFile) header() (peHeader, error) {
	switch oh := p.File.OptionalHeader.(type) {
	case *pe.OptionalHeader64:
		return peHeader{
			is64:                true,
			imageBase:           oh.ImageBase,
			addressOfEntryPoint: oh.AddressOfEntryPoint,
			sizeOfImage:         oh.SizeOfImage,
			dllCharacteristics:  oh.DllCharacteristics,
			dataDirectory:       oh.DataDirectory,
		}, nil
	case *pe.OptionalHeader32:
		return peHeader{
			is64:                false,
			imageBase:           uint64(oh.ImageBase),
			addressOfEntryPoint: oh.AddressOfEntryPoint,
			sizeOfImage:         oh.SizeOfImage,
			dllCharacteristics:  oh.DllCharacteristics,
			dataDirectory:       oh.DataDirectory,
		}, nil
	default:
		return peHeader{}, errors.New("unrecognized PE optional header (not PE32 or PE32+)")
	}
}

// Is64Bit reports whether this is a PE32+ (x64) image.
func (p *PEFile) Is64Bit() (bool, error) {
	h, err := p.header()
	if err != nil {
		return false, err
	}
	return h.is64, nil
}

// ImageBase returns the preferred load address from the PE optional header
// (the Go analogue of pwntools' ELF.address when ASLR is disabled). For a
// live-process-backed PEFile this is still the header's *preferred* base,
// not necessarily where the module actually landed -- use the base passed
// to OpenPEFromProcess (e.g. from ResolveModuleBase) for the real address.
func (p *PEFile) ImageBase() (uint64, error) {
	h, err := p.header()
	if err != nil {
		return 0, err
	}
	return h.imageBase, nil
}

// EntryPoint returns the absolute address of the entry point
// (ImageBase + AddressOfEntryPoint).
func (p *PEFile) EntryPoint() (uint64, error) {
	h, err := p.header()
	if err != nil {
		return 0, err
	}
	return h.imageBase + uint64(h.addressOfEntryPoint), nil
}

// RVAToFileOffset converts a relative virtual address into the offset to
// pass to this PEFile's backing ReaderAt/WriterAt.
//
// For a disk-backed PEFile (OpenPE/OpenPEForWrite) this walks the section
// table to translate an RVA into a PointerToRawData-relative file offset --
// necessary because SectionAlignment and FileAlignment differ, so a
// section's position in the loaded image and its position on disk aren't
// the same number.
//
// For a live-process-backed PEFile (OpenPEFromProcess) the backing
// ReaderAt/WriterAt already treats offset 0 as the module's base address,
// so "offset" already *is* the RVA: that's the entire definition of a
// relative virtual address once the image is actually loaded. This is the
// identity function in that case.
func (p *PEFile) RVAToFileOffset(rva uint32) int64 {
	if p.live {
		return int64(rva)
	}
	for _, sec := range p.File.Sections {
		if rva >= sec.VirtualAddress && rva < sec.VirtualAddress+sec.VirtualSize {
			return int64(rva - sec.VirtualAddress + sec.Offset)
		}
	}
	return 0
}

// readCString reads a null-terminated ASCII string starting at the given
// read offset (file offset, or RVA for a live-backed PEFile -- see
// RVAToFileOffset).
func (p *PEFile) readCString(offset int64) (string, error) {
	var out []byte
	buf := make([]byte, 1)
	for {
		if _, err := p.r.ReadAt(buf, offset); err != nil {
			if err == io.EOF {
				break
			}
			return "", err
		}
		if buf[0] == 0 {
			break
		}
		out = append(out, buf[0])
		offset++
	}
	return string(out), nil
}

// readStructAt fills v (a pointer to a fixed-size struct of fixed-width
// fields) by reading binary.Size(v) bytes at offset through this PEFile's
// backing ReaderAt. The one read-offset helper every struct-shaped PE
// directory parse in this package goes through, so disk and live-process
// backings share the exact same parsing code.
func (p *PEFile) readStructAt(offset int64, v any) error {
	size := binary.Size(v)
	if size < 0 {
		return errors.New("readStructAt: unsupported type")
	}
	buf := make([]byte, size)
	if _, err := p.r.ReadAt(buf, offset); err != nil {
		return err
	}
	return binary.Read(bytes.NewReader(buf), binary.LittleEndian, v)
}

// writeAt writes data at offset through this PEFile's backing WriterAt,
// failing clearly if the PEFile was opened read-only.
func (p *PEFile) writeAt(offset int64, data []byte) error {
	if p.w == nil {
		return errors.New("PEFile is read-only; open with OpenPEForWrite or OpenPEFromProcess(write) for write access")
	}
	_, err := p.w.WriteAt(data, offset)
	return err
}

// GetProcAddress looks up a function's RVA by name, the bare-bones analogue
// of the real WinAPI call of the same name. It's a thin wrapper over the
// canonical export-table walk in GetExport (exports.go); use GetExport
// directly when you need forwarder resolution or ordinal information too.
func (p *PEFile) GetProcAddress(funcName string) (uint64, error) {
	exp, err := p.GetExport(funcName)
	if err != nil {
		return 0, err
	}
	if exp.RVA == 0 {
		return 0, fmt.Errorf("export %q is a forwarder (%s.%s), not a local RVA -- resolve it in the target DLL instead",
			funcName, exp.ForwardTarget, funcName)
	}
	return uint64(exp.RVA), nil
}