学步园

/*****************************************************************
  文件名        : WssFilterDrv.c
  描述          : 跟踪记录现有设备驱动
  作者          : sinister
  最后修改日期  : 2006.1.09

 *****************************************************************/


#include "ntddk.h"
#include "string.h"
#include "stdio.h"

//
// 输入标志
//
#define IN_BUFFER_FLAG  0x925
//
// 输出标志
//
#define OUT_BUFFER_FLAG 0x926

////////////////////////////////////////////////////////////////
// 结构作用 ：设备对象扩展结构，用于保存过滤设备信息
// 注意 ： 需设备对象 DEVICE_OBJECT 有效情况下使用
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
typedef struct _DEVICE_EXTENSION
{
    PDEVICE_OBJECT DeviceObject;  // 过滤设备对象
    PDEVICE_OBJECT TargetDevice;  // 附加后的设备对象
    PFILE_OBJECT FileObject;      // 原设备的文件对象
    ULONG DeviceExtensionFlags;   // 预留的扩展标志
}DEVICE_EXTENSION, * PDEVICE_EXTENSION;

VOID DriverUnload( IN PDRIVER_OBJECT DriverObject );
NTSTATUS DispatchPassThrough( IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp );
NTSTATUS CompletionRoutine( IN PDEVICE_OBJECT DeviceObject,
                            IN PIRP Irp,
                            IN PVOID Context );

NTSTATUS WriteLog( PVOID lpBuffer, ULONG uSize );
NTSTATUS GetSystemBuffer( PIRP Irp, PVOID pBuffer, SIZE_T Length, LONG* pMode );
NTSTATUS FormatMessage( ULONG IoControl,
                        PUCHAR pBuf,
                        ULONG uSize,
                        ULONG dwStatus,
                        ULONG uMode );
NTSTATUS ReadIOCtrlBuffer( PIRP Irp,
                           PIO_STACK_LOCATION pIrpStack,
                           ULONG dwStatus );


/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ：创建过滤设备将其附加到需要跟踪的设备上，保存设备相关
//        信息，返回附加后的驱动对象
// 注意 ： 
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
AttachFilterDevice( IN UNICODE_STRING* DeviceName, // 需要跟踪的设备名
                    IN PDRIVER_OBJECT  DriverObject, // 过滤驱动也就是本驱动的驱动对象
                    OUT PDRIVER_OBJECT* FilterDriverObject ) // 返回附加后的驱动对象
{
  PDEVICE_OBJECT DeviceObject; 
  PDEVICE_OBJECT FilterDeviceObject;
  PDEVICE_OBJECT TargetDevice; 
  PFILE_OBJECT FileObject; 
  PDEVICE_EXTENSION DevExt;

  NTSTATUS ntStatus; 

  //
  // 根据设备名称找到需要附加的设备对象
  //
  ntStatus = IoGetDeviceObjectPointer( DeviceName,
                                       FILE_ALL_ACCESS,
                                       &FileObject,
                                       &DeviceObject ); 

  if ( !NT_SUCCESS( ntStatus ) )
  {
    DbgPrint( "IoGetDeviceObjectPointer() %d/n", ntStatus );
    return ntStatus;
  } 

  //
  // 创建过滤设备对象
  //
  ntStatus = IoCreateDevice( DriverObject,
                             sizeof( DEVICE_EXTENSION ),
                             NULL,
                             FILE_DEVICE_UNKNOWN,
                             0,
                             FALSE,
                             &FilterDeviceObject ); 

  if ( !NT_SUCCESS( ntStatus ) )
  {
    ObDereferenceObject( FileObject ); 
    DbgPrint( "IoCreateDevice() %d!/n", ntStatus );
    return ntStatus;
  } 

  //
  // 得到设备扩展结构，以便下面保存过滤设备信息
  //
  DevExt = ( PDEVICE_EXTENSION ) FilterDeviceObject->DeviceExtension;

  //
  // 将过滤设备对象附加在目标设备对象之上，并返回附加后的原设备对象
  //
  TargetDevice = IoAttachDeviceToDeviceStack( FilterDeviceObject,
                                              DeviceObject ); 
  if ( !TargetDevice )
  {
    ObDereferenceObject( FileObject ); 
    IoDeleteDevice( FilterDeviceObject ); 
    DbgPrint( "IoAttachDeviceToDeviceStack() %d!/n", ntStatus );
    return STATUS_INSUFFICIENT_RESOURCES;
  } 

  //
  // 保存过滤设备信息
  //
  DevExt->DeviceObject = FilterDeviceObject; 
  DevExt->TargetDevice = TargetDevice; 
  DevExt->FileObject = FileObject;

  //
  // 设置过滤设备相关信息与标志
  //
  FilterDeviceObject->DeviceType = TargetDevice->DeviceType; 
  FilterDeviceObject->Characteristics = TargetDevice->Characteristics; 
  FilterDeviceObject->Flags |= ( TargetDevice->Flags & ( DO_DIRECT_IO |
                                                         DO_BUFFERED_IO ) ); 

  //
  // 返回附加后的驱动对象
  //
  *FilterDriverObject = TargetDevice->DriverObject;

  ObDereferenceObject( FileObject ); 

  return STATUS_SUCCESS;
}

/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 得到内核/用户模式 I/O 缓冲区数据，并返回当前数据所
//         在模式
// 注意 ： 在 "过滤模块" 中此函数是为了配合 ReadIOCtrlBuffer 
//         函数得到数据
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
GetSystemBuffer( PIRP Irp, PVOID pBuffer, SIZE_T Length, LONG* pMode )
{
  PUCHAR pSysBuffer = NULL;

  if ( pBuffer == NULL )
  {
    return STATUS_INSUFFICIENT_RESOURCES;
  }

  //
  // 判断当前 I/O 模式与数据存放的位置
  //
  switch ( Irp->RequestorMode )
  {
    case KernelMode:
      *pMode = 0; // 内核模式 I/O

      if ( Irp->UserBuffer )
      {
        pSysBuffer = Irp->UserBuffer;
      }
      else if ( Irp->MdlAddress )
      {
        pSysBuffer = ( PUCHAR )
                     MmGetSystemAddressForMdlSafe( Irp->MdlAddress,
                                                   NormalPagePriority );

        if ( pSysBuffer == NULL )
        {
          return STATUS_INSUFFICIENT_RESOURCES;
        }
      }

      break;

    case UserMode:
      *pMode = 1; // 用户模式 I/O

      if ( Irp->UserBuffer )
      {
        pSysBuffer = Irp->UserBuffer;
      }
      else if ( Irp->MdlAddress )
      {
        pSysBuffer = ( PUCHAR )
                     MmGetSystemAddressForMdlSafe( Irp->MdlAddress,
                                                   NormalPagePriority );

        if ( pSysBuffer == NULL )
        {
          return STATUS_INSUFFICIENT_RESOURCES;
        }
      }
      break;

    default:
      break;
  }

  //
  // 输出得到的缓冲区
  //
  RtlCopyMemory( pBuffer, pSysBuffer, Length );

  return STATUS_SUCCESS;
}

/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 格式化缓冲区数据，并记录日志
// 注意 ： 在 "过滤模块" 中此函数作用是格式化 GetSystemBuffer 
           函数得到的数据。因为此函数调用了日志函数，所以此函数
//         仅能在 IRQL == PASSVIE_LEVEL  级别运行。
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 : sinister
// 修改日期 : 2006.01.03
// 修改内容 : 在函数中增加了记录日志函数
/////////////////////////////////////////////////////////////////
NTSTATUS
FormatMessage( ULONG IoControl,
               PUCHAR pBuf,
               ULONG uSize,
               ULONG dwStatus,
               ULONG uMode )
{
  ULONG uLen;
  ULONG uCount = 0;
  PVOID pBuffer = NULL;
  UCHAR uTitle[100];
  UCHAR uTemp[3];

  if ( pBuf == NULL || uSize < 0 )
  {
    return STATUS_INSUFFICIENT_RESOURCES;
  }

  //
  // 低层驱动返回数据后，设置此标志
  //
  if ( dwStatus == OUT_BUFFER_FLAG )
  {
    if ( uMode == 0 )
    {
      sprintf( uTitle, "KERNEL MODE: IoControl=%x, OUTBUFFER:/n", IoControl );
    }
    else
    {
      sprintf( uTitle, "USER MODE: IoControl=%x, OUTBUFFER:/n", IoControl );
    }
  }

  //
  // 上层向低层驱动发送数据，设置此标志
  //
  else if ( dwStatus == IN_BUFFER_FLAG )
  {
    if ( uMode == 0 )
    {
      sprintf( uTitle, "KERNEL MODE: IoControl=%x, INBUFFER:/n", IoControl );
    }
    else
    {
      sprintf( uTitle, "USER MODE: IoControl=%x, INBUFFER:/n", IoControl );
    }
  }

  //
  // 记录日志，目的是记录分析，所以没有考虑效率。
  //
  WriteLog( uTitle, strlen( uTitle ) );

  for ( uLen = 0; uLen < uSize; uLen++ )
  {
    sprintf( uTemp, "%.2X ", pBuf[uLen] );
    WriteLog( uTemp, sizeof( uTemp ) );
    if ( uLen % 16 == 0 )
    {
      WriteLog( "/n", 1 );
    }
  }

  WriteLog( "/n", 1 );

  return STATUS_SUCCESS;
}

/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 得到 IRP_MJ_DEVICE_CONTROL 的 I/O 数据
// 注意 ： 
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
ReadIOCtrlBuffer( PIRP Irp, PIO_STACK_LOCATION pIrpStack, ULONG dwStatus )
{
  ULONG InBufLen;
  ULONG OutBufLen;
  ULONG IoControl;

  PUCHAR pSysBuffer = NULL;
  KEVENT KrnlEvent;
  NTSTATUS ntStatus;
  ULONG uMode;

  //
  // 得到设备控制字与 I/O 缓冲区长度
  //
  InBufLen = pIrpStack->Parameters.DeviceIoControl.InputBufferLength;
  OutBufLen = pIrpStack->Parameters.DeviceIoControl.OutputBufferLength;
  IoControl = pIrpStack->Parameters.DeviceIoControl.IoControlCode;

  if ( dwStatus == OUT_BUFFER_FLAG )
  {
    //
    // 如果是低层驱动返回的数据，缓冲区长度则要这样获得
    //
    OutBufLen = Irp->IoStatus.Information;

    pSysBuffer = ( PUCHAR ) ExAllocatePool( PagedPool,
                                            OutBufLen + sizeof( PUCHAR ) );
    if ( pSysBuffer == NULL )
    {
      return STATUS_INSUFFICIENT_RESOURCES;
    }

    ntStatus = GetSystemBuffer( Irp, pSysBuffer, OutBufLen, &uMode );
    if ( !NT_SUCCESS( ntStatus ) )
    {
      ExFreePool( pSysBuffer );
      return STATUS_INSUFFICIENT_RESOURCES;
    }

    FormatMessage( IoControl, pSysBuffer, OutBufLen, dwStatus, uMode );
  }
  else
  {
    pSysBuffer = ( PUCHAR ) ExAllocatePool( PagedPool,
                                            InBufLen + sizeof( PUCHAR ) );
    if ( pSysBuffer == NULL )
    {
      return STATUS_INSUFFICIENT_RESOURCES;
    }

    ntStatus = GetSystemBuffer( Irp, pSysBuffer, InBufLen, &uMode );
    if ( !NT_SUCCESS( ntStatus ) )
    {
      ExFreePool( pSysBuffer );
      return STATUS_INSUFFICIENT_RESOURCES;
    }

    FormatMessage( IoControl, pSysBuffer, InBufLen, dwStatus, uMode );
  }

  if ( pSysBuffer )
  {
    ExFreePool( pSysBuffer );
  }

  return Irp->IoStatus.Status;
}

/////////////////////////////////////////////////////////////////
// 函数类型 ：系统函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 替换指定设备的驱动例程使其指向当前过滤驱动的派遣例程
// 注意 ： 
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
DriverEntry( IN PDRIVER_OBJECT DriverObject, IN PUNICODE_STRING RegistryPath )
{
  PDRIVER_OBJECT FilterDrvObj; 
  UNICODE_STRING DeviceName;
  PDRIVER_DISPATCH EmptyDispatch; 
  NTSTATUS ntStatus; 
  ULONG iCount;

  EmptyDispatch = DriverObject->MajorFunction[IRP_MJ_CREATE];

  RtlInitUnicodeString( &DeviceName, L"//Device//FNT0" ); 
  FilterDrvObj = NULL;

  DriverObject->DriverUnload = DriverUnload;

  ntStatus = AttachFilterDevice( &DeviceName, DriverObject, &FilterDrvObj );
  if ( !NT_SUCCESS( ntStatus ) || FilterDrvObj == NULL )
  {
    DbgPrint( "Attach Device to failed!/n" );
    return STATUS_INSUFFICIENT_RESOURCES;
  }

  //
  // 只过滤目标驱动对象使用的例程
  //
  for ( iCount = 0; iCount < IRP_MJ_MAXIMUM_FUNCTION; iCount++ )
  {
    if ( FilterDrvObj->MajorFunction[iCount] != EmptyDispatch &&
         DriverObject->MajorFunction[iCount] == EmptyDispatch )
    {
      DriverObject->MajorFunction[iCount] = DispatchPassThrough;
    }
  } 

  return STATUS_SUCCESS;
} 

/////////////////////////////////////////////////////////////////
// 函数类型 ：系统函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 删除并卸载过滤设备
// 注意 ： 卸载列程中并没有做过多的处理，在过滤驱动中如果当前所
//         附加的设备IRP 很多的情况下，这样做可能会导致 BSOD
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
VOID
DriverUnload( IN PDRIVER_OBJECT DriverObject )
{
  PDEVICE_OBJECT FilterDevice; 
  PDEVICE_OBJECT TargetDevice; 
  PDEVICE_EXTENSION DevExt; 
  NTSTATUS ntStatus; 

  FilterDevice = DriverObject->DeviceObject;
  DevExt = ( PDEVICE_EXTENSION ) FilterDevice->DeviceExtension;
  TargetDevice = DevExt->TargetDevice;

  //
  // 从目标断开附加设备，并删除过滤设备
  //
  IoDetachDevice( TargetDevice ); 
  IoDeleteDevice( FilterDevice ); 

  return;
}


/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 过滤驱动的派遣例程
// 注意 ： 
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
DispatchPassThrough( IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp )
{
  PIO_STACK_LOCATION pIrpStack;
  PDEVICE_OBJECT pDeviceObject;
  PDEVICE_EXTENSION DevExtension = ( PDEVICE_EXTENSION )
                                   DeviceObject->DeviceExtension; 
  NTSTATUS ntStatus; 
  KEVENT KrnlEvent;

  pIrpStack = IoGetCurrentIrpStackLocation( Irp );

  //
  // 这里只过滤了 IRP_MJ_DEVICE_CONTROL 需要过滤更多的控制字可以
  // 在这里增加判断
  //
  if ( pIrpStack->MajorFunction == IRP_MJ_DEVICE_CONTROL )
  {
    //
    // 上层发送数据，设置 IN_BUFFER_FLAG 标志记录。
    //
    ReadIOCtrlBuffer( Irp, pIrpStack, IN_BUFFER_FLAG );

    //
    // 因为是上层过滤驱动，所以必须向低层驱动请求返回数据
    // 这里虽然没有对数据进行修改，但为了以后扩充方便
    // 还是复制了 IRP 堆栈和安装完成列程
    //
    IoCopyCurrentIrpStackLocationToNext( Irp );
    IoSetCompletionRoutine( Irp, CompletionRoutine, NULL, TRUE, TRUE, TRUE );

    ( void ) IoCallDriver( DevExtension->TargetDevice, Irp );
    //
    // 当低层驱动返回数据后，设置 OUT_BUFFER_FLAG 标志，再次记录。
    //
    ntStatus = ReadIOCtrlBuffer( Irp, pIrpStack, OUT_BUFFER_FLAG );

    return ntStatus;
  }

  //
  // 对于其他控制字则直接传递到低层驱动，无须安装完成列程
  //
  else
  {
    IoSkipCurrentIrpStackLocation( Irp );
    return IoCallDriver( DevExtension->TargetDevice, Irp );
  }
} 

/////////////////////////////////////////////////////////////////
// 函数类型 ：系统回调函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 
// 注意 ： 这里没有对返回数据做操作
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
CompletionRoutine( IN PDEVICE_OBJECT DeviceObject,
                   IN PIRP Irp,
                   IN PVOID Context )
{
  if ( Irp->PendingReturned )
  {
    IoMarkIrpPending( Irp );
  } 

  return STATUS_SUCCESS;
} 

/////////////////////////////////////////////////////////////////
// 函数类型 ：自定义工具函数
// 函数模块 : 过滤模块
////////////////////////////////////////////////////////////////
// 功能 ： 记录日志到 C 盘的 WSSFILTER.LOG 中
// 注意 ： 只能在 IRQL == PASSVIE_LEVEL 级别使用
/////////////////////////////////////////////////////////////////
// 作者 : sinister
// 发布版本 : 1.00.00
// 发布日期 : 2005.12.27
/////////////////////////////////////////////////////////////////
// 重   大   修   改   历   史
////////////////////////////////////////////////////////////////
// 修改者 :
// 修改日期 :
// 修改内容 :
/////////////////////////////////////////////////////////////////
NTSTATUS
WriteLog( PVOID lpBuffer, ULONG uSize )
{
  NTSTATUS status; 
  IO_STATUS_BLOCK IoStatusBlock; 
  OBJECT_ATTRIBUTES ObjectAttributes; 
  FILE_POSITION_INFORMATION FilePOTInfo; 
  FILE_STANDARD_INFORMATION FileSTDInfo; 

  ANSI_STRING ansiString; 
  UNICODE_STRING filenameString; 
  HANDLE hFile;
  WCHAR LogFile[] = L"//??//C://WSSFILTER.LOG"; 


  RtlInitUnicodeString( &filenameString, LogFile );     

  InitializeObjectAttributes( &ObjectAttributes,
                              &filenameString,
                              OBJ_CASE_INSENSITIVE,
                              NULL,
                              NULL ); 

  status = ZwCreateFile( &hFile,
                         GENERIC_WRITE | SYNCHRONIZE | GENERIC_READ,
                         &ObjectAttributes,
                         &IoStatusBlock,
                         0,
                         FILE_ATTRIBUTE_NORMAL,
                         FILE_SHARE_DELETE,
                         FILE_OPEN_IF,
                         FILE_SYNCHRONOUS_IO_NONALERT,
                         NULL,
                         0 ); 

  if ( !NT_SUCCESS( status ) )
  {
    return status;
  } 

  ZwQueryInformationFile( hFile,
                          &IoStatusBlock,
                          ( void * ) &FileSTDInfo,
                          sizeof( FileSTDInfo ),
                          FileStandardInformation );


  if ( !NT_SUCCESS( IoStatusBlock.Status ) )
  {
    return IoStatusBlock.Status;
  } 

  FilePOTInfo.CurrentByteOffset = FileSTDInfo.EndOfFile; 

  ZwSetInformationFile( hFile,
                        &IoStatusBlock,
                        ( void * ) &FilePOTInfo,
                        sizeof( FilePOTInfo ),
                        FilePositionInformation ); 


  if ( !NT_SUCCESS( IoStatusBlock.Status ) )
  {
    return IoStatusBlock.Status;
  }

  ZwWriteFile( hFile, 0, 0, 0, &IoStatusBlock, lpBuffer, uSize, NULL, NULL ); 

  ZwClose( hFile ); 

  return IoStatusBlock.Status;
}