#include "ntddk.h"
#include <windef.h>
#include <stdlib.h>
#include "dayed.h"

#define OBJECT_TO_OBJECT_HEADER( o ) CONTAINING_RECORD( (o), OBJECT_HEADER, Body )    
extern POBJECT_TYPE *PsProcessType;
extern POBJECT_TYPE *PsThreadType;
KSPIN_LOCK SDTSpinLock;
KIRQL  oldIrql;

BYTE g_HookCode[5] = { 0xe9, 0, 0, 0, 0 };
BYTE g_OrigCode[5] = { 0 };
BYTE jmp_orig_code[7] = { 0xEA, 0, 0, 0, 0, 0x08, 0x00 }; 
char* MyProtectName = "notepad.exe";
int MyProtectPID=0;

PEPROCESS EPROCESS,ProtectedProcess;

void StartHook ();

NTSYSAPI
NTSTATUS
NTAPI ObOpenObjectByPointer(
  IN PVOID  Object,
  IN ULONG  HandleAttributes,
  IN PACCESS_STATE  PassedAccessState  OPTIONAL,
  IN ACCESS_MASK  DesiredAccess  OPTIONAL,
  IN POBJECT_TYPE  ObjectType  OPTIONAL,
  IN KPROCESSOR_MODE  AccessMode,
  OUT PHANDLE  Handle);

NTKERNELAPI
NTSTATUS
PsLookupProcessByProcessId (
      IN ULONG          ProcessId,
      OUT PEPROCESS     *Process
);

NTKERNELAPI
PEPROCESS
NTAPI
IoThreadToProcess (
    IN PETHREAD Thread
);

void StopHook ()
{
  WPOFF();
  KeAcquireSpinLock( &SDTSpinLock, &oldIrql );
  RtlCopyMemory ( (BYTE*)ObOpenObjectByPointer, g_OrigCode, 5 );
  KeReleaseSpinLock( &SDTSpinLock, oldIrql );
  WPON();
}

__declspec (naked) 
NTSTATUS
Proxy_ObOpenObjectByPointer(
  IN PVOID  Object,
  IN ULONG  HandleAttributes,
  IN PACCESS_STATE  PassedAccessState  OPTIONAL,
  IN ACCESS_MASK  DesiredAccess  OPTIONAL,
  IN POBJECT_TYPE  ObjectType  OPTIONAL,
  IN KPROCESSOR_MODE  AccessMode,
  OUT PHANDLE  Handle)
{
  __asm {  // 共12字节
      _emit 0x90
      _emit 0x90
      _emit 0x90
      _emit 0x90
      _emit 0x90  // 前5字节实现原函数的头5字节功能
      _emit 0x90  // 这个填充jmp
      _emit 0x90
      _emit 0x90
      _emit 0x90
      _emit 0x90  // 这4字节保存原函数+5处的地址
      _emit 0x90  
      _emit 0x90  // 因为是长转移,所以必须是 0x0080
  }
}


NTSTATUS __stdcall 
fake_ObOpenObjectByPointer(
  IN PVOID  Object,
  IN ULONG  HandleAttributes,
  IN PACCESS_STATE  PassedAccessState  OPTIONAL,
  IN ACCESS_MASK  DesiredAccess  OPTIONAL,
  IN POBJECT_TYPE  ObjectType  OPTIONAL,
  IN KPROCESSOR_MODE  AccessMode,
  OUT PHANDLE  Handle)
{
     if ((Object != NULL) && (MmIsAddressValid(Object))) // 地址有效性验证
    {
        if (((POBJECT_HEADER)(OBJECT_TO_OBJECT_HEADER(Object)))->Type == *PsProcessType) // 若为进程对象
        {
            if ((ProtectedProcess !=PsGetCurrentProcess())) // 若操作者不是受保护的进程自己
            {
                if (Object == ProtectedProcess) // 若被操作进程是受保护进程
                {
                    return STATUS_ACCESS_DENIED; // 拒绝访问
                }

             }
         }
        else
             if (OBJECT_TO_OBJECT_HEADER(Object) -> Type == *PsThreadType) // 若为线程对象
            {
                EPROCESS = IoThreadToProcess(Object); // 获取线程对应进程的 EPROCESS
                if (EPROCESS == ProtectedProcess) // 若是受保护进程
                {
                    if ((PsGetCurrentProcess() != ProtectedProcess)) // 若操作者不是受保护进程自己
                    {
                        return STATUS_ACCESS_DENIED; // 拒绝访问
                    }
                 }
             }
    }
    return Proxy_ObOpenObjectByPointer (Object, HandleAttributes,PassedAccessState,DesiredAccess,ObjectType,AccessMode,Handle);
}


void StartHook ()
{ 
  RtlCopyMemory (g_OrigCode, (BYTE*)ObOpenObjectByPointer, 5);
  DbgPrint("g_OrigCode address at %x\n",g_OrigCode);
  *( (ULONG*)(g_HookCode + 1) ) = (ULONG)fake_ObOpenObjectByPointer - (ULONG)ObOpenObjectByPointer - 5;
  DbgPrint("fake_ObOpenObjectByPointer address at %x\n",fake_ObOpenObjectByPointer);
  DbgPrint("ObOpenObjectByPointer address at %x\n",ObOpenObjectByPointer);
  WPOFF();
  KeAcquireSpinLock( &SDTSpinLock, &oldIrql );
  RtlCopyMemory ( (BYTE*)ObOpenObjectByPointer, g_HookCode, 5 );
  *( (ULONG*)(jmp_orig_code + 1) ) = (ULONG) ( (BYTE*)ObOpenObjectByPointer + 5 );
  RtlCopyMemory ( (BYTE*)Proxy_ObOpenObjectByPointer, g_OrigCode, 5);
  RtlCopyMemory ( (BYTE*)Proxy_ObOpenObjectByPointer+ 5, jmp_orig_code, 7);
  KeReleaseSpinLock( &SDTSpinLock, oldIrql );
  WPON();
  DbgPrint("Proxy_ObOpenObjectByPointer address at %x\n",Proxy_ObOpenObjectByPointer);
}

VOID Unload(PDRIVER_OBJECT  DriverObject)
{  
if (MyProtectPID!=0)
{
   StopHook();
}
}

NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING str)
{
  NTSTATUS      ntStatus;
  char ProcessName[256];
  ULONG cbBuffer = 0x8000; 
  PSYSTEM_PROCESS_INFORMATION pInfo;
  VOID* pBuffer = NULL;
  
  DriverObject->DriverUnload = Unload;
          pBuffer = ExAllocatePool (NonPagedPool, cbBuffer); 
        if (pBuffer == NULL) 
        {
            return 1;
        }
        ntStatus = ZwQuerySystemInformation(5, pBuffer, cbBuffer, NULL);
    
        if (!NT_SUCCESS(ntStatus))
        {
            ExFreePool(pBuffer); 
            return 1; 
        }
   
    pInfo = (PSYSTEM_PROCESS_INFORMATION)pBuffer;
    
    while(1){
        LPWSTR pszProcessName = pInfo->ProcessName.Buffer;
        if (pszProcessName == NULL) 
        pszProcessName = L"NULL"; 

        wcstombs(ProcessName,pszProcessName,256); 
    //    DbgPrint("%s\tPid=%d\n",ProcessName,pInfo->ProcessId);
           if(_stricmp(MyProtectName,ProcessName)==0)
           {
              MyProtectPID=pInfo->ProcessId;
              DbgPrint("the MyProtectPID is %d\n",pInfo->ProcessId);
           }
             
        
                if (pInfo->NextEntryDelta == 0) 
            break; 

        pInfo = (PSYSTEM_PROCESS_INFORMATION)(((PUCHAR)pInfo)+ pInfo->NextEntryDelta);
    }
    ExFreePool(pBuffer); 
  
  if (MyProtectPID!=0)
  {
  ntStatus = PsLookupProcessByProcessId(MyProtectPID, &ProtectedProcess);
  if(NT_SUCCESS(ntStatus))
  {
    ObDereferenceObject(ProtectedProcess);
  }
    StartHook(); 
    DbgPrint("ObOpenObjectByPointer address at %x\n",ObOpenObjectByPointer);
    DbgPrint("Hook Start");
    return STATUS_SUCCESS;
    }
    DbgPrint("Can't Hook");
  return STATUS_SUCCESS;
}