#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

#define HOOK_JUMP_LEN 5

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

#ifdef _CPU_X64
static void *FindModuleTextBlankAlign(HMODULE hmodule)
{
    BYTE *p = (BYTE *)hmodule;
    p += ((IMAGE_DOS_HEADER *)p)->e_lfanew + 4; // 根据DOS头获取PE信息偏移量
    p += sizeof(IMAGE_FILE_HEADER) + ((IMAGE_FILE_HEADER *)p)->SizeOfOptionalHeader; // 跳过可选头
    WORD sections = ((IMAGE_FILE_HEADER *)p)->NumberOfSections; // 获取区段长度
    for (int i = 0; i < sections; i++) {
        IMAGE_SECTION_HEADER *psec = (IMAGE_SECTION_HEADER *)p;
        p += sizeof(IMAGE_SECTION_HEADER);
        if (memcmp(psec->Name, ".text", 5) == 0) { // 是否.text段
            BYTE *offset = (BYTE *)hmodule + psec->VirtualAddress + psec->Misc.VirtualSize; // 计算空白区域偏移量
            offset += 16 - (INT_PTR)offset % 16; // 对齐16字节
            long long *buf = (long long *)offset;
            while (buf[0] != 0 || buf[1] != 0) // 找到一块全是0的区域
                buf += 16;
            return (void *)buf;
        }
    }
    return 0;
}
#endif

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

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

#ifdef _CPU_X64
#define ENTRY_LEN 9
#endif
#ifdef _CPU_X86
#define ENTRY_LEN 5
#endif

void dohook()
{
    HMODULE hmodule = GetModuleHandleA("kernelbase.dll");
    hook_func = (void *)GetProcAddress(hmodule, "WriteConsoleA");
    // 允许func_ptr处最前面的5字节内存可读可写可执行
    VirtualProtect(hook_func, HOOK_JUMP_LEN, PAGE_EXECUTE_READWRITE, NULL);
    // 使用VirtualAlloc申请内存，使其可读可写可执行
    old_entry = VirtualAlloc(NULL, 32, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
#ifdef _CPU_X64
    union
    {
        void *ptr;
        struct
        {
            long lo;
            long hi;
        };
    } ptr64;
    void *blank = FindModuleTextBlankAlign(hmodule); // 找到第一处空白区域
    VirtualProtect(blank, 14, PAGE_EXECUTE_READWRITE, NULL); // 可读写
    hook_jump[0] = 0xE9; // 跳转代码
    *(long *)&hook_jump[1] = (BYTE *)blank - (BYTE *)hook_func - 5; // 跳转到空白区域
    ptr64.ptr = (void *)fk_WriteConsoleA;
    BYTE blank_jump[14];
    blank_jump[0] = 0x68; // push xxx
    *(long *)&blank_jump[1] = ptr64.lo; // xxx，即地址的低4位
    blank_jump[5] = 0xC7;
    blank_jump[6] = 0x44;
    blank_jump[7] = 0x24;
    blank_jump[8] = 0x04; // mov dword [rsp+4], yyy
    *(long *)&blank_jump[9] = ptr64.hi; // yyy，即地址的高4位
    blank_jump[13] = 0xC3; // ret
    // 写入真正的跳转代码到空白区域
    WriteProcessMemory(GetCurrentProcess(), blank, &blank_jump, 14, NULL);
    // 保存原来的入口代码
    memcpy(old_entry, hook_func, ENTRY_LEN);
    ptr64.ptr = (BYTE *)hook_func + ENTRY_LEN;
    // 设置新的跳转代码
    BYTE *new_jump = (BYTE *)old_entry + ENTRY_LEN;
    new_jump[0] = 0x68;
    *(long *)(new_jump + 1) = ptr64.lo;
    new_jump[5] = 0xC7;
    new_jump[6] = 0x44;
    new_jump[7] = 0x24;
    new_jump[8] = 0x04;
    *(long *)(new_jump + 9) = ptr64.hi;
    new_jump[13] = 0xC3;
#endif
#ifdef _CPU_X86
    hook_jump[0] = 0xE9; // 跳转代码
    *(long *)&hook_jump[1] = (BYTE *)fk_WriteConsoleA - (BYTE *)hook_func - 5; // 直接到hook的代码
    memcpy(old_entry, hook_func, ENTRY_LEN); // 保存入口
    BYTE *new_jump = (BYTE *)old_entry + ENTRY_LEN;
    *new_jump = 0xE9; // 跳回去的代码
    *(long *)(new_jump + 1) = (BYTE *)hook_func + ENTRY_LEN - new_jump - 5;
#endif
    _WriteConsoleA = (WRITECONSOLEA)old_entry;
    WriteProcessMemory(GetCurrentProcess(), hook_func, &hook_jump, HOOK_JUMP_LEN, NULL);
}

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

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()
{
    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);
}