内核中的注册表操作

#include <ntifs.h>
#include <windef.h>

/*
Zw层次的注册表操作很简单。
但是做与不做还是有点区别的。

made by correy
made at 2014.07.24
homepage:http://correy.webs.com
*/

NTSTATUS ZwEnumerateKeyEx(IN UNICODE_STRING * Name)
    /*
    显示一个注册表的键下的：子键，名字，类型，数据。
    注意：没有递归显示。
    */
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    OBJECT_ATTRIBUTES        ObjectAttributes;
    HANDLE  KeyHandle;
    PKEY_FULL_INFORMATION pfi;    
    ULONG  ResultLength;
    ULONG i = 0;

    InitializeObjectAttributes(&ObjectAttributes, Name, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
    Status = ZwOpenKey(&KeyHandle, KEY_READ, &ObjectAttributes);
    if( !NT_SUCCESS( Status ) )
    {
        return Status;
    }

    /*
    注意ZwQueryKey的第一个参数。
    The KeyHandle passed to ZwQueryKey must have been opened with KEY_QUERY_VALUE access. 
    This is accomplished by passing KEY_QUERY_VALUE, KEY_READ, or KEY_ALL_ACCESS as the DesiredAccess parameter to ZwCreateKey or ZwOpenKey.
    */

    // 第一次调用是为了获取需要的长度
    Status = ZwQueryKey(KeyHandle, KeyFullInformation, NULL, 0, &ResultLength);
    if( !NT_SUCCESS( Status ) )
    {
        if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
        {
            //在下面申请内存。
        }
        else
        {
            ZwClose(KeyHandle);
            return Status;
        }
    }

    //ResultLength += MAX_PATH ;
    //ResultLength *= 2;//多申请一半。
    pfi = (PKEY_FULL_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength);
    if (pfi == NULL)
    {
        //If ExAllocatePool returns NULL, the caller should return the NTSTATUS value STATUS_INSUFFICIENT_RESOURCES or should delay processing to another point in time.
        Status = STATUS_INSUFFICIENT_RESOURCES;
        ZwClose(KeyHandle);
        return Status;
    }

    // 第二次调用是为了获取数据
    Status = ZwQueryKey(KeyHandle, KeyFullInformation, pfi, ResultLength, &ResultLength);//少了赋值。这等低级的错误。
    if( !NT_SUCCESS( Status ) )
    {
        ExFreePool(pfi);
        ZwClose(KeyHandle);
        return Status;
    }

    //枚举子键。
    for (i = 0; i < pfi->SubKeys; i++)
    {     
        PKEY_BASIC_INFORMATION pbi;
        UNICODE_STRING us;

        // 获取第i个子项的长度
        Status = ZwEnumerateKey(KeyHandle, i, KeyBasicInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                break;
            }
        }

        pbi = (PKEY_BASIC_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pbi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }

        // 获取第i个子项的数据
        Status = ZwEnumerateKey(KeyHandle, i, KeyBasicInformation, pbi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pbi);
            break;
        }

        us.Buffer = pbi->Name;
        us.Length = (USHORT)pbi->NameLength;
        us.MaximumLength = us.Length;
        
        DbgPrint("subkey:%wZ\n", &us); 

        /*
        在这里组合字符串，可以考虑递归枚举。
        */

        ExFreePool(pbi);// 释放内存
    }

    //枚举名字，类型，数据。
    for (i = 0; i < pfi->Values; i++) //可以考虑用ZwQueryValueKey获取数量。MSDN关于这个成员的解释是：The number of value entries for this key.
    {     
        PKEY_VALUE_BASIC_INFORMATION pkvbi;
        UNICODE_STRING us;
        PKEY_VALUE_PARTIAL_INFORMATION pkvpi;
        UNICODE_STRING data;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        // 获取名字及类型。
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValueBasicInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                break;
            }
        }
        pkvbi = (PKEY_VALUE_BASIC_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pkvbi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValueBasicInformation, pkvbi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pkvbi);
            break;
        }

        us.Buffer = pkvbi->Name;
        us.Length = (USHORT)pkvbi->NameLength;
        us.MaximumLength = us.Length;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
        // 获取数据
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValuePartialInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                ExFreePool(pkvbi);
                break;
            }
        }
        pkvpi = (PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pkvpi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            ExFreePool(pkvbi);
            break;
        }
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValuePartialInformation, pkvpi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pkvpi);
            ExFreePool(pkvbi);
            break;
        }

        data.Buffer = (PWCH)pkvpi->Data;//有的数据可能无法显示。
        data.Length = (USHORT)pkvpi->DataLength;
        data.MaximumLength = data.Length;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        DbgPrint("name:%wZ,type:%d,data:%wZ\n", &us, pkvbi->Type, &data); 

        ExFreePool(pkvbi);// 释放内存
        ExFreePool(pkvpi);
    }

    ExFreePool(pfi);
    ZwClose(KeyHandle);

    return Status;
}


DRIVER_UNLOAD Unload;
VOID Unload(__in PDRIVER_OBJECT DriverObject)
{

}


DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry( __in struct _DRIVER_OBJECT  * DriverObject, __in PUNICODE_STRING  RegistryPath)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL; 
    UNICODE_STRING test = RTL_CONSTANT_STRING(L"\\REGISTRY\\MACHINE\\SYSTEM\\CurrentControlSet\\Control");//\\Session Manager
    
    KdBreakPoint();

    DriverObject->DriverUnload = Unload;  

    status = ZwEnumerateKeyEx(&test);
    if( !NT_SUCCESS( status ) )
    {
        DbgPrint("ZwEnumerateKeyEx fail with 0x%x\n", status); 
    }   

    return status;//STATUS_SUCCESS
} 

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include <ntifs.h>
#include <windef.h>

#define TAG 'tset' //test


NTSTATUS ZwCopyKey(IN UNICODE_STRING * Name, IN UNICODE_STRING * Name2)
    /*
    复制一个注册表的键下的：子键，名字，类型，数据。
    注意：
    1.没有递归复制。
    2.没有复制（安全）属性。
    3.没有对参数的有效性进行检查。字符串的末尾不要带L'\\'.
    4.确认使用前这两个路径是存在的。
    5.更多的缺陷，请你补充纠正。更多的功能等待你的发挥。
    */
{
    NTSTATUS Status = STATUS_UNSUCCESSFUL;
    OBJECT_ATTRIBUTES        ObjectAttributes;
    HANDLE  KeyHandle;
    HANDLE  KeyHandle3;
    PKEY_FULL_INFORMATION pfi;    
    ULONG  ResultLength;
    ULONG i = 0;

    InitializeObjectAttributes(&ObjectAttributes, Name, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
    Status = ZwOpenKey(&KeyHandle, KEY_ALL_ACCESS, &ObjectAttributes);
    if( !NT_SUCCESS( Status ) )
    {
        return Status;
    }

    /*
    注意ZwQueryKey的第一个参数。
    The KeyHandle passed to ZwQueryKey must have been opened with KEY_QUERY_VALUE access. 
    This is accomplished by passing KEY_QUERY_VALUE, KEY_READ, or KEY_ALL_ACCESS as the DesiredAccess parameter to ZwCreateKey or ZwOpenKey.
    */

    // 第一次调用是为了获取需要的长度
    Status = ZwQueryKey(KeyHandle, KeyFullInformation, NULL, 0, &ResultLength);
    if( !NT_SUCCESS( Status ) )
    {
        if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
        {
            //在下面申请内存。
        }
        else
        {
            ZwClose(KeyHandle);
            return Status;
        }
    }

    //ResultLength += MAX_PATH ;
    //ResultLength *= 2;//多申请一半。
    pfi = (PKEY_FULL_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength);
    if (pfi == NULL)
    {
        //If ExAllocatePool returns NULL, the caller should return the NTSTATUS value STATUS_INSUFFICIENT_RESOURCES or should delay processing to another point in time.
        Status = STATUS_INSUFFICIENT_RESOURCES;
        ZwClose(KeyHandle);
        return Status;
    }

    // 第二次调用是为了获取数据
    Status = ZwQueryKey(KeyHandle, KeyFullInformation, pfi, ResultLength, &ResultLength);//少了赋值。这等低级的错误。
    if( !NT_SUCCESS( Status ) )
    {
        ExFreePool(pfi);
        ZwClose(KeyHandle);
        return Status;
    }

    //枚举子键。
    for (i = 0; i < pfi->SubKeys; i++)
    {     
        PKEY_BASIC_INFORMATION pbi;
        UNICODE_STRING us;
        UNICODE_STRING new_key ;
        OBJECT_ATTRIBUTES ob;
        HANDLE KeyHandle2 = 0;
        PUNICODE_STRING Class = NULL;
        ULONG  Disposition;

        // 获取第i个子项的长度
        Status = ZwEnumerateKey(KeyHandle, i, KeyBasicInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                break;
            }
        }

        pbi = (PKEY_BASIC_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pbi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }

        // 获取第i个子项的数据
        Status = ZwEnumerateKey(KeyHandle, i, KeyBasicInformation, pbi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pbi);
            break;
        }

        us.Buffer = pbi->Name;
        us.Length = (USHORT)pbi->NameLength;
        us.MaximumLength = us.Length;
        
        DbgPrint("subkey:%wZ\n", &us); 

        //开始新建。

        new_key.Buffer = (wchar_t *)ExAllocatePoolWithTag(NonPagedPool, MAX_PATH, TAG);
        if (new_key.Buffer == NULL) { 
            ExFreePool(pbi);
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }
        RtlZeroMemory(new_key.Buffer, MAX_PATH);
        RtlInitEmptyUnicodeString(&new_key, new_key.Buffer,MAX_PATH);

        RtlCopyUnicodeString(&new_key,Name2); 

        Status = RtlAppendUnicodeToString(&new_key, L"\\");
        if (!NT_SUCCESS (Status)) {
            ExFreePool(new_key.Buffer);
            ExFreePool(pbi);
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }         

        Status = RtlAppendUnicodeStringToString(&new_key, &us);
        if (!NT_SUCCESS (Status)) {
            ExFreePool(new_key.Buffer);
            ExFreePool(pbi);
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        } 

        InitializeObjectAttributes(&ob, &new_key, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, 0, 0);
        Status = ZwCreateKey(&KeyHandle2, KEY_WRITE, &ob, 0, Class, REG_OPTION_NON_VOLATILE, &Disposition);//KEY_ALL_ACCESS KEY_READ
        if (!NT_SUCCESS (Status)) 
        {
            //如果子键已经存在，返回正确。
            ExFreePool(new_key.Buffer);
            ExFreePool(pbi);
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }
        else
        {
            if (KeyHandle2)//断言FileHandle不等于0也不是无效的句柄。
            {
                Status = ZwClose(KeyHandle2);
                if (!NT_SUCCESS (Status)) 
                {
                    KdPrint(("ZwClose fail with 0x%x.\n", Status));
                }
            }
        }

        /*
        在这里组合字符串，可以考虑递归。
        */

        ExFreePool(pbi);// 释放内存
        ExFreePool(new_key.Buffer);
    }

    //处理上面失败的情况。主要是for 循环。
    if( !NT_SUCCESS( Status ) )
    {
        ExFreePool(pfi);
        ZwClose(KeyHandle);
        return Status;
    }

    InitializeObjectAttributes(&ObjectAttributes, Name2, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);
    Status = ZwOpenKey(&KeyHandle3, KEY_ALL_ACCESS, &ObjectAttributes);
    if( !NT_SUCCESS( Status ) )
    {
        ExFreePool(pfi);
        ZwClose(KeyHandle);
        return Status;
    }

    //枚举名字，类型，数据。
    for (i = 0; i < pfi->Values; i++) //可以考虑用ZwQueryValueKey获取数量。MSDN关于这个成员的解释是：The number of value entries for this key.
    {     
        PKEY_VALUE_BASIC_INFORMATION pkvbi;
        UNICODE_STRING us;
        PKEY_VALUE_PARTIAL_INFORMATION pkvpi;
        UNICODE_STRING data;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        // 获取名字及类型。
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValueBasicInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                break;
            }
        }
        pkvbi = (PKEY_VALUE_BASIC_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pkvbi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValueBasicInformation, pkvbi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pkvbi);
            break;
        }

        us.Buffer = pkvbi->Name;
        us.Length = (USHORT)pkvbi->NameLength;
        us.MaximumLength = us.Length;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
        // 获取数据
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValuePartialInformation, NULL, 0, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            if (Status == STATUS_BUFFER_TOO_SMALL  || Status == STATUS_BUFFER_OVERFLOW) //STATUS_BUFFER_OVERFLOW这个情况应该不会发生在这种情况下。
            {
                //在下面申请内存。
            }
            else
            {
                ExFreePool(pkvbi);
                break;
            }
        }
        pkvpi = (PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePool(NonPagedPool, ResultLength); 
        if (pkvpi == NULL)
        {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            ExFreePool(pkvbi);
            break;
        }
        Status = ZwEnumerateValueKey(KeyHandle, i, KeyValuePartialInformation, pkvpi, ResultLength, &ResultLength);
        if( !NT_SUCCESS( Status ) )
        {
            ExFreePool(pkvpi);
            ExFreePool(pkvbi);
            break;
        }

        data.Buffer = (PWCH)pkvpi->Data;//有的数据可能无法显示。
        data.Length = (USHORT)pkvpi->DataLength;
        data.MaximumLength = data.Length;

        //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

        Status = ZwSetValueKey(KeyHandle3, &us, 0, pkvbi->Type, data.Buffer, data.Length);
        if( !NT_SUCCESS( Status ) ) //如果句柄的权限是KEY_READ这里成功，但是实际是没有成功的。
        {
            ExFreePool(pkvpi);
            ExFreePool(pkvbi);
            break;
        }

        DbgPrint("name:%wZ,type:%d,data:%wZ\n", &us, pkvbi->Type, &data); 

        ExFreePool(pkvbi);// 释放内存
        ExFreePool(pkvpi);
    }

    ExFreePool(pfi);
    ZwClose(KeyHandle);
    ZwClose(KeyHandle3);

    return Status;
}


DRIVER_UNLOAD Unload;
VOID Unload(__in PDRIVER_OBJECT DriverObject)
{

}


DRIVER_INITIALIZE DriverEntry;
NTSTATUS DriverEntry( __in struct _DRIVER_OBJECT  * DriverObject, __in PUNICODE_STRING  RegistryPath)
{
    NTSTATUS status = STATUS_UNSUCCESSFUL; 
    UNICODE_STRING test = RTL_CONSTANT_STRING(L"\\REGISTRY\\MACHINE\\SYSTEM\\CurrentControlSet\\Control\\Session Manager");
    UNICODE_STRING test2 = RTL_CONSTANT_STRING(L"\\REGISTRY\\MACHINE\\SYSTEM\\CurrentControlSet\\Control\\Session Manager2");
    
    KdBreakPoint();

    DriverObject->DriverUnload = Unload;  

    status = ZwCopyKey(&test, &test2);
    if( !NT_SUCCESS( status ) )
    {
        DbgPrint("ZwEnumerateKeyEx fail with 0x%x\n", status); 
    }   

    return status;//STATUS_SUCCESS
}