package main

import (
	"encoding/binary"
	"fmt"
	"os"
	"strings"
	"unsafe"

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

const compareBytes = 16

// --- PE Structure Definitions (x64) ---

type IMAGE_DOS_HEADER struct {
	EMagic  uint16
	Pad     [28]uint16
	ELfanew int32
}

type IMAGE_DATA_DIRECTORY struct {
	VirtualAddress uint32
	Size           uint32
}

type IMAGE_FILE_HEADER struct {
	Machine              uint16
	NumberOfSections     uint16
	TimeDateStamp        uint32
	PointerToSymbolTable uint32
	NumberOfSymbols      uint32
	SizeOfOptionalHeader uint16
	Characteristics      uint16
}

type IMAGE_OPTIONAL_HEADER64 struct {
	Magic                       uint16
	MajorLinkerVersion          uint8
	MinorLinkerVersion          uint8
	SizeOfCode                  uint32
	SizeOfInitializedData       uint32
	SizeOfUninitializedData     uint32
	AddressOfEntryPoint         uint32
	BaseOfCode                  uint32
	ImageBase                   uint64
	SectionAlignment            uint32
	FileAlignment               uint32
	MajorOperatingSystemVersion uint16
	MinorOperatingSystemVersion uint16
	MajorImageVersion           uint16
	MinorImageVersion           uint16
	MajorSubsystemVersion       uint16
	MinorSubsystemVersion       uint16
	Win32VersionValue           uint32
	SizeOfImage                 uint32
	SizeOfHeaders               uint32
	CheckSum                    uint32
	Subsystem                   uint16
	DllCharacteristics          uint16
	SizeOfStackReserve          uint64
	SizeOfStackCommit           uint64
	SizeOfHeapReserve           uint64
	SizeOfHeapCommit            uint64
	LoaderFlags                 uint32
	NumberOfRvaAndSizes         uint32
	DataDirectory               [16]IMAGE_DATA_DIRECTORY
}

type IMAGE_NT_HEADERS64 struct {
	Signature      uint32
	FileHeader     IMAGE_FILE_HEADER
	OptionalHeader IMAGE_OPTIONAL_HEADER64
}

type IMAGE_EXPORT_DIRECTORY struct {
	Characteristics       uint32
	TimeDateStamp         uint32
	MajorVersion          uint16
	MinorVersion          uint16
	Name                  uint32
	Base                  uint32
	NumberOfFunctions     uint32
	NumberOfNames         uint32
	AddressOfFunctions    uint32
	AddressOfNames        uint32
	AddressOfNameOrdinals uint32
}

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

// --- Hook Detection Logic ---

type HookResult struct {
	FunctionName string
	HookType     string
	MemoryBytes  []byte
	DiskBytes    []byte
	TargetAddr   uintptr
	TargetModule string
}

func classifyHook(mem, disk []byte) string {
	if len(mem) < 6 || len(disk) < 6 {
		return ""
	}

	match := true
	n := len(mem)
	if len(disk) < n {
		n = len(disk)
	}
	for i := 0; i < n; i++ {
		if mem[i] != disk[i] {
			match = false
			break
		}
	}
	if match {
		return ""
	}

	switch {
	case mem[0] == 0xE9:
		return "JMP rel32 (E9)"
	case mem[0] == 0xCC:
		return "INT3 Breakpoint"
	case mem[0] == 0xFF && mem[1] == 0x25:
		return "JMP [RIP+disp32]"
	case mem[0] == 0x48 && mem[1] == 0xB8:
		return "MOV RAX + JMP RAX"
	case mem[0] == 0x68 && len(mem) > 5 && mem[5] == 0xC7:
		return "PUSH/MOV/RET"
	default:
		return "Unknown Patch"
	}
}

func resolveJmpRel32(hookAddr uintptr, memBytes []byte) uintptr {
	if len(memBytes) < 5 || memBytes[0] != 0xE9 {
		return 0
	}
	offset := int32(binary.LittleEndian.Uint32(memBytes[1:5]))
	return hookAddr + 5 + uintptr(offset)
}

var (
	modKernel32            = windows.NewLazySystemDLL("kernel32.dll")
	procGetModuleHandleExA = modKernel32.NewProc("GetModuleHandleExA")
	procGetModuleFileNameA = modKernel32.NewProc("GetModuleFileNameA")
)

func getModuleFromAddress(addr uintptr) string {
	const GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS = 0x00000004
	const GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT = 0x00000002

	var hMod uintptr
	ret, _, _ := procGetModuleHandleExA.Call(
		GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS|
			GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
		addr,
		uintptr(unsafe.Pointer(&hMod)),
	)
	if ret == 0 || hMod == 0 {
		return "(unbacked memory)"
	}

	buf := make([]byte, 260)
	n, _, _ := procGetModuleFileNameA.Call(
		hMod,
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(len(buf)),
	)
	if n == 0 {
		return "(unknown module)"
	}
	for i := 0; i < len(buf); i++ {
		if buf[i] == 0 {
			return string(buf[:i])
		}
	}
	return string(buf[:n])
}

func bytesToHex(b []byte) string {
	parts := make([]string, len(b))
	for i, v := range b {
		parts[i] = fmt.Sprintf("%02X", v)
	}
	return strings.Join(parts, " ")
}

// rvaToFileOffset converts an RVA to a raw file offset using section headers.
// This is critical: on disk, data lives at PointerToRawData offsets, not at
// the VirtualAddress offsets used after the loader maps the PE into memory.
func rvaToFileOffset(rva uint32, sections []IMAGE_SECTION_HEADER) (uint32, bool) {
	for _, s := range sections {
		secEnd := s.VirtualAddress + s.SizeOfRawData
		if s.SizeOfRawData == 0 {
			secEnd = s.VirtualAddress + s.VirtualSize
		}
		if rva >= s.VirtualAddress && rva < secEnd {
			return rva - s.VirtualAddress + s.PointerToRawData, true
		}
	}
	return 0, false
}

func readCString(data []byte, offset uint32) string {
	if int(offset) >= len(data) {
		return ""
	}
	end := offset
	for int(end) < len(data) && data[end] != 0 {
		end++
	}
	return string(data[offset:end])
}

func scanModule(moduleName, diskPath string) ([]HookResult, int, error) {
	// Get loaded module base address
	modNameUTF16, _ := windows.UTF16PtrFromString(moduleName)
	var hModule windows.Handle
	err := windows.GetModuleHandleEx(0, modNameUTF16, &hModule)
	if err != nil {
		return nil, 0, fmt.Errorf("GetModuleHandleEx(%s): %w", moduleName, err)
	}
	moduleBase := uintptr(hModule)

	// Read clean copy from disk
	diskData, err := os.ReadFile(diskPath)
	if err != nil {
		return nil, 0, fmt.Errorf("reading %s: %w", diskPath, err)
	}

	if len(diskData) < int(unsafe.Sizeof(IMAGE_DOS_HEADER{})) {
		return nil, 0, fmt.Errorf("file too small for DOS header")
	}

	// Parse DOS header
	dosHdr := (*IMAGE_DOS_HEADER)(unsafe.Pointer(&diskData[0]))
	if dosHdr.EMagic != 0x5A4D {
		return nil, 0, fmt.Errorf("invalid DOS signature")
	}

	if int(dosHdr.ELfanew)+int(unsafe.Sizeof(IMAGE_NT_HEADERS64{})) > len(diskData) {
		return nil, 0, fmt.Errorf("e_lfanew out of bounds")
	}

	// Parse NT headers
	ntHdr := (*IMAGE_NT_HEADERS64)(unsafe.Pointer(&diskData[dosHdr.ELfanew]))
	if ntHdr.Signature != 0x00004550 {
		return nil, 0, fmt.Errorf("invalid PE signature")
	}

	// Parse section headers (right after the optional header)
	sectionOffset := int(dosHdr.ELfanew) + 4 +
		int(unsafe.Sizeof(IMAGE_FILE_HEADER{})) +
		int(ntHdr.FileHeader.SizeOfOptionalHeader)
	numSections := int(ntHdr.FileHeader.NumberOfSections)
	sections := make([]IMAGE_SECTION_HEADER, numSections)
	for i := 0; i < numSections; i++ {
		off := sectionOffset + i*int(unsafe.Sizeof(IMAGE_SECTION_HEADER{}))
		if off+int(unsafe.Sizeof(IMAGE_SECTION_HEADER{})) > len(diskData) {
			break
		}
		sections[i] = *(*IMAGE_SECTION_HEADER)(unsafe.Pointer(&diskData[off]))
	}

	// Locate export directory
	exportDirRVA := ntHdr.OptionalHeader.DataDirectory[0].VirtualAddress
	exportDirSize := ntHdr.OptionalHeader.DataDirectory[0].Size
	if exportDirRVA == 0 {
		return nil, 0, fmt.Errorf("no export directory")
	}

	exportDirFileOff, ok := rvaToFileOffset(exportDirRVA, sections)
	if !ok {
		return nil, 0, fmt.Errorf("cannot resolve export dir RVA")
	}

	if int(exportDirFileOff)+int(unsafe.Sizeof(IMAGE_EXPORT_DIRECTORY{})) > len(diskData) {
		return nil, 0, fmt.Errorf("export directory out of bounds")
	}

	exportDir := (*IMAGE_EXPORT_DIRECTORY)(unsafe.Pointer(&diskData[exportDirFileOff]))

	// Resolve name/ordinal/function arrays via RVA-to-file-offset translation
	namesOff, ok := rvaToFileOffset(exportDir.AddressOfNames, sections)
	if !ok {
		return nil, 0, fmt.Errorf("cannot resolve AddressOfNames RVA")
	}
	ordinalsOff, ok := rvaToFileOffset(exportDir.AddressOfNameOrdinals, sections)
	if !ok {
		return nil, 0, fmt.Errorf("cannot resolve AddressOfNameOrdinals RVA")
	}
	functionsOff, ok := rvaToFileOffset(exportDir.AddressOfFunctions, sections)
	if !ok {
		return nil, 0, fmt.Errorf("cannot resolve AddressOfFunctions RVA")
	}

	// Build executable section RVA ranges to filter out data exports
	// (e.g. NlsAnsiCodePage, LdrSystemDllInitBlock live in .data/.rdata)
	type rvaRange struct{ start, end uint32 }
	var execRanges []rvaRange
	for _, s := range sections {
		const IMAGE_SCN_MEM_EXECUTE = 0x20000000
		if s.Characteristics&IMAGE_SCN_MEM_EXECUTE != 0 {
			execRanges = append(execRanges,
				rvaRange{s.VirtualAddress, s.VirtualAddress + s.VirtualSize})
		}
	}
	isExecutableRVA := func(rva uint32) bool {
		for _, r := range execRanges {
			if rva >= r.start && rva < r.end {
				return true
			}
		}
		return false
	}

	totalExports := int(exportDir.NumberOfNames)
	var results []HookResult
	scannedCode := 0

	for i := 0; i < totalExports; i++ {
		// Read name RVA
		nameRVAOff := int(namesOff) + i*4
		if nameRVAOff+4 > len(diskData) {
			continue
		}
		nameRVA := binary.LittleEndian.Uint32(diskData[nameRVAOff : nameRVAOff+4])

		nameFileOff, ok := rvaToFileOffset(nameRVA, sections)
		if !ok {
			continue
		}
		funcName := readCString(diskData, nameFileOff)
		if funcName == "" {
			continue
		}

		// Read ordinal
		ordOff := int(ordinalsOff) + i*2
		if ordOff+2 > len(diskData) {
			continue
		}
		ordinal := binary.LittleEndian.Uint16(diskData[ordOff : ordOff+2])

		// Read function RVA
		funcRVAOff := int(functionsOff) + int(ordinal)*4
		if funcRVAOff+4 > len(diskData) {
			continue
		}
		funcRVA := binary.LittleEndian.Uint32(diskData[funcRVAOff : funcRVAOff+4])

		// Skip forwarded exports (RVA within export directory)
		if funcRVA >= exportDirRVA && funcRVA < exportDirRVA+exportDirSize {
			continue
		}

		// Skip data exports: only scan functions in executable sections
		if !isExecutableRVA(funcRVA) {
			continue
		}

		scannedCode++

		// Resolve file offset for the function body
		funcFileOff, ok := rvaToFileOffset(funcRVA, sections)
		if !ok {
			continue
		}
		if int(funcFileOff)+compareBytes > len(diskData) {
			continue
		}

		// Read disk bytes
		diskBytes := make([]byte, compareBytes)
		copy(diskBytes, diskData[funcFileOff:funcFileOff+compareBytes])

		// Read memory bytes (module base + RVA = in-memory address)
		memAddr := moduleBase + uintptr(funcRVA)
		memBytes := make([]byte, compareBytes)

		// Safe memory read using ReadProcessMemory on self
		var bytesRead uintptr
		currentProcess, _ := windows.GetCurrentProcess()
		err := windows.ReadProcessMemory(
			currentProcess, memAddr,
			&memBytes[0], uintptr(compareBytes), &bytesRead,
		)
		if err != nil || bytesRead < compareBytes {
			// Fallback: direct pointer read (may work if memory is readable)
			func() {
				defer func() { recover() }()
				src := unsafe.Slice((*byte)(unsafe.Pointer(memAddr)), compareBytes)
				copy(memBytes, src)
			}()
		}

		hookType := classifyHook(memBytes, diskBytes)
		if hookType == "" {
			continue
		}

		result := HookResult{
			FunctionName: funcName,
			HookType:     hookType,
			MemoryBytes:  memBytes,
			DiskBytes:    diskBytes,
		}

		if memBytes[0] == 0xE9 {
			result.TargetAddr = resolveJmpRel32(memAddr, memBytes)
			if result.TargetAddr != 0 {
				result.TargetModule = getModuleFromAddress(result.TargetAddr)
			}
			// Filter out intra-module JMPs (internal control flow, not EDR hooks)
			var modInfo windows.ModuleInfo
			err := windows.GetModuleInformation(
				currentProcess, hModule, &modInfo,
				uint32(unsafe.Sizeof(modInfo)),
			)
			if err == nil {
				modStart := uintptr(modInfo.BaseOfDll)
				modEnd := modStart + uintptr(modInfo.SizeOfImage)
				if result.TargetAddr >= modStart && result.TargetAddr < modEnd {
					continue // JMP lands inside the same DLL - not a hook
				}
			}
		}

		results = append(results, result)
	}

	return results, scannedCode, nil
}

func main() {
	fmt.Println("=== EDR User-Mode Hook Scanner (Go) ===")
	fmt.Println("Compares loaded modules against clean disk baselines")
	fmt.Println()

	targets := []struct {
		module string
		path   string
	}{
		{"ntdll.dll", `C:\Windows\System32\ntdll.dll`},
		{"kernel32.dll", `C:\Windows\System32\kernel32.dll`},
		{"kernelbase.dll", `C:\Windows\System32\kernelbase.dll`},
	}

	totalHooks := 0

	for _, t := range targets {
		results, numExports, err := scanModule(t.module, t.path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "[-] Error scanning %s: %v\n", t.module, err)
			continue
		}

		fmt.Printf("\n[+] %s: %d hooks detected (%d code exports scanned)\n",
			t.module, len(results), numExports)
		fmt.Println(strings.Repeat("=", 60))

		for _, r := range results {
			fmt.Printf("[!] HOOKED: %s\n", r.FunctionName)
			fmt.Printf("    Type:   %s\n", r.HookType)
			fmt.Printf("    Memory: %s\n", bytesToHex(r.MemoryBytes))
			fmt.Printf("    Disk:   %s\n", bytesToHex(r.DiskBytes))
			if r.TargetAddr != 0 {
				fmt.Printf("    Target: 0x%X -> %s\n", r.TargetAddr, r.TargetModule)
			}
			fmt.Println()
		}

		totalHooks += len(results)
	}

	fmt.Printf("\n[+] SUMMARY: %d total hooks detected across all scanned modules\n",
		totalHooks)
	if totalHooks == 0 {
		fmt.Println("[*] No hooks found. Either no EDR is active, " +
			"or it uses kernel-only telemetry.")
	}
}