#include <windows.h>
#include <stdio.h>

#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64) || defined(_M_AMD64)
#define _CPU_X64
#elif defined(__i386__) || defined(_M_IX86)
#define _CPU_X86
#else
#error "Unsupported CPU"
#endif

#ifdef _CPU_X64
#define HOOK_JUMP_LEN 14
#endif
#ifdef _CPU_X86
#define HOOK_JUMP_LEN 5
#endif

WINBOOL WINAPI fk_WriteConsoleA(HANDLE, CONST VOID *, DWORD, LPDWORD, LPVOID);

HANDLE hstdout = NULL;
void *hook_func = NULL;
char hook_jump[HOOK_JUMP_LEN];
char old_entry[HOOK_JUMP_LEN];

void inithook()
{
    HMODULE hmodule = GetModuleHandleA("kernelbase.dll");
    hook_func = (void *)GetProcAddress(hmodule, "WriteConsoleA");
    VirtualProtect(hook_func, HOOK_JUMP_LEN, PAGE_EXECUTE_READWRITE, NULL);
#ifdef _CPU_X64
    union
    {
        void *ptr;
        struct
        {
            long lo;
            long hi;
        };
    } ptr64;
    ptr64.ptr = (void *)fk_WriteConsoleA;
    hook_jump[0] = 0x68; // push xxx
    *(long *)&hook_jump[1] = ptr64.lo; // xxx，即地址的低4位
    hook_jump[5] = 0xC7;
    hook_jump[6] = 0x44;
    hook_jump[7] = 0x24;
    hook_jump[8] = 0x04; // mov dword [rsp+4], yyy
    *(long *)&hook_jump[9] = ptr64.hi; // yyy，即地址的高4位
    hook_jump[13] = 0xC3; // ret
    ReadProcessMemory(GetCurrentProcess(), hook_func, old_entry, HOOK_JUMP_LEN, NULL);
#endif
#ifdef _CPU_X86
    hook_jump[0] = 0xE9;
    *(long *)&hook_jump[1] = (BYTE *)fk_WriteConsoleA - (BYTE *)hook_func - 5;
    ReadProcessMemory(GetCurrentProcess(), hook_func, old_entry, HOOK_JUMP_LEN, NULL);
#endif
}

void dohook()
{
    WriteProcessMemory(GetCurrentProcess(), hook_func, hook_jump, HOOK_JUMP_LEN, NULL);
}

void unhook()
{
    WriteProcessMemory(GetCurrentProcess(), hook_func, old_entry, HOOK_JUMP_LEN, NULL);
}

WINBOOL WINAPI fk_WriteConsoleA(HANDLE hConsoleOutput, CONST VOID *lpBuffer, DWORD nNumberOfCharsToWrite, LPDWORD lpNumberOfCharsWritten, LPVOID lpReserved)
{
    unhook();
    char buf[128];
    strcpy(buf, (char *)lpBuffer);
    buf[nNumberOfCharsToWrite - 1] = '\0';
    strcat(buf, "\t[hook]\n");
    int len = nNumberOfCharsToWrite + 8;
    WINBOOL result = WriteConsoleA(hConsoleOutput, buf, len, NULL, NULL);
    dohook();
    return result;
}

DWORD WINAPI thread_writehello(void *stdh)
{
    DWORD id = GetCurrentThreadId();
    char str[64];
    for (int i = 0; i < 10; i++) {
        int len = sprintf(str, "%d: Hello World %d\n", id, i);
        WriteConsoleA(stdh, str, len, NULL, NULL);
    }
    return 0;
}

#define THREAD_COUNT 5

int main()
{
    inithook();
    dohook();
    hstdout = GetStdHandle(-11);
    HANDLE hthreads[THREAD_COUNT];
    for (int i = 0; i < THREAD_COUNT; i++)
        hthreads[i] = CreateThread(NULL, 0, thread_writehello, hstdout, CREATE_SUSPENDED, NULL);
    for (int i = 0; i < THREAD_COUNT; i++)
        ResumeThread(hthreads[i]);
    for (int i = 0; i < THREAD_COUNT; i++)
        WaitForSingleObject(hthreads[i], 1000);
    for (int i = 0; i < THREAD_COUNT; i++)
        CloseHandle(hthreads[i]);
    WriteConsoleA(hstdout, "Must hook\n", 10, NULL, NULL);
    unhook();
    WriteConsoleA(hstdout, "Not hook\n", 9, NULL, NULL);
}