package winpwn

import (
	"debug/pe"
	"encoding/binary"
	"errors"
	"os"
	"unsafe"
)

// OpenPEForWrite opens a PE file read-write, for the patching methods below.
// Plain OpenPE is read-only by design; writing is opt-in so a script can't
// accidentally corrupt a target binary it only meant to inspect.
func OpenPEForWrite(path string) (*PEFile, error) {
	fd, err := os.OpenFile(path, os.O_RDWR, 0)
	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, w: fd, closer: fd, size: info.Size()}, nil
}

// PatchBytes overwrites the file's contents at the given RVA with data, the
// general-purpose "patch on the fly" primitive.
func (p *PEFile) PatchBytes(rva uint32, data []byte) error {
	offset := p.RVAToFileOffset(rva)
	if offset == 0 {
		return errors.New("RVA does not map to any section")
	}
	return p.PatchBytesAtOffset(offset, data)
}

// PatchBytesAtOffset overwrites the file's contents at a raw file offset.
func (p *PEFile) PatchBytesAtOffset(offset int64, data []byte) error {
	return p.writeAt(offset, data)
}

// peHeaderOffset reads e_lfanew (at the fixed DOS-header offset 0x3C) to
// find where the "PE\0\0" header begins.
func (p *PEFile) peHeaderOffset() (int64, error) {
	var lfanew uint32
	if err := p.readStructAt(0x3C, &lfanew); err != nil {
		return 0, err
	}
	return int64(lfanew), nil
}

// coffHeaderOffset returns the file offset of the COFF File Header, right
// after the 4-byte "PE\0\0" signature.
func (p *PEFile) coffHeaderOffset() (int64, error) {
	peOffset, err := p.peHeaderOffset()
	if err != nil {
		return 0, err
	}
	return peOffset + 4, nil
}

// optionalHeaderOffset returns the file offset of the Optional Header,
// right after the fixed 20-byte COFF File Header.
func (p *PEFile) optionalHeaderOffset() (int64, error) {
	coffOffset, err := p.coffHeaderOffset()
	if err != nil {
		return 0, err
	}
	return coffOffset + 20, nil
}

// sectionHeaderTableOffset returns the file offset of the first
// IMAGE_SECTION_HEADER entry, right after the Optional Header.
func (p *PEFile) sectionHeaderTableOffset() (int64, error) {
	optOffset, err := p.optionalHeaderOffset()
	if err != nil {
		return 0, err
	}
	return optOffset + int64(p.File.FileHeader.SizeOfOptionalHeader), nil
}

// imageSectionHeaderSize and the byte offset of the Characteristics field
// within it (IMAGE_SECTION_HEADER: Name[8] + 6 DWORDs + 2 WORDs + Characteristics DWORD).
const (
	imageSectionHeaderSize               = 40
	imageSectionHeaderCharacteristicsOff = 36
)

// SetSectionCharacteristics overwrites a section's Characteristics flags
// directly in the section header — e.g. to flip on IMAGE_SCN_MEM_EXECUTE for
// a section you want to use as shellcode landing space. Requires a PEFile
// opened with OpenPEForWrite.
func (p *PEFile) SetSectionCharacteristics(name string, characteristics uint32) error {
	tableOffset, err := p.sectionHeaderTableOffset()
	if err != nil {
		return err
	}
	for i, sec := range p.File.Sections {
		if sec.Name != name {
			continue
		}
		headerOffset := tableOffset + int64(i)*imageSectionHeaderSize
		var buf [4]byte
		binary.LittleEndian.PutUint32(buf[:], characteristics)
		return p.PatchBytesAtOffset(headerOffset+imageSectionHeaderCharacteristicsOff, buf[:])
	}
	return errors.New("section not found: " + name)
}

// MakeSectionExecutable ORs in IMAGE_SCN_MEM_EXECUTE on top of a section's
// existing characteristics (e.g. "make .data executable" for a quick and
// dirty shellcode-in-data-section trick).
func (p *PEFile) MakeSectionExecutable(name string) error {
	return p.orSectionCharacteristics(name, imageSCNMemExecute)
}

// MakeSectionWritable ORs in IMAGE_SCN_MEM_WRITE on top of a section's
// existing characteristics.
func (p *PEFile) MakeSectionWritable(name string) error {
	return p.orSectionCharacteristics(name, imageSCNMemWrite)
}

func (p *PEFile) orSectionCharacteristics(name string, flag uint32) error {
	for _, sec := range p.File.Sections {
		if sec.Name == name {
			return p.SetSectionCharacteristics(name, sec.Characteristics|flag)
		}
	}
	return errors.New("section not found: " + name)
}

// imageTLSDirectory mirrors winnt.h's IMAGE_TLS_DIRECTORY32/64: same field
// order in both, only pointer-sized members change width. All fields here
// are absolute VAs, not RVAs — the one PE directory that isn't RVA-based.
type imageTLSDirectory64 struct {
	StartAddressOfRawData uint64
	EndAddressOfRawData   uint64
	AddressOfIndex        uint64
	AddressOfCallBacks    uint64
}

type imageTLSDirectory32 struct {
	StartAddressOfRawData uint32
	EndAddressOfRawData   uint32
	AddressOfIndex        uint32
	AddressOfCallBacks    uint32
}

// DisableTLSCallbacks zeroes the AddressOfCallBacks field of the TLS
// Directory, so the loader never walks (and never invokes) the callback
// array at all — the one-field patch that defeats TLS-callback-based
// anti-debug/anti-instrumentation tricks that fire before your entry point
// or your debugger's first breakpoint gets a chance to run.
func (p *PEFile) DisableTLSCallbacks() error {
	h, err := p.header()
	if err != nil {
		return err
	}

	const dirEntryTLS = 9
	dir := h.dataDirectory[dirEntryTLS]
	if dir.VirtualAddress == 0 {
		return errors.New("no TLS directory present")
	}
	offset := p.RVAToFileOffset(dir.VirtualAddress)
	if offset == 0 {
		return errors.New("failed to map TLS directory RVA to file offset")
	}

	if h.is64 {
		zeros := make([]byte, 8)
		return p.PatchBytesAtOffset(offset+int64(unsafe.Offsetof(imageTLSDirectory64{}.AddressOfCallBacks)), zeros)
	}
	zeros := make([]byte, 4)
	return p.PatchBytesAtOffset(offset+int64(unsafe.Offsetof(imageTLSDirectory32{}.AddressOfCallBacks)), zeros)
}

// RecalculateChecksum recomputes and writes the Optional Header's PE
// checksum (the algorithm behind imagehlp's CheckSumMappedFile/MapFileAndCheckSum),
// so a binary you've patched on disk still passes loaders/AV/signing tools
// that validate it.
func (p *PEFile) RecalculateChecksum() error {
	h, err := p.header()
	if err != nil {
		return err
	}
	optOffset, err := p.optionalHeaderOffset()
	if err != nil {
		return err
	}
	checksumOffset := optOffset
	if h.is64 {
		checksumOffset += int64(unsafe.Offsetof(pe.OptionalHeader64{}.CheckSum))
	} else {
		checksumOffset += int64(unsafe.Offsetof(pe.OptionalHeader32{}.CheckSum))
	}

	data := make([]byte, p.size)
	if _, err := p.r.ReadAt(data, 0); err != nil {
		return err
	}

	checksum := peChecksum(data, checksumOffset)

	var buf [4]byte
	binary.LittleEndian.PutUint32(buf[:], checksum)
	return p.PatchBytesAtOffset(checksumOffset, buf[:])
}

// peChecksum implements the PE checksum algorithm: sum the file as 16-bit
// little-endian words (treating the existing 4-byte checksum field as if it
// contributed zero), fold carries back into the low 16 bits, then add the
// file size.
func peChecksum(data []byte, checksumFieldOffset int64) uint32 {
	var checksum uint32

	n := len(data)
	for i := 0; i < n; i += 2 {
		if int64(i) == checksumFieldOffset || int64(i) == checksumFieldOffset+2 {
			continue // skip the checksum field's own two words
		}
		var word uint32
		if i+1 < n {
			word = uint32(data[i]) | uint32(data[i+1])<<8
		} else {
			word = uint32(data[i]) // trailing odd byte
		}
		checksum = (checksum & 0xFFFF) + word + (checksum >> 16)
	}
	checksum = (checksum & 0xFFFF) + (checksum >> 16)
	checksum += uint32(n)
	return checksum
}