#include <ntdll_windows.h>
#include <ntdll.h>

// extern "C"
// UINT_PTR
// NTAPI
// MwGenericCall(
//   ULONG SyscallNumber,
//   ULONG ArgumentCount,
//   va_list ArgumentList
//   );

#pragma section(".text")
__declspec(allocate(".text"))
UCHAR MwSyscallCode[] = {             //
                                  // ;
                                  // ; Save r12 to the shadow space.
                                  // ; This register will be held as a temporary
                                  // ; stack pointer for the syscall.
                                  // ;
  0x4C, 0x89, 0x64, 0x24, 0x08,   // mov     qword ptr [rsp + 8], r12
                                  //
                                  // ;
                                  // ; eax = SyscallNumber
                                  // ; r10 = ArgumentCount
                                  // ; r11 = ArgumentList
                                  // ; r12 = ArgumentList - sizeof(PVOID)
                                  // ;       (space for return address)
                                  // ;
                                  //
  0x8B, 0xC1,                     // mov     eax,  ecx
  0x44, 0x8B, 0xD2,               // mov     r10d, edx
  0x4D, 0x8B, 0xD8,               // mov     r11,  r8
  0x4D, 0x8D, 0x63, 0xF8,         // lea     r12,  qword ptr [r11 - 8]
                                  //
                                  // ;
                                  // ; Assign first 4 arguments (rcx, rdx, r8, r9)
                                  // ; from the ArgumentList.
                                  // ;
  0x41, 0x83, 0xFA, 0x01,         // cmp     r10d, 1
  0x49, 0x0F, 0x43, 0x0B,         // cmovae  rcx, qword ptr [r11 +  0]
                                  //
  0x41, 0x83, 0xFA, 0x02,         // cmp     r10d, 2
  0x49, 0x0F, 0x43, 0x53, 0x08,   // cmovae  rdx, qword ptr [r11 +  8]
                                  //
  0x41, 0x83, 0xFA, 0x03,         // cmp     r10d, 3
  0x4D, 0x0F, 0x43, 0x43, 0x10,   // cmovae  r8,  qword ptr [r11 + 16]
                                  //
  0x41, 0x83, 0xFA, 0x04,         // cmp     r10d, 4
  0x4D, 0x0F, 0x43, 0x4B, 0x18,   // cmovae  r9,  qword ptr [r11 + 24]
                                  //
                                  // ;
                                  // ; The syscall handler expects the first
                                  // ; parameter in the r10 register.
                                  // ;
  0x4C, 0x8B, 0xD1,               // mov     r10, rcx
                                  //
                                  // ;
                                  // ; Stack pivot, syscall, restore stack.
                                  // ;
  0x4C, 0x87, 0xE4,               // xchg    r12, rsp
  0x0F, 0x05,                     // syscall
  0x4C, 0x87, 0xE4,               // xchg    r12, rsp
                                  //
                                  // ;
                                  // ; Restore r12 register from the shadow space.
                                  // ;
  0x4C, 0x8B, 0x64, 0x24, 0x08,   // mov     r12, qword ptr [rsp + 8]
                                  //
  0xC3,                           // retn
};                                //

typedef UINT_PTR (NTAPI* GENERIC_CALL_FN)(
  _In_ ULONG SyscallNumber,
  _In_ ULONG ArgumentCount,
  _In_ va_list ArgumentList
  );

GENERIC_CALL_FN MwGenericCall = (GENERIC_CALL_FN)&MwSyscallCode;

static const unsigned char ASM_MOV_EAX_IMM = 0xb8; // followed by DWORD
static const unsigned char ASM_RET = 0xc3;
static const unsigned char ASM_MOV_R10_RCX[] = { 0x4c, 0x8b, 0xD1 };
static const unsigned char ASM_SYSCALL[] = { 0x0f, 0x05 };

static const unsigned char ASM_SYSCALL_STUB[] = {
  0x4c, 0x8b, 0xD1, // mov r10, rcx
  0xb8              // mov eax, imm32 - follows DWORD
};

NTSTATUS
NTAPI
MwDisasmGetSyscallNumber(
  VOID* FunctionAddress,
  ULONG* SyscallNumber
  )
{
  NTSTATUS Status = STATUS_UNSUCCESSFUL;

  if ((*(PULONG)FunctionAddress) == (*(PULONG)ASM_SYSCALL_STUB))
  {
    *SyscallNumber = *(ULONG*)((BYTE*)FunctionAddress + sizeof(ASM_SYSCALL_STUB));

    Status = STATUS_SUCCESS;
  }

  return Status;
}

NTSTATUS
NTAPI
MwGetFunctionAddress(
  VOID* ImageBase,
  const CHAR* FunctionName,
  VOID** FunctionAddress,
  ULONG* SyscallNumber
  )
{
  TEB* Teb = NtCurrentTeb();

  LIST_ENTRY* ModuleListHead = &Teb->ProcessEnvironmentBlock->Ldr->InLoadOrderModuleList;
  LIST_ENTRY* ModuleListNextEntry;

  BYTE* NtdllBase;

  for (
    ModuleListNextEntry = ModuleListHead;
    ModuleListNextEntry->Flink != ModuleListHead;
    ModuleListNextEntry = ModuleListNextEntry->Flink
    )
  {
    LDR_DATA_TABLE_ENTRY* LdrEntry = (LDR_DATA_TABLE_ENTRY*)ModuleListNextEntry->Flink;
    if (!wcscmp(LdrEntry->BaseDllName.Buffer, L"ntdll.dll"))
    {
      NtdllBase = (BYTE*)LdrEntry->DllBase;
      break;
    }
  }

  IMAGE_DOS_HEADER* DosHeader = (IMAGE_DOS_HEADER*)NtdllBase;
  IMAGE_NT_HEADERS* NtHeader = (IMAGE_NT_HEADERS*)(NtdllBase + DosHeader->e_lfanew);
  IMAGE_DATA_DIRECTORY* ExportDataDirectory = (IMAGE_DATA_DIRECTORY*)&NtHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
  IMAGE_EXPORT_DIRECTORY* ExportDirectory = (IMAGE_EXPORT_DIRECTORY*)(NtdllBase + ExportDataDirectory->VirtualAddress);

  DWORD* FunctionPtrRvaArray = (DWORD*)(NtdllBase + ExportDirectory->AddressOfFunctions);
  DWORD* FunctionNameRvaArray = (DWORD*)(NtdllBase + ExportDirectory->AddressOfNames);
  USHORT* FunctionOrdinalRvaArray = (USHORT*)(NtdllBase + ExportDirectory->AddressOfNameOrdinals);

  int FunctionNameRvaArraySize = ExportDirectory->NumberOfNames;

  CHAR* ExportFunctionName;
  VOID* ExportFunctionPtr;

  for (int i = 0; i < FunctionNameRvaArraySize; i++)
  {
    ExportFunctionName = (CHAR*)(NtdllBase + FunctionNameRvaArray[i]);
    ExportFunctionPtr = (VOID*)(NtdllBase + FunctionPtrRvaArray[FunctionOrdinalRvaArray[i]]);

    if (!strcmp(ExportFunctionName, FunctionName))
    {
      MwDisasmGetSyscallNumber(ExportFunctionPtr, SyscallNumber);
      *FunctionAddress = ExportFunctionPtr;

      return STATUS_SUCCESS;
    }
  }

  return STATUS_UNSUCCESSFUL;
}

NTSTATUS
NTAPI
MwCallNtdll(
  const CHAR* FunctionName,
  ULONG ArgumentCount,
  ...
  )
{
  VOID* FunctionAddress;
  ULONG SyscallNumber;
  MwGetFunctionAddress(NULL, FunctionName, &FunctionAddress, &SyscallNumber);

  va_list ArgumentList;
  va_start(ArgumentList, ArgumentCount);

  UINT_PTR Result = MwGenericCall(SyscallNumber, ArgumentCount, ArgumentList);

  va_end(ArgumentList);

  return (NTSTATUS)Result;
}

//
// Variadic macro rape.
//

#define MW_MACRO_EVAL(...)                                                    __VA_ARGS__
#define MW_MACRO_EXPAND(expr)                                                 expr
#define MW_VA_NUM_ARGS_IMPL(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, N, ...)  N
#define MW_VA_NUM_ARGS(...)                                                   MW_MACRO_EXPAND(MW_VA_NUM_ARGS_IMPL(__VA_ARGS__, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1))

#define MW_CALL_WITH_NUM_ARGS(FunctionName, ArgumentCount, ...)               MwCallNtdll(FunctionName, ArgumentCount, __VA_ARGS__)

#define MW_CALL_VA_STR(FunctionName, ...)                                     MW_CALL_WITH_NUM_ARGS(FunctionName, MW_VA_NUM_ARGS(__VA_ARGS__), __VA_ARGS__)
#define MW_CALL_VA(FunctionName, ...)                                         MW_CALL_WITH_NUM_ARGS(#FunctionName, MW_VA_NUM_ARGS(__VA_ARGS__), __VA_ARGS__)
#define MW_CALL_STR(FunctionName, Args)                                       MW_CALL_VA_STR(FunctionName,  MW_MACRO_EVAL Args)
#define MW_CALL(FunctionName, Args)                                           MW_CALL_VA_STR(#FunctionName, MW_MACRO_EVAL Args)

#define NtCall MW_CALL

void NtMain()
{
  NTSTATUS Status;

  HANDLE SectionHandle;
  LARGE_INTEGER MaximumSize;
  MaximumSize.QuadPart = 4096;
  Status = NtCall(NtCreateSection, (&SectionHandle, SECTION_ALL_ACCESS, NULL, &MaximumSize, PAGE_EXECUTE_READWRITE, SEC_COMMIT, NULL));

  PVOID BaseAddress = NULL;
  SIZE_T ViewSize = 4096;
  Status = NtCall(NtMapViewOfSection, (SectionHandle, NtCurrentProcess(), &BaseAddress, 0, 4096, NULL, &ViewSize, ViewUnmap, 0, PAGE_READWRITE));

  RtlCopyMemory(BaseAddress, "HELLO", 6);
  NtTestAlert();
}