///////////////////////////////////////////////////////////////////////////////
///
/// Copyright (c) 2008 - <dts>
///
/// Original filename: ZwHook.c
/// Project          : ZwHook
/// Date of creation : 2008-11-19
/// Author(s)        : 梧桐
///
/// Purpose          : <description>
///
/// Revisions:
///  0000 [2008-11-19] Initial revision.
///
///////////////////////////////////////////////////////////////////////////////

#include "ntddk.h"

#pragma pack(1)
typedef struct ServiceDescriptorEntry
{
        unsigned int *ServiceTableBase;
        unsigned int *ServiceCounterTableBase; //Used only in checked build
        unsigned int NumberOfServices;
        unsigned char *ParamTableBase;
} ServiceDescriptorTableEntry_t, *PServiceDescriptorTableEntry_t;

#pragma pack()

__declspec(dllimport)  ServiceDescriptorTableEntry_t KeServiceDescriptorTable;

#define SYSTEMSERVICE(_function)  KeServiceDescriptorTable.ServiceTableBase[ *(PULONG)((PUCHAR)_function+1)]

PMDL  g_pmdlSystemCall;
PVOID *MappedSystemCallTable;

#define SYSCALL_INDEX(_Function) *(PULONG)((PUCHAR)_Function+1)

#define HOOK_SYSCALL(_Function, _Hook)  \
       (PVOID) InterlockedExchange( (PLONG) &MappedSystemCallTable[SYSCALL_INDEX(_Function)], (LONG) _Hook)

#define UNHOOK_SYSCALL(_Function, _Hook)  \
       InterlockedExchange( (PLONG) &MappedSystemCallTable[SYSCALL_INDEX(_Function)], (LONG) _Hook)

NTSYSAPI
NTSTATUS
NTAPI  ZwSetInformationThread(
                IN HANDLE  ThreadHandle,
                IN THREADINFOCLASS  ThreadInformationClass,
                IN PVOID  ThreadInformation,
                IN ULONG  ThreadInformationLength
                );

typedef NTSTATUS (*ZWSETINFORMATIONTHREAD)(
  IN HANDLE  ThreadHandle,
  IN THREADINFOCLASS  ThreadInformationClass,
  IN PVOID  ThreadInformation,
  IN ULONG  ThreadInformationLength
  );

ZWSETINFORMATIONTHREAD        OldZwSetInformationThread;

NTSTATUS NewZwSetInformationThread(
                   IN HANDLE  ThreadHandle,
                   IN THREADINFOCLASS  ThreadInformationClass,
                   IN PVOID  ThreadInformation,
                   IN ULONG  ThreadInformationLength)
{ 
  NTSTATUS ntStatus;

  if(ThreadInformationClass == 17) //ANTI-DEBUG
    ntStatus =  STATUS_SUCCESS;
  else
    ntStatus = ((ZWSETINFORMATIONTHREAD)(OldZwSetInformationThread))(
    ThreadHandle,
    ThreadInformationClass,
    ThreadInformation,
    ThreadInformationLength);
  return ntStatus;
}

VOID OnUnload(IN PDRIVER_OBJECT DriverObject)
{
   DbgPrint("ROOTKIT: OnUnload called\n");

   // unhook system calls
   UNHOOK_SYSCALL( ZwSetInformationThread, OldZwSetInformationThread);

   // Unlock and Free MDL
   if(g_pmdlSystemCall)
   {
      MmUnmapLockedPages(MappedSystemCallTable, g_pmdlSystemCall);
      IoFreeMdl(g_pmdlSystemCall);
   }
}


NTSTATUS DriverEntry(IN PDRIVER_OBJECT theDriverObject, 
           IN PUNICODE_STRING theRegistryPath)
{
   // Register a dispatch function for Unload
   theDriverObject->DriverUnload  = OnUnload; 

   // save old system call locations
   OldZwSetInformationThread =(ZWSETINFORMATIONTHREAD)(SYSTEMSERVICE(ZwSetInformationThread));

   // Map the memory into our domain so we can change the permissions on the MDL
   g_pmdlSystemCall = MmCreateMdl(NULL, KeServiceDescriptorTable.ServiceTableBase, KeServiceDescriptorTable.NumberOfServices*4);
   if(!g_pmdlSystemCall)
      return STATUS_UNSUCCESSFUL;

   MmBuildMdlForNonPagedPool(g_pmdlSystemCall);

   // Change the flags of the MDL
   g_pmdlSystemCall->MdlFlags = g_pmdlSystemCall->MdlFlags | MDL_MAPPED_TO_SYSTEM_VA;

   MappedSystemCallTable = MmMapLockedPages(g_pmdlSystemCall, KernelMode);

   // hook system calls
   HOOK_SYSCALL( ZwSetInformationThread, NewZwSetInformationThread);
                              
   return STATUS_SUCCESS;
}