#include "ntddk.h"
#include "main.h"
#include "LDasm.h"

PKSERVICE_TABLE_DESCRIPTOR KeServiceDescriptorTableShadow = NULL;
PVOID pWin32kBase = NULL;
ULONG g_NtUserSendInput = 0;

VOID GetWin32kBase(PDRIVER_OBJECT pDriverObj)
{
PLIST_ENTRY pList = NULL;
//InitializeListHead(&pList);

PLDR_DATA_TABLE_ENTRY Ldr = NULL;

pList = ( (PLIST_ENTRY)pDriverObj->DriverSection )->Flink;
do {
Ldr = CONTAINING_RECORD(pList,LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
if (Ldr->EntryPoint &&Ldr->FullDllName.Buffer)
{
//DbgPrint(Ldr->FullDllName.Buffer);

if ( !_wcsicmp(Ldr->FullDllName.Buffer, L"\\systemroot\\system32\\win32k.sys") )
{
//比较模块名字，应该比较FullDllName，单单比较BaseDllName很可能会遇到同名文件。
pWin32kBase = Ldr->DllBase;//保存模块基址
break;
}
}
pList = pList->Flink;//下一个链表
} while ( pList != ((LIST_ENTRY*)pDriverObj->DriverSection)->Flink );
}


VOID GetSSDTShadowBase()
{
UCHAR *cPtr;
UCHAR *pOpcode;
ULONG Length;
for (cPtr = (PUCHAR)KeAddSystemServiceTable;
cPtr < (PUCHAR)KeAddSystemServiceTable + PAGE_SIZE;
cPtr += Length)
{
Length = SizeOfCode(cPtr, &pOpcode);

if (!Length) break;

if ( *(PUSHORT)cPtr == 0x888D )
{
KeServiceDescriptorTableShadow = (PKSERVICE_TABLE_DESCRIPTOR)(*(ULONG *)((ULONG)pOpcode + 2));
break;
}
}
}


NTSTATUS DispatchCreate(PDEVICE_OBJECT pDevObj, PIRP pIrp)
{
if (g_NtUserSendInput==0)
{
GET_INPUT();
if (g_NtUserSendInput!=0)
{
DbgPrint("Get OK\n");
}
}
else
{
//在这里恢复？创建的时候？
RestoreShadow();
DbgPrint("hook OK\n");
}

RestoreShadow();
pIrp->IoStatus.Status = STATUS_SUCCESS;
pIrp->IoStatus.Information = 0;
DbgPrint("IRP_MJ_CREATE\n");
IoCompleteRequest(pIrp, IO_NO_INCREMENT);
return STATUS_SUCCESS;
}

NTSTATUS DispatchClose(PDEVICE_OBJECT pDevObj, PIRP pIrp)
{
pIrp->IoStatus.Status = STATUS_SUCCESS;
pIrp->IoStatus.Information = 0;
DbgPrint("IRP_MJ_CLOSE\n");
IoCompleteRequest(pIrp, IO_NO_INCREMENT);
return STATUS_SUCCESS;
}

VOID DriverUnload(PDRIVER_OBJECT pDriverObj)
{
DeleteDrive(pDriverObj);
DbgPrint("Unload");
}

NTSTATUS DeleteDrive(PDRIVER_OBJECT pDriverobj)
{
UNICODE_STRING strLink;
RtlInitUnicodeString(&strLink, LINK_NAME);
IoDeleteSymbolicLink(&strLink);
IoDeleteDevice(pDriverobj->DeviceObject);
DbgPrint("Unloaded\n");
return STATUS_SUCCESS;
}
NTSTATUS Createdrive(PDRIVER_OBJECT pDriverobj)
{
NTSTATUS status = STATUS_SUCCESS;
UNICODE_STRING ustrLinkName;
UNICODE_STRING ustrDevName;
PDEVICE_OBJECT pDevObj;
RtlInitUnicodeString(&ustrDevName, DEVICE_NAME);

status = IoCreateDevice(pDriverobj,
0,
&ustrDevName,
FILE_DEVICE_UNKNOWN,
0,
FALSE,
&pDevObj);

DbgPrint("Device Name %S", ustrDevName.Buffer);

if (!NT_SUCCESS(status))
{
DbgPrint("IoCreateDevice = 0x%x\n", status);
return status;
}

RtlInitUnicodeString(&ustrLinkName, LINK_NAME);

status = IoCreateSymbolicLink(&ustrLinkName, &ustrDevName);
if (!NT_SUCCESS(status))
{
DbgPrint("IoCreateSymbolicLink = 0x%x\n", status);
IoDeleteDevice(pDevObj);
return status;
}

DbgPrint("SymbolicLink:%S", ustrLinkName.Buffer);
return STATUS_SUCCESS;
}
NTSTATUS DriverEntry(PDRIVER_OBJECT pDriverObj, PUNICODE_STRING pRegistryString)
{
pDriverObj->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
pDriverObj->MajorFunction[IRP_MJ_CLOSE] = DispatchClose;
pDriverObj->DriverUnload = DriverUnload;

if(Createdrive(pDriverObj)==STATUS_SUCCESS)
{
GetWin32kBase(pDriverObj);
DbgPrint("Win32Base:%X",pWin32kBase);
GetSSDTShadowBase();
DbgPrint("KeServiceDescriptorTableShadow:%X",KeServiceDescriptorTableShadow);
}
return STATUS_SUCCESS;
}



NTSTATUS GET_INPUT()
{
NTSTATUS status;
HANDLE hFile;//文件句柄
OBJECT_ATTRIBUTES ObjAttr;
UNICODE_STRING ustrWin32k;
IO_STATUS_BLOCK ioStatus;
ULONG ulShadowRaw = 0;
ULONG ulShadowBase = 0;
PVOID PoolArea = NULL;
FILE_POSITION_INFORMATION fpi;
LARGE_INTEGER Offset;
ULONG OrigAddress = 0;
ULONG CurAddress = 0;
ULONG i = 0;
ULONG ulCount = 0;
PULONG pAddr;
ULONG hookAddr=0;

if ( pWin32kBase == NULL ||
KeServiceDescriptorTableShadow == NULL)
{
DbgPrint("Error.");
return STATUS_UNSUCCESSFUL;
}
//索引为1的项目？

ulCount = KeServiceDescriptorTableShadow[1].Limit;//Linit就是表中函数的个数

DbgPrint("Count Of Shadow : %d\n", ulCount );

ulShadowBase = *(ULONG*)&KeServiceDescriptorTableShadow[1].Base;//得到基址

DbgPrint("ulShadowBase = 0x%X\n",ulShadowBase);
//镜像中的偏移,file offset???
ulShadowRaw = ulShadowBase - (ULONG)pWin32kBase;
//ulShadowRaw = RVAToRaw(pWin32kBase,ulShadowBase);

DbgPrint("ulShadowRaw = 0x%X\n",ulShadowRaw);

RtlInitUnicodeString(&ustrWin32k, L"\\SystemRoot\\System32\\win32k.sys");
//分配空间
PoolArea = ExAllocatePool( PagedPool, sizeof(ULONG) * ulCount );
//分配空间，用于保存读取到的数据，因为每个地址的长度sizeof(ULONG),个数是ulCount，所以相乘

if (!PoolArea) {
DbgPrint("PoolArea is null\n");
return STATUS_UNSUCCESSFUL;
}

RtlZeroMemory(&ObjAttr, sizeof(ObjAttr) );
//获取Win32k.sys的属性
InitializeObjectAttributes(
&ObjAttr,
&ustrWin32k,
OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
NULL,
NULL);
//打开文件win32K.SYS
status = IoCreateFile(
&hFile,
FILE_READ_ATTRIBUTES,
&ObjAttr,
&ioStatus,
0,
FILE_ATTRIBUTE_NORMAL,
FILE_SHARE_READ,
FILE_OPEN,
0,
NULL,
0,
CreateFileTypeNone,
NULL,
IO_NO_PARAMETER_CHECKING);

if ( !NT_SUCCESS(status) ) {
DbgPrint("IoCreateFile Error : 0x%X", status);
goto __exit;
}

//设置文件偏移
Offset.LowPart = ulShadowRaw;
Offset.HighPart = 0;
//开始读取数据
status = ZwReadFile (
hFile,
NULL,
NULL,
NULL,
&ioStatus,
//从文件读出到分配空间
PoolArea,
ulCount*sizeof(ULONG),
//偏移
&Offset,
NULL);

if ( !NT_SUCCESS(status) ) {
DbgPrint("ZwReadFile Error : 0x%X");
goto __exit;
}
//改变指针类型
pAddr = (PULONG)PoolArea;
//比较原始地址与当前的地址并且输出调试
_asm
{
CLI
MOV EAX,CR0
AND EAX,NOT 10000H
MOV CR0,EAX
}

OrigAddress = *pAddr;//指向原始地址
CurAddress = KeServiceDescriptorTableShadow[1].Base[502];//读取当前地址

if (g_NtUserSendInput==0)
{
__asm
{
MOV EAX,g_NtUserSendInput
CMP EAX,0
JNE L_Get1
MOV EAX,CurAddress
MOV BX,WORD PTR DS:[EAX]
CMP BX,0x186A
JNE L_Get1
MOV EAX,DWORD PTR DS:[EAX+3]
MOV g_NtUserSendInput,EAX
L_Get1:
}
DbgPrint("获取SendInput:%X\r\n",g_NtUserSendInput);
}




_asm
{
MOV EAX,CR0
OR EAX,10000h
MOV CR0,EAX
STI
}

__exit:
if (PoolArea) {
ExFreePool(PoolArea);
//释放空间
}
if (hFile) {
ZwClose(hFile);
//关闭句柄
}
return status;
}

NTSTATUS RestoreShadow()
{
NTSTATUS status;
HANDLE hFile;//文件句柄
OBJECT_ATTRIBUTES ObjAttr;
UNICODE_STRING ustrWin32k;
IO_STATUS_BLOCK ioStatus;
ULONG ulShadowRaw = 0;
ULONG ulShadowBase = 0;
PVOID PoolArea = NULL;
FILE_POSITION_INFORMATION fpi;
LARGE_INTEGER Offset;
ULONG OrigAddress = 0;
ULONG CurAddress = 0;
ULONG i = 0;
ULONG ulCount = 0;
PULONG pAddr;
ULONG hookAddr=0;

if ( pWin32kBase == NULL ||
KeServiceDescriptorTableShadow == NULL)
{
DbgPrint("Error.");
return STATUS_UNSUCCESSFUL;
}
//索引为1的项目？

ulCount = KeServiceDescriptorTableShadow[1].Limit;//Linit就是表中函数的个数

DbgPrint("Count Of Shadow : %d\n", ulCount );

ulShadowBase = *(ULONG*)&KeServiceDescriptorTableShadow[1].Base;//得到基址

DbgPrint("ulShadowBase = 0x%X\n",ulShadowBase);
//镜像中的偏移,file offset???
ulShadowRaw = ulShadowBase - (ULONG)pWin32kBase;
//ulShadowRaw = RVAToRaw(pWin32kBase,ulShadowBase);

DbgPrint("ulShadowRaw = 0x%X\n",ulShadowRaw);

RtlInitUnicodeString(&ustrWin32k, L"\\SystemRoot\\System32\\win32k.sys");
//分配空间
PoolArea = ExAllocatePool( PagedPool, sizeof(ULONG) * ulCount );
//分配空间，用于保存读取到的数据，因为每个地址的长度sizeof(ULONG),个数是ulCount，所以相乘

if (!PoolArea) {
DbgPrint("PoolArea is null\n");
return STATUS_UNSUCCESSFUL;
}

RtlZeroMemory(&ObjAttr, sizeof(ObjAttr) );
//获取Win32k.sys的属性
InitializeObjectAttributes(
&ObjAttr,
&ustrWin32k,
OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
NULL,
NULL);
//打开文件win32K.SYS
status = IoCreateFile(
&hFile,
FILE_READ_ATTRIBUTES,
&ObjAttr,
&ioStatus,
0,
FILE_ATTRIBUTE_NORMAL,
FILE_SHARE_READ,
FILE_OPEN,
0,
NULL,
0,
CreateFileTypeNone,
NULL,
IO_NO_PARAMETER_CHECKING);

if ( !NT_SUCCESS(status) ) {
DbgPrint("IoCreateFile Error : 0x%X", status);
goto __exit;
}

//设置文件偏移
Offset.LowPart = ulShadowRaw;
Offset.HighPart = 0;
//开始读取数据
status = ZwReadFile (
hFile,
NULL,
NULL,
NULL,
&ioStatus,
//从文件读出到分配空间
PoolArea,
ulCount*sizeof(ULONG),
//偏移
&Offset,
NULL);

if ( !NT_SUCCESS(status) ) {
DbgPrint("ZwReadFile Error : 0x%X");
goto __exit;
}
//改变指针类型
pAddr = (PULONG)PoolArea;
//比较原始地址与当前的地址并且输出调试
_asm
{
CLI
MOV EAX,CR0
AND EAX,NOT 10000H
MOV CR0,EAX
}

//OrigAddress = *pAddr;//指向原始地址
CurAddress = KeServiceDescriptorTableShadow[1].Base[502];//读取当前地址


_asm
{
MOV EAX,CR0
OR EAX,10000h
MOV CR0,EAX
STI
}

__exit:
if (PoolArea) {
ExFreePool(PoolArea);
//释放空间
}
if (hFile) {
ZwClose(hFile);
//关闭句柄
}
return status;
}