自己调用NTDLL函数

一、概述

在DLL初始化的时候有时不能调用其它系统DLL的函数,以免导致问题,但有时候又必须要调用怎么办?一种办法就是自己直接调用NTDLL接口,这样肯定没有问题。

下面我写个自己调用Registry的封装类,用来代替原本系统注册表API。

 

二、申明NTDLL导出函数

在自己工程中需要调用NTDLL导出函数,可以通GetProcessAddr来获取函数地址再调用,也可以通过导入库的方式(这种需要ntdll.lib文件,在DDK中有),我这里采用第二种方式。

函数和类型声明:

NTDllDecl.h

#pragma once
#include <windows.h>
#include <winternl.h>

namespace NT {
    extern "C" {

// macro definition
#ifndef NTSTATUS
#define NTSTATUS    LONG                                
#endif

#ifndef NTAPI
#define NTAPI        __stdcall                                
#endif

#ifndef NTSYSAPI 
#define NTSYSAPI    __declspec(dllimport)
#endif

#ifndef STATUS_BUFFER_OVERFLOW
#define STATUS_BUFFER_OVERFLOW           ((NTSTATUS)0x80000005L)
#endif

typedef enum _KEY_INFORMATION_CLASS {
    KeyBasicInformation,
    KeyNodeInformation,
    KeyFullInformation,
    KeyNameInformation,
    KeyCachedInformation,
    KeyFlagsInformation,
    KeyVirtualizationInformation,
    KeyHandleTagsInformation,
    MaxKeyInfoClass  // MaxKeyInfoClass should always be the last enum
} KEY_INFORMATION_CLASS;

typedef enum _KEY_VALUE_INFORMATION_CLASS {
    KeyValueBasicInformation,
    KeyValueFullInformation,
    KeyValuePartialInformation,
    KeyValueFullInformationAlign64,
    KeyValuePartialInformationAlign64,
    MaxKeyValueInfoClass  // MaxKeyValueInfoClass should always be the last enum
} KEY_VALUE_INFORMATION_CLASS;

typedef struct _KEY_BASIC_INFORMATION {
    LARGE_INTEGER LastWriteTime;
    ULONG   TitleIndex;
    ULONG   NameLength;
    WCHAR   Name[1];            // Variable length string
} KEY_BASIC_INFORMATION, *PKEY_BASIC_INFORMATION;

typedef struct _KEY_NODE_INFORMATION {
    LARGE_INTEGER LastWriteTime;
    ULONG   TitleIndex;
    ULONG   ClassOffset;
    ULONG   ClassLength;
    ULONG   NameLength;
    WCHAR   Name[1];            // Variable length string
//          Class[1];           // Variable length string not declared
} KEY_NODE_INFORMATION, *PKEY_NODE_INFORMATION;

typedef struct _KEY_FULL_INFORMATION {
    LARGE_INTEGER LastWriteTime;
    ULONG   TitleIndex;
    ULONG   ClassOffset;
    ULONG   ClassLength;
    ULONG   SubKeys;
    ULONG   MaxNameLen;
    ULONG   MaxClassLen;
    ULONG   Values;
    ULONG   MaxValueNameLen;
    ULONG   MaxValueDataLen;
    WCHAR   Class[1];           // Variable length
} KEY_FULL_INFORMATION, *PKEY_FULL_INFORMATION;

typedef struct _KEY_VALUE_BASIC_INFORMATION {
    ULONG   TitleIndex;
    ULONG   Type;
    ULONG   NameLength;
    WCHAR   Name[1];            // Variable size
} KEY_VALUE_BASIC_INFORMATION, *PKEY_VALUE_BASIC_INFORMATION;

typedef struct _KEY_VALUE_PARTIAL_INFORMATION {
    ULONG   TitleIndex;
    ULONG   Type;
    ULONG   DataLength;
    UCHAR   Data[1];            // Variable size
} KEY_VALUE_PARTIAL_INFORMATION, *PKEY_VALUE_PARTIAL_INFORMATION;

// Registry
NTSYSAPI
NTSTATUS
NTAPI
ZwCreateKey(
    OUT PHANDLE KeyHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    IN ULONG TitleIndex,
    IN PUNICODE_STRING Class,
    IN ULONG CreateOptions,
    OUT PULONG Disposition
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwOpenKey(
    OUT PHANDLE KeyHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwQueryKey(
    IN HANDLE KeyHandle,
    IN KEY_INFORMATION_CLASS KeyInformationClass,
    OUT PVOID KeyInformation,
    IN ULONG KeyInformationLength,
    OUT PULONG ResultLength
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwSetValueKey(
    IN HANDLE KeyHandle,
    IN PUNICODE_STRING ValueName,
    IN ULONG TitleIndex,
    IN ULONG Type,
    IN PVOID Data,
    IN ULONG DataSize
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwQueryValueKey(
    IN HANDLE KeyHandle,
    IN PUNICODE_STRING ValueName,
    IN KEY_VALUE_INFORMATION_CLASS KeyValueInformationClass,
    OUT PVOID KeyValueInformation,
    IN ULONG KeyValueInformationLength,
    OUT PULONG ResultLength
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwEnumerateKey(
    IN HANDLE KeyHandle,
    IN ULONG Index,
    IN KEY_INFORMATION_CLASS KeyInformationClass,
    OUT PVOID KeyInformation,
    IN ULONG KeyInformationLength,
    OUT PULONG ResultLength
    );

NTSYSAPI
NTSTATUS
NTAPI
ZwClose(
    IN HANDLE Handle
    );

NTSYSAPI
NTSTATUS 
NTAPI
RtlFormatCurrentUserKeyPath( 
    OUT PUNICODE_STRING RegistryPath 
    );


    }
};

 

三、封装自己注册表函数

自己定义一个类用来封装注册表函数,我这里只封装了几个需要用的注册表函数(参考ReatOS和Windows代码),其余的可以后续再添加。

注册封装类声明:

NTRegistry.h

//********************************************************************  
//    文件名:        NTRegistry.h
//    文件描述:      Ntdll导出注册表函数定义头文件
//    作者:           
//    版本:          1.0
// 
//    修改历史:      
//    备注:          
//*********************************************************************
#pragma once

#ifndef STATIC
#define STATIC        static
#endif

#ifndef NTSTATUS
#define NTSTATUS    LONG
#endif


class NTRegistry
{
public:
    
    STATIC LONG RegCreateKey(
                HKEY hKey,
                LPCWSTR lpSubKey,
                DWORD Reserved,
                LPWSTR lpClass,
                DWORD dwOptions,
                REGSAM samDesired,
                LPSECURITY_ATTRIBUTES lpSecurityAttributes,
                PHKEY phkResult,
                LPDWORD lpdwDisposition);

    STATIC LONG RegOpenKey(
                HKEY hKey,
                LPCWSTR lpSubKey,
                DWORD ulOptions,
                REGSAM samDesired,
                PHKEY phkResult);

    STATIC LONG RegQueryInfoKey(
                HKEY hKey,
                LPWSTR lpClass,
                LPDWORD lpcbClass,
                LPDWORD lpReserved,
                LPDWORD lpcSubKeys,
                LPDWORD lpcbMaxSubKeyLen,
                LPDWORD lpcbMaxClassLen,
                LPDWORD lpcValues,
                LPDWORD lpcbMaxValueNameLen,
                LPDWORD lpcbMaxValueLen,
                LPDWORD lpcbSecurityDescriptor,
                PFILETIME lpftLastWriteTime
                );

    STATIC LONG RegSetValue( 
                HKEY hKey,
                LPCWSTR lpValueName,
                DWORD Reserved,
                DWORD dwType,
                CONST BYTE* lpData,
                DWORD cbData);

    STATIC LONG RegQueryValue(
                HKEY hKey,
                LPCWSTR lpValueName,
                LPDWORD lpReserved,
                LPDWORD lpType,
                LPBYTE lpData,
                LPDWORD lpcbData
                );

    STATIC LONG RegEnumKey(
                HKEY hKey,
                DWORD dwIndex,
                LPWSTR lpName,
                LPDWORD lpcbName,
                LPDWORD lpReserved,
                LPWSTR lpClass,
                LPDWORD lpcbClass,
                PFILETIME lpftLastWriteTime);

    STATIC LONG RegCloseKey(HKEY hKey);

private:

    STATIC NTSTATUS OpenCurrentUser(
                IN ACCESS_MASK DesiredAccess,
                OUT PHANDLE KeyHandle);

    STATIC NTSTATUS OpenLocalMachineKey(
                OUT PHANDLE KeyHandle);

    STATIC NTSTATUS OpenCurrentConfigKey (PHANDLE KeyHandle);

    STATIC NTSTATUS OpenUsersKey(PHANDLE KeyHandle);

    STATIC NTSTATUS OpenClassesRootKey(PHANDLE KeyHandle);

    STATIC NTSTATUS MapDefaultKey(
                OUT PHANDLE RealKey, 
                IN HKEY Key);
};

 

注册封装类定义:

NTRegistry.cpp

#include "StdAfx.h"
#include "NTRegistry.h"
#include "NTDllDecl.h"

#define ARGUMENT_PRESENT( arg )                                            \
    ((( PVOID ) arg ) != (( PVOID ) NULL ))

#define DEFAULT_KEY_NAME_SIZE                    128
#define DEFAULT_CLASS_SIZE                        128
#define DEFAULT_VALUE_SIZE                        128

#define REG_MAX_NAME_SIZE                        256
#define REG_MAX_DATA_SIZE                        2048

#ifndef STATUS_SUCCESS
#define STATUS_SUCCESS                          ((NTSTATUS)0x00000000L) 
#endif

#define RTL_CONSTANT_STRING(__SOURCE_STRING__)                            \
{                                                                        \
    sizeof(__SOURCE_STRING__) - sizeof((__SOURCE_STRING__)[0]),            \
    sizeof(__SOURCE_STRING__),                                            \
    (__SOURCE_STRING__)                                                    \
}

#define IsPredefKey(HKey)                                               \
    (((ULONG_PTR)(HKey) & 0xF0000000) == 0x80000000)

#define ClosePredefKey(Handle)                                          \
    if ((ULONG_PTR)Handle & 0x1) {                                      \
        NT::ZwClose(Handle);                                            \
    }

NTSTATUS
NTRegistry::OpenCurrentUser(
    IN ACCESS_MASK DesiredAccess,
    OUT PHANDLE KeyHandle)
{
    OBJECT_ATTRIBUTES ObjectAttributes;
    UNICODE_STRING KeyPath;
    NTSTATUS Status;

    /* Get the user key */
    Status = NT::RtlFormatCurrentUserKeyPath(&KeyPath);
    if (NT_SUCCESS(Status))
    {
        /* Initialize the attributes and open it */
        InitializeObjectAttributes(&ObjectAttributes,
                                   &KeyPath,
                                   OBJ_CASE_INSENSITIVE,
                                   NULL,
                                   NULL);
        Status = NT::ZwOpenKey(KeyHandle, DesiredAccess, &ObjectAttributes);

        /* Free the path and return success if it worked */
        RtlFreeUnicodeString(&KeyPath);        
    }

    return Status;
}

NTSTATUS 
NTRegistry::OpenLocalMachineKey(
    OUT PHANDLE KeyHandle)
{
    OBJECT_ATTRIBUTES Attributes;
    UNICODE_STRING KeyName = RTL_CONSTANT_STRING(L"\\Registry\\Machine");
    NTSTATUS Status;

    InitializeObjectAttributes(
        &Attributes,
        &KeyName,
        OBJ_CASE_INSENSITIVE,
        NULL,
        NULL);
    Status = NT::ZwOpenKey(KeyHandle, MAXIMUM_ALLOWED, &Attributes);

    return Status;
}

NTSTATUS
NTRegistry::OpenClassesRootKey(PHANDLE KeyHandle)
{
    OBJECT_ATTRIBUTES Attributes;
    UNICODE_STRING KeyName = RTL_CONSTANT_STRING(L"\\Registry\\Machine\\Software\\CLASSES");

    InitializeObjectAttributes(
        &Attributes,
        &KeyName,
        OBJ_CASE_INSENSITIVE,
        NULL,
        NULL);

    return NT::ZwOpenKey(KeyHandle,
                     MAXIMUM_ALLOWED,
                     &Attributes);
}

NTSTATUS
NTRegistry::OpenUsersKey(PHANDLE KeyHandle)
{
    OBJECT_ATTRIBUTES Attributes;
    UNICODE_STRING KeyName = RTL_CONSTANT_STRING(L"\\Registry\\User");
   
    InitializeObjectAttributes(&Attributes,
                               &KeyName,
                               OBJ_CASE_INSENSITIVE,
                               NULL,
                               NULL);
    return NT::ZwOpenKey(KeyHandle,
                     MAXIMUM_ALLOWED,
                     &Attributes);
}

NTSTATUS
NTRegistry::OpenCurrentConfigKey (PHANDLE KeyHandle)
{
    OBJECT_ATTRIBUTES Attributes;
    UNICODE_STRING KeyName =
        RTL_CONSTANT_STRING(L"\\Registry\\Machine\\System\\CurrentControlSet\\Hardware Profiles\\Current");

    InitializeObjectAttributes(&Attributes,
                               &KeyName,
                               OBJ_CASE_INSENSITIVE,
                               NULL,
                               NULL);
    return NT::ZwOpenKey(KeyHandle,
                     MAXIMUM_ALLOWED,
                     &Attributes);
}

NTSTATUS
NTRegistry::MapDefaultKey(
    OUT PHANDLE RealKey,
    IN HKEY Key)
{
    NTSTATUS Status;
    if (!IsPredefKey(Key))
    {
        *RealKey = (HANDLE)((ULONG_PTR)Key & ~0x1);
        return STATUS_SUCCESS;
    }

    switch ((LONG)Key)
    {
    case HKEY_CURRENT_USER:
        {
            Status = OpenCurrentUser(MAXIMUM_ALLOWED, RealKey);
        }
        break;

    case HKEY_LOCAL_MACHINE:
        {
            Status = OpenLocalMachineKey(RealKey);
        }
        break;

    case HKEY_CLASSES_ROOT:
        {
            Status = OpenClassesRootKey(RealKey);
        }
        break;

    case HKEY_USERS:
        {
            Status = OpenUsersKey(RealKey);
        }
        break;

    case HKEY_CURRENT_CONFIG:
        {
            Status = OpenCurrentConfigKey(RealKey);
        }
        break;

    default:
        Status = STATUS_INVALID_PARAMETER;
        break;
    }

    return Status;
}

LONG 
NTRegistry::RegCreateKey(
    HKEY hKey,
    LPCWSTR lpSubKey,
    DWORD Reserved,
    LPWSTR lpClass,
    DWORD dwOptions,
    REGSAM samDesired,
    LPSECURITY_ATTRIBUTES lpSecurityAttributes,
    PHKEY phkResult,
    LPDWORD lpdwDisposition)
{
    UNICODE_STRING SubKeyString;
    UNICODE_STRING ClassString;
    OBJECT_ATTRIBUTES ObjectAttributes;
    HANDLE ParentKey = NULL;
    NTSTATUS Status;

    if (lpSecurityAttributes && lpSecurityAttributes->nLength != sizeof(SECURITY_ATTRIBUTES))
    {
        return ERROR_INVALID_USER_BUFFER;
    }

    /* get the real parent key */
    Status = MapDefaultKey(&ParentKey,
                           hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    if (!phkResult)  
    {  
        return ERROR_INVALID_PARAMETER;  
    }

    RtlInitUnicodeString(&ClassString, lpClass);
    RtlInitUnicodeString(&SubKeyString, lpSubKey);

    InitializeObjectAttributes(
        &ObjectAttributes,
        &SubKeyString,
        OBJ_CASE_INSENSITIVE,
        (HANDLE)ParentKey,
        lpSecurityAttributes ? (PSECURITY_DESCRIPTOR)lpSecurityAttributes->lpSecurityDescriptor : NULL);

    Status = NT::ZwCreateKey(
        (PHANDLE)phkResult,
        samDesired,
        &ObjectAttributes,
        0,
        (lpClass == NULL) ? NULL : &ClassString,                    /* optional*/
        dwOptions,
        lpdwDisposition
        );         

    ClosePredefKey(ParentKey);    

    return RtlNtStatusToDosError( Status );
}

LONG NTRegistry::RegOpenKey(
    HKEY hKey,
    LPCWSTR lpSubKey,
    DWORD ulOptions,
    REGSAM samDesired,
    PHKEY phkResult)
{
    OBJECT_ATTRIBUTES ObjectAttributes;
    UNICODE_STRING SubKeyString;
    HANDLE KeyHandle;
    NTSTATUS Status;

    if (!phkResult)
    {
        return ERROR_INVALID_PARAMETER;
    }

    Status = MapDefaultKey(&KeyHandle, hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    if (lpSubKey != NULL)
        RtlInitUnicodeString(&SubKeyString, (LPWSTR)lpSubKey);
    else
        RtlInitUnicodeString(&SubKeyString, (LPWSTR)L"");
        
    InitializeObjectAttributes(
        &ObjectAttributes,
        &SubKeyString,
        OBJ_CASE_INSENSITIVE,
        KeyHandle,
        NULL
        );

    Status = NT::ZwOpenKey(
        (PHANDLE)phkResult,
        samDesired,
        &ObjectAttributes
        );

    ClosePredefKey(KeyHandle);

    return RtlNtStatusToDosError( Status );
}

LONG NTRegistry::RegQueryInfoKey(
    HKEY hKey,
    LPWSTR lpClass,
    LPDWORD lpcbClass,
    LPDWORD lpReserved,
    LPDWORD lpcSubKeys,
    LPDWORD lpcbMaxSubKeyLen,
    LPDWORD lpcbMaxClassLen,
    LPDWORD lpcValues,
    LPDWORD lpcbMaxValueNameLen,
    LPDWORD lpcbMaxValueLen,
    LPDWORD lpcbSecurityDescriptor,
    PFILETIME lpftLastWriteTime
    )
{
    NT::KEY_FULL_INFORMATION FullInfoBuffer;
    NT::PKEY_FULL_INFORMATION FullInfo;
    ULONG FullInfoSize;
    ULONG ClassLength = 0;
    HANDLE KeyHandle;
    NTSTATUS Status;
    ULONG Length;
    LONG ErrorCode = ERROR_SUCCESS;

    if ((lpClass) && (!lpcbClass))
    {
        return ERROR_INVALID_PARAMETER;
    }

    Status = MapDefaultKey(&KeyHandle,
                           hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    if (lpClass != NULL)
    {
        if (*lpcbClass > 0)
        {
            ClassLength = min(*lpcbClass - 1, REG_MAX_NAME_SIZE) * sizeof(WCHAR);
        }
        else
        {
            ClassLength = 0;
        }

        FullInfoSize = sizeof(NT::KEY_FULL_INFORMATION) + ((ClassLength + 3) & ~3);
        FullInfo = (NT::PKEY_FULL_INFORMATION)HeapAlloc(GetProcessHeap(),
                             0,
                             FullInfoSize);
        if (FullInfo == NULL)
        {
            ErrorCode = ERROR_OUTOFMEMORY;
            goto Cleanup;
        }

        FullInfo->ClassLength = ClassLength;
    }
    else
    {
        FullInfoSize = sizeof(NT::KEY_FULL_INFORMATION);
        FullInfo = &FullInfoBuffer;
        FullInfo->ClassLength = 0;
    }
    FullInfo->ClassOffset = FIELD_OFFSET(NT::KEY_FULL_INFORMATION, Class);

    Status = NT::ZwQueryKey(KeyHandle,
                        NT::KeyFullInformation,
                        FullInfo,
                        FullInfoSize,
                        &Length);
    if (!NT_SUCCESS(Status))
    {
        if (lpClass != NULL)
        {
            HeapFree(GetProcessHeap(),
                        0,
                        FullInfo);
        }

        ErrorCode = RtlNtStatusToDosError(Status);
        goto Cleanup;
    }

    if (lpcSubKeys != NULL)
    {
        *lpcSubKeys = FullInfo->SubKeys;
    }

    if (lpcbMaxSubKeyLen != NULL)
    {
        *lpcbMaxSubKeyLen = FullInfo->MaxNameLen / sizeof(WCHAR) + 1;
    }

    if (lpcbMaxClassLen != NULL)
    {
        *lpcbMaxClassLen = FullInfo->MaxClassLen / sizeof(WCHAR) + 1;
    }

    if (lpcValues != NULL)
    {
        *lpcValues = FullInfo->Values;
    }

    if (lpcbMaxValueNameLen != NULL)
    {
        *lpcbMaxValueNameLen = FullInfo->MaxValueNameLen / sizeof(WCHAR) + 1;
    }

    if (lpcbMaxValueLen != NULL)
    {
        *lpcbMaxValueLen = FullInfo->MaxValueDataLen;
    }

    if (lpftLastWriteTime != NULL)
    {
        lpftLastWriteTime->dwLowDateTime = FullInfo->LastWriteTime.u.LowPart;
        lpftLastWriteTime->dwHighDateTime = FullInfo->LastWriteTime.u.HighPart;
    }

    if (lpClass != NULL)
    {
        if (FullInfo->ClassLength > ClassLength)
        {
            ErrorCode = ERROR_BUFFER_OVERFLOW;
        }
        else
        {
            RtlCopyMemory(lpClass,
                            FullInfo->Class,
                            FullInfo->ClassLength);
            *lpcbClass = FullInfo->ClassLength / sizeof(WCHAR);
            lpClass[*lpcbClass] = 0;
        }

        HeapFree(GetProcessHeap(),
                    0,
                    FullInfo);
    }

Cleanup:

    ClosePredefKey(KeyHandle);

    return ErrorCode;
}

LONG NTRegistry::RegSetValue(
    HKEY hKey,
    LPCWSTR lpValueName,
    DWORD Reserved,
    DWORD dwType,
    CONST BYTE* lpData,
    DWORD cbData)
{
    UNICODE_STRING ValueName;
    PUNICODE_STRING pValueName;
    HANDLE KeyHandle;
    NTSTATUS Status;

    Status = MapDefaultKey(&KeyHandle,
                           hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    RtlInitUnicodeString(&ValueName, lpValueName);
    pValueName = &ValueName;

    if ((dwType == REG_SZ) || (dwType == REG_EXPAND_SZ) || (dwType == REG_MULTI_SZ) && (cbData != 0))
    {
        PWSTR pwsData = (PWSTR)lpData;

        if((pwsData[cbData / sizeof(WCHAR) - 1] != L'\0') &&
            (pwsData[cbData / sizeof(WCHAR)] == L'\0'))
        {
            /* Increment length if last character is not zero and next is zero */
            cbData += sizeof(WCHAR);
        }
    }

    Status = NT::ZwSetValueKey(
        hKey,
        &ValueName,
        0,
        dwType,
        (PVOID)lpData,
        cbData);

    ClosePredefKey(KeyHandle);

    return RtlNtStatusToDosError( Status );
}

LONG NTRegistry::RegQueryValue(
    HKEY hKey,
    LPCWSTR lpValueName,
    LPDWORD lpReserved,
    LPDWORD lpType,
    LPBYTE lpData,
    LPDWORD lpcbData
    )
{
    NTSTATUS Status;
    UNICODE_STRING ValueName;
    ULONG    BufferLength;
    NT::KEY_VALUE_INFORMATION_CLASS KeyValueInformationClass;
    PVOID    KeyValueInformation;
    ULONG    ResultLength;
    HANDLE    KeyHandle;
    BYTE    PrivateKeyValueInformation[ sizeof( NT::KEY_VALUE_PARTIAL_INFORMATION ) +
                                        DEFAULT_VALUE_SIZE ];

    //
    // Validate dependency between lpData and lpcbData parameters.
    //

    if( ARGUMENT_PRESENT( lpReserved ) ||
        (ARGUMENT_PRESENT( lpData ) && ( ! ARGUMENT_PRESENT( lpcbData )))) {
        return ERROR_INVALID_PARAMETER;
    }

    Status = MapDefaultKey(&KeyHandle,
                           hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    //
    // Convert the value name to a counted Unicode string.
    //

    RtlInitUnicodeString( &ValueName, lpValueName );
    
    //
    //  First we assume that the information we want will fit on
    //  PrivateKeyValueInformattion
    //

    KeyValueInformationClass = ( ARGUMENT_PRESENT( lpcbData ))?
                               NT::KeyValuePartialInformation :
                               NT::KeyValueBasicInformation;

    KeyValueInformation = PrivateKeyValueInformation;
    BufferLength = sizeof( PrivateKeyValueInformation );    

    Status = NT::ZwQueryValueKey(
                KeyHandle,
                &ValueName,
                KeyValueInformationClass,
                KeyValueInformation,
                BufferLength,
                &ResultLength);

    if( ( Status == STATUS_BUFFER_OVERFLOW ) &&
        ( !ARGUMENT_PRESENT( lpData ) ) ) 
    {
        
        Status = STATUS_SUCCESS;
    }

    if( Status == STATUS_BUFFER_OVERFLOW ) 
    {
        //
        //  The buffer defined in the stack wasn't big enough to hold
        //  the Value information.
        //  If the caller's buffer is big enough to hold the value data
        //  then allocate a new buffer, and call the NT API again.
        //
        if( ( ( KeyValueInformationClass == NT::KeyValuePartialInformation ) &&
              ( ARGUMENT_PRESENT( lpData ) ) &&
              ( *lpcbData >= (( NT::PKEY_VALUE_PARTIAL_INFORMATION )KeyValueInformation )->DataLength)) ) 
        {
            BufferLength = ResultLength;

            KeyValueInformation = HeapAlloc( GetProcessHeap( ), 0, BufferLength);
            //
            // If the memory allocation fails, return a Registry error.
            //

            if( ! KeyValueInformation ) 
            {
                return ERROR_OUTOFMEMORY;
            }

            //
            // Query for the necessary information about the supplied value.
            //

            Status = NT::ZwQueryValueKey( 
                            KeyHandle,
                            &ValueName,
                            KeyValueInformationClass,
                            KeyValueInformation,
                            BufferLength,
                            &ResultLength
                            );
        }
    }

    if( NT_SUCCESS( Status ) &&
        ARGUMENT_PRESENT( lpData ) ) 
    {
        //
        //  If requested, copy the value data
        //
        if( *lpcbData >= (( NT::PKEY_VALUE_PARTIAL_INFORMATION )
                            KeyValueInformation )->DataLength ) 
        {

            RtlMoveMemory( lpData, (( NT::PKEY_VALUE_PARTIAL_INFORMATION )KeyValueInformation )->Data,
                           (( NT::PKEY_VALUE_PARTIAL_INFORMATION )KeyValueInformation )->DataLength);

        } 
        else 
        {
            Status = STATUS_BUFFER_OVERFLOW;
        }
    }

    //
    // Certain information is returned on success or in the case of
    // NtQueryValueKey returning STATUS_BUFFER_OVERFLOW.  This information
    // is always available because we always pass the minimum size required for
    // the NtQueryValueKey API.
    //

    if( NT_SUCCESS( Status ) ||
        ( Status == STATUS_BUFFER_OVERFLOW ) ) {

        if( KeyValueInformationClass == NT::KeyValueBasicInformation ) {

            //
            // If requested, return the value type.
            //

            if( ARGUMENT_PRESENT( lpType )) {

                *lpType = (( NT::PKEY_VALUE_BASIC_INFORMATION )
                            KeyValueInformation )->Type;
            }

        } else {

            //
            // If requested, return the value type.
            //

            if( ARGUMENT_PRESENT( lpType )) {

                *lpType = (( NT::PKEY_VALUE_PARTIAL_INFORMATION )
                            KeyValueInformation )->Type;
            }

            //
            // Return the value data size
            //
            *lpcbData = (( NT::PKEY_VALUE_PARTIAL_INFORMATION )
                            KeyValueInformation )->DataLength;
        }
    }

    //
    // Transmit all of the data back to the client.
    //

    if( ARGUMENT_PRESENT( lpcbData ) ) {
        if( NT_SUCCESS( Status ) &&
            ARGUMENT_PRESENT( lpData ) ) {
            *lpcbData = (( NT::PKEY_VALUE_PARTIAL_INFORMATION )
                        KeyValueInformation )->DataLength;
        } else {
            //
            // The API failed, so make sure that no data is transmitted back
            // to the client. This ensures that the client stub will not
            // attempt to unmarshall data that doesn't exist.
            //

            *lpcbData = 0;
        }
    }

    //
    //  If memory was allocated, then free it
    //
    if( KeyValueInformation != PrivateKeyValueInformation ) {

        HeapFree( GetProcessHeap( ), 0, KeyValueInformation );
    }
    
    return RtlNtStatusToDosError( Status );
}

LONG NTRegistry::RegEnumKey(
    HKEY hKey,
    DWORD dwIndex,
    LPWSTR lpName,
    LPDWORD lpcName,
    LPDWORD lpReserved,
    LPWSTR lpClass,
    LPDWORD lpcClass,
    PFILETIME lpftLastWriteTime)
{
    NTSTATUS Status;
    NT::KEY_INFORMATION_CLASS  KeyInformationClass;
    UNICODE_STRING NameString;
    UNICODE_STRING ClassString;
    PVOID  KeyInformation;
    ULONG  BufferLength;
    ULONG  ResultLength;
    HANDLE KeyHandle;
    BYTE   PrivateKeyInformation[ sizeof( NT::KEY_NODE_INFORMATION ) +
                                        DEFAULT_KEY_NAME_SIZE +
                                        DEFAULT_CLASS_SIZE ];


    if( !ARGUMENT_PRESENT( lpName ) ||
        (ARGUMENT_PRESENT( lpClass ) && ( ! ARGUMENT_PRESENT( lpcClass )))) {
        return ERROR_INVALID_PARAMETER;
    }

    //
    // Use the supplied name string buffer as the buffer in a counted
    // Unicode string.
    //
    NameString.Length           = 0;
    if ((*lpcName << 1) > 0xFFFE) {
        NameString.MaximumLength    = ( USHORT ) 0xFFFE;
    } else {
        NameString.MaximumLength    = ( USHORT )( *lpcName << 1 );
    }    
    NameString.Buffer           = lpName;

    //
    // If supplied use the supplied name string buffer as the buffer in a
    // counted Unicode string.
    //
    if (ARGUMENT_PRESENT( lpClass ))
        RtlInitUnicodeString(&ClassString, (LPWSTR)lpClass);
    else
        RtlInitUnicodeString(&ClassString, (LPWSTR)L"");    

    //
    //  First we assume that the information we want will fit on
    //  PrivateKeyValueInformattion
    //

    KeyInformationClass = (ARGUMENT_PRESENT( lpClass )) ?  NT::KeyNodeInformation : NT::KeyBasicInformation;

    KeyInformation = PrivateKeyInformation;
    BufferLength = sizeof( PrivateKeyInformation );

    Status = MapDefaultKey(&KeyHandle,
                           hKey);
    if (!NT_SUCCESS(Status))
    {
        return RtlNtStatusToDosError(Status);
    }

    Status = NT::ZwEnumerateKey(
        KeyHandle, 
        dwIndex,
        KeyInformationClass,
        KeyInformation,
        BufferLength,
        &ResultLength);

    if( Status == STATUS_BUFFER_OVERFLOW ) {
        //
        //  The buffer defined in the stack wasn't big enough to hold
        //  the Key information.
        //  If the caller's buffer are big enough to hold the key name
        //  and key class, then allocate a new buffer, and call the
        //  NT API again.
        //
        if( ( ( KeyInformationClass == NT::KeyBasicInformation ) &&
              ( (ULONG)(NameString.MaximumLength) >= (( NT::PKEY_BASIC_INFORMATION )KeyInformation )->NameLength + sizeof(UNICODE_NULL)) ) ||
            ( ( KeyInformationClass == NT::KeyNodeInformation ) &&
              ( (ULONG)(NameString.MaximumLength) >= (( NT::PKEY_NODE_INFORMATION )KeyInformation )->NameLength + sizeof(UNICODE_NULL)) &&
              ( ARGUMENT_PRESENT( lpClass )) && 
              ( (ULONG)(ClassString.MaximumLength) >= (( NT::PKEY_NODE_INFORMATION )KeyInformation )->ClassLength + sizeof(UNICODE_NULL))
            ) ) 
        {
            BufferLength = ResultLength;

            KeyInformation = HeapAlloc( GetProcessHeap( ), 0, BufferLength);
            //
            // If the memory allocation fails, return a Registry error.
            //

            if( ! KeyInformation ) {
                return ERROR_OUTOFMEMORY;
            }

            //
            // Query for the necessary information about the supplied key.
            // This may or may not include the class depending on lpClass->Buffer
            // as determined above.
            //

            {
                Status = NT::ZwEnumerateKey(
                            KeyHandle, dwIndex,
                            KeyInformationClass,
                            KeyInformation,
                            BufferLength,
                            &ResultLength);
            }

        }
    }

    if( NT_SUCCESS( Status ) ) 
    {

        if( KeyInformationClass == NT::KeyBasicInformation ) 
        {
            //
            // Return the name length and the name of the key.
            // Note that the NUL byte is included so that RPC copies the
            // correct number of bytes. It is decremented on the client
            // side.
            //

            if( (ULONG)(NameString.MaximumLength) >=
                 (( NT::PKEY_BASIC_INFORMATION )KeyInformation )->NameLength + sizeof( UNICODE_NULL ) ) 
            {

                NameString.Length = ( USHORT )(( NT::PKEY_BASIC_INFORMATION )KeyInformation )->NameLength;

                RtlMoveMemory( NameString.Buffer, 
                    (( NT::PKEY_BASIC_INFORMATION )KeyInformation )->Name,  
                    NameString.Length);

                //
                // NUL terminate the value name.
                //

                NameString.Buffer[ NameString.Length >> 1 ] = UNICODE_NULL;
                if( ARGUMENT_PRESENT( lpcName )) {

                    *lpcName = NameString.Length >> 1;
                }

            } else {

                Status = STATUS_BUFFER_OVERFLOW;
            }

            //
            // If requested, return the last write time.
            //

            if( ARGUMENT_PRESENT( lpftLastWriteTime )) {

                *lpftLastWriteTime = *( PFILETIME )&(( NT::PKEY_BASIC_INFORMATION ) KeyInformation )->LastWriteTime;
            }

        } else {
            //
            // Return the name length and the name of the key.
            // Note that the NUL byte is included so that RPC copies the
            // correct number of bytes. It is decremented on the client
            // side.
            //

            if( ( (ULONG)(NameString.MaximumLength) >=
                  (( NT::PKEY_NODE_INFORMATION )
                   KeyInformation )->NameLength + sizeof( UNICODE_NULL ) ) &&
                ( (ULONG)(ClassString.MaximumLength) >=
                  (( NT::PKEY_NODE_INFORMATION )
                   KeyInformation )->ClassLength + sizeof( UNICODE_NULL) )
              ) {
                //
                //  Copy the key name
                //
                NameString.Length = ( USHORT )
                                 (( NT::PKEY_NODE_INFORMATION )
                                 KeyInformation )->NameLength;
                
                RtlMoveMemory( NameString.Buffer,
                               (( NT::PKEY_NODE_INFORMATION )KeyInformation )->Name,
                               NameString.Length
                             );

                //
                // NUL terminate the key name.
                //

                NameString.Buffer[ NameString.Length >> 1 ] = UNICODE_NULL;                
                if( ARGUMENT_PRESENT( lpcName )) {

                    *lpcName = NameString.Length >> 1;
                }


                //
                //  Copy the key class
                //

                ClassString.Length = (USHORT)
                    ((( NT::PKEY_NODE_INFORMATION ) KeyInformation )->ClassLength );

                RtlMoveMemory(
                    ClassString.Buffer,
                    ( PBYTE ) KeyInformation
                    + (( NT::PKEY_NODE_INFORMATION ) KeyInformation )->ClassOffset,
                    ClassString.Length
                    );

                //
                // NUL terminate the class.
                //

                ClassString.Buffer[ ClassString.Length >> 1 ] = UNICODE_NULL;
                if( ARGUMENT_PRESENT( lpcClass )) {

                    *lpcClass = ClassString.Length >> 1;
                }

            } else {
                Status = STATUS_BUFFER_OVERFLOW;
            }

            //
            // If requested, return the last write time.
            //

            if( ARGUMENT_PRESENT( lpftLastWriteTime )) {

                *lpftLastWriteTime
                = *( PFILETIME )
                &(( NT::PKEY_NODE_INFORMATION ) KeyInformation )
                ->LastWriteTime;
            }
        }
    }
    
    if( KeyInformation != PrivateKeyInformation ) {
        //
        // Free the buffer allocated.
        //

        HeapFree( GetProcessHeap( ), 0, KeyInformation );
    }

    return RtlNtStatusToDosError( Status );
}

LONG NTRegistry::RegCloseKey(HKEY hKey)
{
    NTSTATUS Status;
    /* don't close null handle or a pseudo handle */
    if ((!hKey) || (((ULONG_PTR)hKey & 0xF0000000) == 0x80000000))
    {
        return ERROR_INVALID_HANDLE;
    }

    Status = NT::ZwClose( hKey );

    return RtlNtStatusToDosError( Status );
}

 

四、项目设置

为了使用NTDLL导出函数,我们这里还需要用到ntdll.lib,需要在工程属性中添加导入ntdll.lib,只需要把ntdll.lib从DDK中拷贝出来放到你自己第三方库位置。

 

五、测试程序

测试了个获取Adobe安装目录的程序

NtFunctionTest.cpp

#include "stdafx.h"
#include <gtest/gtest.h>
#include <string>
#include <NT/NtRegistry.h>

TEST(GsKernelTest, tRegistryKey)
{

    HKEY hKey = NULL;
    LONG RetVal    = NTRegistry::RegOpenKey(
                    HKEY_LOCAL_MACHINE, 
                    L"SOFTWARE\\Adobe\\Acrobat Reader",         
                    REG_OPTION_NON_VOLATILE,
                    KEY_READ,        
                    &hKey);
    ASSERT_TRUE(RetVal==ERROR_SUCCESS);
    if ( ERROR_SUCCESS != RetVal )
    {
        return;
    }

    DWORD cSubKeys = 0;
    RetVal = NTRegistry::RegQueryInfoKey(hKey, NULL, NULL, NULL, &cSubKeys, NULL, NULL, NULL, NULL, NULL, NULL, NULL );
    ASSERT_TRUE(RetVal==ERROR_SUCCESS);
    if ( ERROR_SUCCESS != RetVal )
    {
        NTRegistry::RegCloseKey(hKey);
        return;
    }

    TCHAR szKeyName[MAX_PATH];
    for ( DWORD i=0; i < cSubKeys; ++i )
    {
        DWORD dwKeyNameLength = sizeof(szKeyName)/sizeof(szKeyName[0]);
        RetVal = NTRegistry::RegEnumKey( hKey, i, szKeyName, &dwKeyNameLength, NULL, NULL, NULL,NULL);
        if ( ERROR_SUCCESS != RetVal )
        {
            continue;
        }

        double fVersion = _tstof(szKeyName);
        DWORD dwVesion = (DWORD)(fVersion*10);

        if ( dwVesion > 120 || dwVesion < 60  )
        {
            continue;
        }        

        /// 拷贝文件
        TCHAR szAdobePath[MAX_PATH];
        _tcscat_s(szKeyName,_countof(szKeyName), TEXT("\\InstallPath"));
        HKEY hChildKey = NULL;
        RetVal = NTRegistry::RegOpenKey(hKey, szKeyName, 0, KEY_READ, &hChildKey);
        if ( ERROR_SUCCESS != RetVal )
        {
            continue;
        }

        DWORD cKey = sizeof(szAdobePath);
        RetVal = NTRegistry::RegQueryValue(hChildKey, NULL, 0, NULL, (LPBYTE)szAdobePath, &cKey );
        if ( ERROR_SUCCESS == RetVal )
        {
            wcout << L"InstallPath: " << szAdobePath << endl;            
        }
        
        NTRegistry::RegCloseKey(hChildKey);

    }

    NTRegistry::RegCloseKey(hKey);
}

我的测试代码都是放到测试工程中的,采用的是gTest框架。

不采用导入ntdll.lib库方式的我后面再补充,只要通过GetProcessAddr获取函数地址就可以了。

 

Book16  看书、学习、写代码
posted @ 2015-09-25 16:20  Quincy  阅读(5443)  评论(0编辑  收藏  举报