#include <ntddk.h>

#pragma pack(1)
typedef struct ServiceDescriptorEntry {
	unsigned int *ServiceTableBase;
	unsigned int *ServiceCounterTableBase; //仅适用于checked build版本
	unsigned int NumberOfServices;
	unsigned char *ParamTableBase;
} ServiceDescriptorTableEntry_t, *PServiceDescriptorTableEntry_t;
#pragma pack()

__declspec(dllimport) ServiceDescriptorTableEntry_t KeServiceDescriptorTable;

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


//global
ULONG	g_ntopenprocess;

typedef NTSTATUS (*NTOPENPROCESS) (
    __out PHANDLE ProcessHandle,
    __in ACCESS_MASK DesiredAccess,
    __in POBJECT_ATTRIBUTES ObjectAttributes,
    __in_opt PCLIENT_ID ClientId
    );

void PageProtectOn()
{
	__asm{//恢复内存保护  
		mov  eax,cr0
		or   eax,10000h
		mov  cr0,eax
		sti
	}
}

void PageProtectOff()
{
	__asm{//去掉内存保护
		cli
		mov  eax,cr0
		and  eax,not 10000h
		mov  cr0,eax
	}
}

BOOLEAN ProtectProcess(HANDLE ProcessId,char *str_ProtectObjName)
{
	NTSTATUS	status;
	PEPROCESS	process_obj;

	if (!MmIsAddressValid(str_ProtectObjName))
	{
		return FALSE;
	}
	
	if (ProcessId==0)
	{
		return FALSE;
	}

	status = PsLookupProcessByProcessId(ProcessId,&process_obj);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("error code:%X---ProcessId:%d",status,ProcessId));
		return FALSE;
	}

	if (strstr((char*)process_obj+0x16c,str_ProtectObjName)!=0)
	{
		ObDereferenceObject(process_obj);
		return TRUE;
	}

	ObDereferenceObject(process_obj);
	return FALSE;
}

NTSTATUS NewNtOpenProcess (
    __out PHANDLE ProcessHandle,
    __in ACCESS_MASK DesiredAccess,
    __in POBJECT_ATTRIBUTES ObjectAttributes,
    __in_opt PCLIENT_ID ClientId
    )
{
	if (ProtectProcess(ClientId->UniqueProcess,"notepad.exe"))
	{
		KdPrint(("%s",(char*)PsGetCurrentProcess()+0x16c));
		return STATUS_UNSUCCESSFUL;
	}

// 	if (ClientId->UniqueProcess==2604)
// 	{
// 		KdPrint(("%s",(char*)PsGetCurrentProcess()+0x16c));
// 		return STATUS_UNSUCCESSFUL;
// 	}

	return ((NTOPENPROCESS)g_ntopenprocess)(ProcessHandle,DesiredAccess,ObjectAttributes,ClientId);
}

NTSTATUS HookNtOpenProcess()
{
	NTSTATUS	status;

	status = STATUS_SUCCESS;
	
	PageProtectOff();

	g_ntopenprocess = KeServiceDescriptorTable.ServiceTableBase[190];
	
	KeServiceDescriptorTable.ServiceTableBase[190] = (unsigned int)NewNtOpenProcess;

	PageProtectOn();

	return status;
}

VOID UnHookNtOpenProcess()
{
	PageProtectOff();
	
	KeServiceDescriptorTable.ServiceTableBase[190] = (unsigned int)g_ntopenprocess;

	PageProtectOn();
}

VOID MyUnload(PDRIVER_OBJECT pDriverObject)
{
	UnHookNtOpenProcess();
}

NTSTATUS	DriverEntry(PDRIVER_OBJECT	pDriverObject,PUNICODE_STRING Reg_Path)
{
	HookNtOpenProcess();
	pDriverObject->DriverUnload = MyUnload;
	return STATUS_SUCCESS;
}