package manalocator

import (
	"errors"
	"fmt"
	"unsafe"

	"github.com/almounah/superdeye/internal/utils/helper"
	"github.com/almounah/superdeye/internal/utils/superdwindows"
)

var RANGE int = 200

/*
 * ssn in the ssn number
 * cleanSyscall is the Syscall address found in NTDLL
 **/
type SuperdSyscallTool struct {
	Ssn                       uint32
	SyscallInstructionAddress superdwindows.PVOID
}

func LookupSSNAndTrampoline(syscallName string, hModule superdwindows.HANDLE) (superdSyscallTool SuperdSyscallTool, err error) {
	pBase := unsafe.Pointer(hModule)
	pImgExportDir, err := helper.GetImageExportDirectory(hModule)
    if err != nil {
        fmt.Println("Messed Up Getting Image Export Directory")
        return SuperdSyscallTool{}, err
    }
	numFunction := pImgExportDir.NumberOfFunctions

	AddressOfFuntionArray := unsafe.Slice((*superdwindows.DWORD)(unsafe.Pointer(uintptr(pBase)+uintptr(pImgExportDir.AddressOfFunctions))), pImgExportDir.NumberOfFunctions)
	AddressOfNamesArray := unsafe.Slice((*superdwindows.DWORD)(unsafe.Pointer(uintptr(pBase)+uintptr(pImgExportDir.AddressOfNames))), pImgExportDir.NumberOfFunctions)
	AddressOfNameOrdinalArray := unsafe.Slice((*superdwindows.WORD)(unsafe.Pointer(uintptr(pBase)+uintptr(pImgExportDir.AddressOfNameOrdinals))), pImgExportDir.NumberOfFunctions)

	for num := superdwindows.DWORD(0); num < numFunction; num++ {
		functionNameRVA := AddressOfNamesArray[num]
		s := helper.NameRvaToString(uintptr(pBase), functionNameRVA)
		if s == syscallName {
			functionAddress := uintptr(pBase) + uintptr(AddressOfFuntionArray[AddressOfNameOrdinalArray[num]])

			if checkIfCleanSSN(functionAddress) {
				low := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 4)))
				high := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 5)))
				ssn := uint32((high << 8) | low)
				syscallAddress, _ := findSyscallAddress(functionAddress)
				res := SuperdSyscallTool{ssn, syscallAddress}
				return res, nil
			}

			// Search the neighbors if SSN is hooked
			for neighborIndex := 1; neighborIndex < RANGE; neighborIndex++ {
				upNeighborFunctionAddress := uintptr(unsafe.Add(unsafe.Pointer(functionAddress), neighborIndex*32))
				if checkIfCleanSSN(upNeighborFunctionAddress) {
					low := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(upNeighborFunctionAddress), 4)))
					high := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(upNeighborFunctionAddress), 5)))
					ssn := uint32((high<<8)|low) - uint32(neighborIndex)
					syscallAddress, _ := findSyscallAddress(upNeighborFunctionAddress)
					res := SuperdSyscallTool{ssn, syscallAddress}
					return res, nil
				}
				downNeighborFunctionAddress := uintptr(unsafe.Add(unsafe.Pointer(functionAddress), -neighborIndex*32))
				if checkIfCleanSSN(downNeighborFunctionAddress) {
					low := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(downNeighborFunctionAddress), 4)))
					high := uint32(*(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(downNeighborFunctionAddress), 5)))
					ssn := uint32((high<<8)|low) + uint32(neighborIndex)
					syscallAddress, _ := findSyscallAddress(downNeighborFunctionAddress)
					res := SuperdSyscallTool{ssn, syscallAddress}
					return res, nil
				}

			}

		}
	}
	return SuperdSyscallTool{}, errors.New("Did not find SSN. Double Check Syscall Name then report as a Github issue as this is very interesting.")
}

func checkIfCleanSSN(functionAddress uintptr) bool {
	val0 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 0))
	val1 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 1))
	val2 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 2))
	val3 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 3))
	val6 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 6))
	val7 := *(*superdwindows.BYTE)(unsafe.Add(unsafe.Pointer(functionAddress), 7))

	if val0 == 0x4c && val1 == 0x8b && val2 == 0xd1 && val3 == 0xb8 && val6 == 0x00 && val7 == 0x00 {
		return true
	}

	return false
}

func findSyscallAddress(functionAddress uintptr) (pAddress superdwindows.PVOID, err error) {
	for z := 1; z < RANGE; z++ {
		pAddress := unsafe.Add(unsafe.Pointer(functionAddress), z)
		pAddressNext := unsafe.Add(unsafe.Pointer(functionAddress), z+1)

		value := *(*superdwindows.BYTE)(pAddress)
		valueNext := *(*superdwindows.BYTE)(pAddressNext)

		if value == 0x0F && valueNext == 0x05 {
			return superdwindows.PVOID(pAddress), nil
		}
	}
	return 0, errors.New("Not Found Syscall instruction in Ntdll")
}
