用ZwQueryDirectoryObject遍历驱动对象

#include <ntifs.h>
#include <windef.h>
#include <Aux_klib.h> //source里面要有TARGETLIBS=$(DDK_LIB_PATH)\Aux_klib.lib

#define TAG 'tset' //test

/*
文件名：ZwEnumerateObject.C
功能：遍历驱动对象。

made by correy
made at 2015.05.19
*/

//\WRK-v1.2\public\internal\base\inc\zwapi.h
//https://msdn.microsoft.com/en-us/library/bb470238(v=vs.85).aspx
NTSYSAPI
NTSTATUS
NTAPI
ZwQueryDirectoryObject (
    __in HANDLE DirectoryHandle,
    __out_bcount_opt(Length) PVOID Buffer,
    __in ULONG Length,
    __in BOOLEAN ReturnSingleEntry,
    __in BOOLEAN RestartScan,
    __inout PULONG Context,
    __out_opt PULONG ReturnLength
    );

//\WRK-v1.2\public\sdk\inc\ntobapi.h
//https://msdn.microsoft.com/en-us/library/bb470238(v=vs.85).aspx
typedef struct _OBJECT_DIRECTORY_INFORMATION {
    UNICODE_STRING Name;
    UNICODE_STRING TypeName;
} OBJECT_DIRECTORY_INFORMATION, *POBJECT_DIRECTORY_INFORMATION;


DRIVER_UNLOAD Unload;
VOID Unload(__in PDRIVER_OBJECT DriverObject)
{  
   
}


NTSTATUS ZwEnumerateObject(IN UNICODE_STRING * directory)
    /*
    功能：枚举一个DirectoryObject下的对象。
    */
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    OBJECT_ATTRIBUTES ob;
    HANDLE FileHandle = 0;
    IO_STATUS_BLOCK  IoStatusBlock = {0};
    PVOID FileInformation = 0;
    ULONG Length = sizeof (FILE_DIRECTORY_INFORMATION);//这个数设置的太小会导致ZwQueryDirectoryFile蓝屏。  
    BOOLEAN           RestartScan;
    ULONG             Context = 0;
    ULONG             ReturnedLength;
    UNICODE_STRING driver = RTL_CONSTANT_STRING(L"Driver");

    InitializeObjectAttributes(&ob, directory, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, 0, 0);
    status = ZwOpenDirectoryObject(&FileHandle, GENERIC_READ | SYNCHRONIZE, &ob);
    if (!NT_SUCCESS (status))
    {
        KdPrint(("ZwOpenDirectoryObject fail with 0x%x.\n", status));
        return status;
    }

    Length = Length + 520;//为何加这个数字，请看ZwEnumerateFile1的说明。
    FileInformation = ExAllocatePoolWithTag(NonPagedPool, Length, TAG);
    if (FileInformation == NULL) {
        status = STATUS_UNSUCCESSFUL;
        DbgPrint("发生错误的文件为:%s, 代码行为:%d\n", __FILE__, __LINE__);
        ZwClose(FileHandle);
        return status;
    }
    RtlZeroMemory(FileInformation, Length);

    //RestartScan = FALSE;//为TRUE会导致死循环;
    //status = ZwQueryDirectoryObject( FileHandle, FileInformation, Length, TRUE, RestartScan, &Context, &ReturnedLength );
    //if (!NT_SUCCESS (status)) //此时也会得到数据。
    //{
    //    KdPrint(("ZwQueryDirectoryFile fail with 0x%x.\n", status));//STATUS_BUFFER_TOO_SMALL == C0000023
    //    ExFreePoolWithTag(FileInformation, TAG);
    //    ZwClose(FileHandle);
    //    return status;
    //}

    do
    {
        UNICODE_STRING FileName ={0};
        POBJECT_DIRECTORY_INFORMATION podi = 0;
        UNICODE_STRING FullName = {0};

        RestartScan = FALSE;//为TRUE会导致死循环;
        status = ZwQueryDirectoryObject( FileHandle, FileInformation, Length, TRUE, RestartScan, &Context, &ReturnedLength );
        if (status != STATUS_NO_MORE_FILES && status != STATUS_SUCCESS)
        {
            break;//这里好像没有走过。
        }

        podi = (POBJECT_DIRECTORY_INFORMATION)FileInformation;

        //不是驱动对象就放过。
        if (RtlCompareUnicodeString(&podi->TypeName, &driver, TRUE ) != 0 )
        {
            continue;
        }

        //申请要显示的内存，另一思路是格式化。
        FullName.MaximumLength = (USHORT)Length + directory->MaximumLength;
        FullName.Buffer = ExAllocatePoolWithTag(NonPagedPool, FullName.MaximumLength, TAG);
        if (FullName.Buffer == NULL) {            
            DbgPrint("发生错误的文件为:%s, 代码行为:%d\n", __FILE__, __LINE__);
            status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }
        RtlZeroMemory(FullName.Buffer, FullName.MaximumLength);

        RtlCopyUnicodeString(&FullName, directory);

        status = RtlAppendUnicodeToString(&FullName, L"\\");
        if (!NT_SUCCESS (status)) {
            KdPrint(("RtlAppendUnicodeToString fail with 0x%X!\n", status));
            ExFreePoolWithTag(FullName.Buffer, TAG);
            break;
        }

        status = RtlAppendUnicodeStringToString(&FullName, &podi->Name);
        if (!NT_SUCCESS (status)) {
            KdPrint(("RtlAppendUnicodeStringToString fail with 0x%X. 文件为:%s, 代码行为:%d\n", status, __FILE__, __LINE__));
            ExFreePoolWithTag(FullName.Buffer, TAG);
            break;
        }

        //KdPrint(("Name %wZ\n", &podi->Name));
        KdPrint(("Name %wZ\n", &FullName));

        ExFreePoolWithTag(FullName.Buffer, TAG);      

    } while (status != STATUS_NO_MORE_FILES);

    if (status = STATUS_NO_MORE_FILES)
    {
        status = STATUS_SUCCESS;
    }

    if (FileInformation)
    {
        ExFreePoolWithTag(FileInformation, TAG);
        FileInformation = NULL;
    }

    ZwClose(FileHandle);

    return status;
}


NTSTATUS ZwEnumerateDriverObject()
    /*
    功能：枚举系统的驱动对象。
    记得：有一个函数可以实现此功能，当时是热衷于枚举系统的模块，所以当时没有注意那个函数，现在有想不起了，一晃又几年过去了。
          记得devicetree就是通过枚举driver和FileSystem下的对象实现的，甚至是WINOBJ等类似的对象浏览器也是通过这实现的。
          看来ZwQueryDirectoryObject的功能应该不小，现在微软在应用层已经公开此函数，在内核使用此函数应该没有问题。
    说明：FileSystem下虽然还有子目录对象，但是下没有驱动，所以没有递归。可能原理也应该没有。
          这个有啥好处呢？因为不包括NT模块（就是那个几个NTOS*.EXE）和HAL.DLL(也包括几个类型的)，所以可以根据此驱动对象进一步获取分发函数，设备对象等操作。此操作只判断存在，余下的还要自己动手。
    */
{
    NTSTATUS status = STATUS_UNSUCCESSFUL;
    UNICODE_STRING directory  = RTL_CONSTANT_STRING(L"\\driver");
    UNICODE_STRING FileSystem  = RTL_CONSTANT_STRING(L"\\FileSystem");

    status = ZwEnumerateObject(&directory);
    if (!NT_SUCCESS (status))
    {
        return status;
    }

    status = ZwEnumerateObject(&FileSystem);

    return status;
}


#pragma INITCODE
DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry(__in struct _DRIVER_OBJECT * DriverObject, __in PUNICODE_STRING RegistryPath)
{
    KdBreakPoint();

    DriverObject->DriverUnload = Unload;

    return ZwEnumerateDriverObject();
}