// Registry.c
//
// 
// Requires DDK Only
// File created on 2/2/2005
//

#include "pch.h"
#ifdef CPR_WMI_TRACE
#include "Registry.tmh"
#endif

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CprRegQueryValueKey
//      Queries the value of a value key in the registry
//
//  Arguments:
//      IN  RegKeyHandle
//              Handle to the root key
//
//      IN  SubKeyName
//              Optional subkey path string
//
//      IN  ValueName
//              Value name string
//
//      IN OUT RegValuePtr
//              Pointer to buffer that will contain the
//              registry value.
//
//  Return Value:
//      NTSTATUS
//
NTSTATUS CprRegQueryValueKey(
    IN  HANDLE      RegKeyHandle,
    IN  PWSTR       SubKeyName,
    IN  PWSTR       ValueName,
	IN OUT PVOID	RegValuePtr
    )
{
    NTSTATUS                        status;
    PKEY_VALUE_PARTIAL_INFORMATION  buffer;
    ULONG                           length;
    UNICODE_STRING                  regPath;
    UNICODE_STRING                  name;
    OBJECT_ATTRIBUTES               objAttributes;
    HANDLE                          hReg;
    BOOLEAN                         bFreeHandle;

    // Callers of ZwQueryValueKey must be at PASSIVE_LEVEL IRQL
    ASSERT(KeGetCurrentIrql() == PASSIVE_LEVEL);

    //ASSERT(RegKeyHandle != NULL);
    ASSERT(ValueName != NULL);
	ASSERT(RegValuePtr != NULL);
    
    buffer = NULL;
    bFreeHandle = FALSE;

    do
    {
        // Check for subkey path
        if (SubKeyName != NULL)
        {
            // Open a new handle
            RtlInitUnicodeString(&regPath, SubKeyName);

            // Initialize a new object attributes
            InitializeObjectAttributes(
                &objAttributes,
                &regPath,
                OBJ_KERNEL_HANDLE | OBJ_CASE_INSENSITIVE,
                RegKeyHandle,
                NULL
                );

            status = ZwOpenKey(&hReg, KEY_ALL_ACCESS, &objAttributes);

            if (!NT_SUCCESS(status))
            {
                CprDebugPrint1(NULL, DBG_PNP, DBG_WARN, __FUNCTION__ ": ZwOpenKey failed %x", status);
                break;
            }

            // Indicate that we need to free a handle here
            bFreeHandle = TRUE;
        }
        else
        {
            hReg = RegKeyHandle;
        }

        RtlInitUnicodeString(&name, ValueName);

        status = ZwQueryValueKey(
                    hReg,
                    &name,
                    KeyValuePartialInformation,
                    NULL,
                    0,
                    &length
                    );
    
        if ((status != STATUS_BUFFER_TOO_SMALL) && (status != STATUS_BUFFER_OVERFLOW))
        {
            CprDebugPrint1(NULL, DBG_PNP, DBG_WARN, __FUNCTION__ ": ZwQueryValueKey failed %x", status);
            break;
        }

        buffer = 
            (PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, length, CPR_POOL_TAG_REG_KEY);

        if (buffer == NULL)
        {
            status = STATUS_INSUFFICIENT_RESOURCES;
            break;
        }

        status = ZwQueryValueKey(
                    hReg,
                    &name,
                    KeyValuePartialInformation,
                    buffer,
                    length,
                    &length
                    );

        if (!NT_SUCCESS(status))
        {
            CprDebugPrint1(NULL, DBG_PNP, DBG_WARN, __FUNCTION__ ": ZwQueryValueKey failed %x", status);
            break;
        }
    }
    while (FALSE);

    // Allocate a buffer to return
    if (NT_SUCCESS(status))
    {
		char *tmpPtr;

        // Zero terminate strings just for ease of handling
        if ((buffer->Type == REG_EXPAND_SZ) ||
            (buffer->Type == REG_MULTI_SZ) ||
            (buffer->Type == REG_SZ))
        {
            // Allocate buffer
            tmpPtr = ExAllocatePoolWithTag(
                                        PagedPool, 
                                        buffer->DataLength + sizeof(WCHAR), 
                                        CPR_POOL_TAG_REG_QUERY
                                        );

			if (tmpPtr == NULL)
			{
				status = STATUS_INSUFFICIENT_RESOURCES;
			}
			else
			{
	            RtlZeroMemory(tmpPtr, buffer->DataLength + sizeof(WCHAR));
				// Copy the registry data to the pointer in the return buffer
				RtlCopyMemory(tmpPtr, (PVOID)buffer->Data, buffer->DataLength);
				*(char**)RegValuePtr = tmpPtr;
			}
		}
		else
		{
			// Copy the registry data to the return buffer
			RtlCopyMemory(RegValuePtr, (PVOID)buffer->Data, buffer->DataLength);
		}
    }

    if (buffer != NULL)
    {
        // Free our allocated memory
        ExFreePool(buffer);
    }

    if (bFreeHandle)
    {
        // Close our reg key handle
        ZwClose(hReg);
    }

    return status;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CprRegEnumerateKeys
//      Enumerates and print names of subkeys using a given registry key handle.
//
//  Arguments:
//      IN  RegKeyHandle
//              Handle to root key
//
//  Return Value:
//      none
//
VOID CprRegEnumerateKeys(
    IN  HANDLE RegKeyHandle
    )
{
    NTSTATUS                status;
    ULONG                   index;
    PKEY_BASIC_INFORMATION  regBuffer;
    PWCHAR                  nameBuffer;
    ULONG                   length;

    status = STATUS_SUCCESS;
    index = 0;
    regBuffer = NULL;
    nameBuffer = NULL;

    while (status != STATUS_NO_MORE_ENTRIES)
    {
        // Get the buffer size necessary
        status = ZwEnumerateKey(
                    RegKeyHandle,
                    index,
                    KeyBasicInformation,
                    NULL,
                    0,
                    &length
                    );

        if ((status != STATUS_BUFFER_TOO_SMALL) && (status != STATUS_BUFFER_OVERFLOW))
        {
            if (status != STATUS_NO_MORE_ENTRIES)
            {
                CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateKey failed %x", status);
            }
            else
            {
                CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": Enumerated %d keys", index);
            }

            break;
        }

        regBuffer = 
            (PKEY_BASIC_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, length, CPR_POOL_TAG_ENUM_KEYS);

        if (regBuffer == NULL)
        {
            continue;
        }

        // Now actually attempt to get subkey info
        status = ZwEnumerateKey(
                    RegKeyHandle,
                    index,
                    KeyBasicInformation,
                    regBuffer,
                    length,
                    &length
                    );

        if (!NT_SUCCESS(status))
        {
            CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateKey failed %x", status);

            // Free our temporary storage
            ExFreePool(regBuffer);

            continue;
        }

        // Allocate a buffer for the display name
        nameBuffer = (PWCHAR)ExAllocatePoolWithTag(
                                    PagedPool, 
                                    regBuffer->NameLength + sizeof(WCHAR), 
                                    CPR_POOL_TAG_DISPLAY_NAME
                                    );

        if (nameBuffer == NULL)
        {
            // Free our temporary storage
            ExFreePool(regBuffer);

            continue;
        }

        // NULL terminate the string
        RtlZeroMemory(nameBuffer, regBuffer->NameLength + sizeof(WCHAR));

        // Copy the name over
        RtlCopyMemory(nameBuffer, regBuffer->Name, regBuffer->NameLength);
        
        CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateKey returned %S", nameBuffer);

        // Free both buffers
        ExFreePool(regBuffer);
        ExFreePool(nameBuffer);

        // Increment our index
        ++index;
    }

    return;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
//  CprRegEnumerateValueKeys
//      Enumerates and print names of sub value keys using a given registry key handle.
//
//  Arguments:
//      IN  RegKeyHandle
//              Handle to root key
//
//  Return Value:
//      none
//
VOID CprRegEnumerateValueKeys(
    IN  HANDLE RegKeyHandle
    )
{
    NTSTATUS                        status;
    ULONG                           index;
    PKEY_VALUE_BASIC_INFORMATION    regBuffer;
    PWCHAR                          nameBuffer;
    ULONG                           length;

    status = STATUS_SUCCESS;
    index = 0;
    regBuffer = NULL;
    nameBuffer = NULL;

    while (status != STATUS_NO_MORE_ENTRIES)
    {
        // Get the buffer size necessary
        status = ZwEnumerateValueKey(
                    RegKeyHandle,
                    index,
                    KeyValueBasicInformation,
                    NULL,
                    0,
                    &length
                    );

        if ((status != STATUS_BUFFER_TOO_SMALL) && (status != STATUS_BUFFER_OVERFLOW))
        {
            if (status != STATUS_NO_MORE_ENTRIES)
            {
                CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateValueKey failed %x", status);
            }
            else
            {
                CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": Enumerated %d value keys", index);
            }

            break;
        }

        regBuffer = 
            (PKEY_VALUE_BASIC_INFORMATION)ExAllocatePoolWithTag(NonPagedPool, length, CPR_POOL_TAG_REG_BUF);

        if (regBuffer == NULL)
        {
            continue;
        }

        // Now actually attempt to get subkey info
        status = ZwEnumerateValueKey(
                    RegKeyHandle,
                    index,
                    KeyValueBasicInformation,
                    regBuffer,
                    length,
                    &length
                    );

        if (!NT_SUCCESS(status))
        {
            CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateValueKey failed %x", status);

            // Free our temporary storage
            ExFreePool(regBuffer);

            continue;
        }

        // Allocate a buffer for the display name
        nameBuffer = (PWCHAR)ExAllocatePoolWithTag(
                                    PagedPool, 
                                    regBuffer->NameLength + sizeof(WCHAR), 
                                    CPR_POOL_TAG_NAME_BUF
                                    );

        if (nameBuffer == NULL)
        {
            // Free our temporary storage
            ExFreePool(regBuffer);

            continue;
        }

        // NULL terminate the string
        RtlZeroMemory(nameBuffer, regBuffer->NameLength + sizeof(WCHAR));

        // Copy the name over
        RtlCopyMemory(nameBuffer, regBuffer->Name, regBuffer->NameLength);
        
        CprDebugPrint1(NULL, DBG_PNP, DBG_INFO, __FUNCTION__ ": ZwEnumerateValueKey returned %S", nameBuffer);

        // Free both buffers
        ExFreePool(regBuffer);
        ExFreePool(nameBuffer);

        // Increment our index
        ++index;
    }

    return;
}