#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <Winternl.h>
#include <assert.h>
#include <Tlhelp32.h>

#pragma pack(push, 1)
struct far_jmp
{
    BYTE PushOp;
    PVOID PushArg;
    BYTE RetOp;
};

struct OldCode
{
    DWORD One;
    WORD Two;
};
#pragma pack(pop)

void StopThreads()
{
    DWORD currTh;
    HANDLE thrHandle;
    HANDLE h;
    DWORD currPr;
    THREADENTRY32 Thread;

    currTh = GetCurrentThreadId();
    currPr = GetCurrentProcessId();
    h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if(h != INVALID_HANDLE_VALUE) {
        Thread.dwSize = sizeof(THREADENTRY32);
        if(Thread32First(h, &Thread)) {
            do {
                if(Thread.th32ThreadID != currTh && Thread.th32OwnerProcessID == currPr) {
                    thrHandle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, Thread.th32ThreadID);
                    if(thrHandle > 0) {
                        SuspendThread(thrHandle);
                        CloseHandle(thrHandle);
                    }
                }
            }
            while(!Thread32Next(h, &Thread));
        }
        CloseHandle(h);
    }
}

void RunThreads()
{
    DWORD currTh;
    HANDLE thrHandle;
    HANDLE h;
    DWORD currPr;
    THREADENTRY32 Thread;

    currTh = GetCurrentThreadId();
    currPr = GetCurrentProcessId();
    h = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if(h != INVALID_HANDLE_VALUE) {
        Thread.dwSize = sizeof(THREADENTRY32);
        if(Thread32First(h, &Thread)) {
            do {
                if(Thread.th32ThreadID != currTh && Thread.th32OwnerProcessID == currPr) {
                    thrHandle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, Thread.th32ThreadID);
                    if(thrHandle > 0) {
                        ResumeThread(thrHandle);
                        CloseHandle(thrHandle);
                    }
                }
            }
            while(!Thread32Next(h, &Thread));
        }
        CloseHandle(h);
    }
}

HANDLE RegQueryCurrProc;
PVOID AdrRegQuery;
OldCode OldRegQuery;
far_jmp JmpRegQuery;

typedef NTSTATUS (WINAPI *NewOpenKeyFun)(PHANDLE KeyHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES ObjectAttributes);

// My new function
NTSTATUS WINAPI NOpenKey(PHANDLE KeyHandle, ACCESS_MASK DesiredAccess, POBJECT_ATTRIBUTES ObjectAttributes)
{
    FILE* outfile = fopen("syscalls.log", "a");


    buffer = (char*)malloc(ObjectAttributes->ObjectName->Length + 1);
    wctomb(buffer, *ObjectAttributes->ObjectName->Buffer);
    buffer[ObjectAttributes->ObjectName->Length] = '\0';
    fprintf(outfile, "%s\n", buffer);
    fclose(outfile);

    WriteProcessMemory(GetCurrentProcess(), AdrRegQuery, &OldRegQuery, sizeof(OldCode), &Written);

    //NewOpenKeyFun ZwOpenKeyAddress =      (NewOpenKeyFun)GetProcAddress(GetModuleHandle(L"Ntdll.dll"), "ZwOpenKey");
    NewOpenKeyFun ZwOpenKeyAddress = (NewOpenKeyFun)AdrRegQuery;

    NTSTATUS result = (*ZwOpenKeyAddress)(KeyHandle, DesiredAccess, ObjectAttributes);


    WriteProcessMemory(GetCurrentProcess(), AdrRegQuery, &JmpRegQuery, sizeof(far_jmp), &Written);
    return result;
}

void SetRegQueryHook()
{
    DWORD Written;
    AdrRegQuery = GetProcAddress(GetModuleHandle(L"Ntdll.dll"), "ZwOpenKey");
    JmpRegQuery.PushOp = 0x68;
        JmpRegQuery.PushArg = NOpenKey;
        JmpRegQuery.RetOp = 0xC3;
    ReadProcessMemory(GetCurrentProcess(), AdrRegQuery, &OldRegQuery, sizeof(OldCode), &Written);
    WriteProcessMemory(GetCurrentProcess(), AdrRegQuery, &JmpRegQuery, sizeof(far_jmp), &Written);
}

BOOL APIENTRY DllMain( HANDLE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
                     )
{
    if( (ul_reason_for_call == DLL_PROCESS_ATTACH)) {
        StopThreads();
        SetRegQueryHook();
        RunThreads();
    }
    return TRUE;
}