/*++
*******************************************************************************
* Copyright (c) 1990-98 Digi International Inc., All Rights Reserved				*
* 																										*
* This software contains proprietary and confidential information of Digi		*
* International Inc.  By accepting transfer of this copy, Recipient agrees		*
* to retain this software in confidence, to prevent disclosure to others,		*
* and to make no use of this software other than that for which it was			*
* delivered.  This is an unpublished copyrighted work of Digi International	*
* Inc.  Except as permitted by federal law, 17 USC 117, copying is strictly	*
* prohibited.																						*
* 																										*
* Restricted Rights Legend																		*
*																										*
* Use, duplication, or disclosure by the Government is subject to					*
* restrictions set forth in sub-paragraph (c)(1)(ii) of The Rights in			*
* Technical Data and Computer Software clause at DFARS 252.227-7031 or			*
* subparagraphs (c)(1) and (2) of the Commercial Computer Software -				*
* Restricted Rights at 48 CFR 52.227-19, as applicable.								*
*																										*
* Digi International Inc. 11001 Bren Road East, Minnetonka, MN 55343				*
*																										*
*******************************************************************************

Module Name:

   digireg.c

Abstract:

   This module contains routines to read values from the registry.

Author:

   Michael Zarns

--*/
#include "digilib.h"


NTSTATUS
DigiRegOpenKey(
	OUT PHANDLE phKey,
	IN PCWSTR   KeyName )
/*++

Description:

	Opens a registry key, returning a handle to be used in subsequent accesses of that 
	key.

Parameters:

	phKey		
		Points to a returned handle if this call is successful.

	KeyName
		The registry key to open

Return Value:

	STATUS_SUCCESS,
	STATUS_INSUFFICIENT_RESOURCES, 
	STATUS_INVALID_HANDLE,
	STATUS_ACCESS_DENIED.

--*/
{
	NTSTATUS Status;
	OBJECT_ATTRIBUTES Attrib;
	UNICODE_STRING Name;
	HANDLE hKey;

	RtlInitUnicodeString( &Name, KeyName );
	InitializeObjectAttributes( &Attrib, &Name, OBJ_CASE_INSENSITIVE, NULL, NULL );
	Status = ZwOpenKey( &hKey, KEY_ALL_ACCESS, &Attrib );

	if ( NT_SUCCESS( Status ) )
		*phKey = hKey;

	return Status;
}


NTSTATUS
DigiRegOpenSubkey(
	OUT PHANDLE phSubKey,
	IN HANDLE   hKeyParent,
	IN PCWSTR   SubKeyName )
/*++

Description:

	Opens a registry key, returning a handle to be used in subsequent accesses of that 
	key.

Parameters:

	phSubKey
		Points to a returned handle if this call is successful.

	hKeyParent
		The parent key handle.

	SubKeyName
		The registry subkey to open

Return Value:

	STATUS_SUCCESS,
	STATUS_INSUFFICIENT_RESOURCES, 
	STATUS_INVALID_HANDLE,
	STATUS_ACCESS_DENIED.

--*/
{
	NTSTATUS Status;
	OBJECT_ATTRIBUTES Attrib;
	HANDLE hSubKey;
	UNICODE_STRING Name;

	RtlInitUnicodeString( &Name, SubKeyName );
	InitializeObjectAttributes( &Attrib, &Name, OBJ_CASE_INSENSITIVE, hKeyParent, NULL );
	Status = ZwOpenKey( &hSubKey, KEY_ALL_ACCESS, &Attrib );

	if ( NT_SUCCESS( Status ) )
		*phSubKey = hSubKey;

	return Status;
}


VOID
DigiRegCloseKey( IN HANDLE hKey )
/*++

Description:

	Closes the registry key, freeing resources used for that key.

Parameters:

	pwszKey	
		The registry key to open

	phKey		
		Points to a returned handle if this call is successful.

Return Value:

	None

--*/
{
	if ( hKey != NULL )
		ZwClose( hKey );
}


NTSTATUS
DigiRegCreateKey(
	OUT PHANDLE phKey,
	IN PCWSTR   KeyName )
/*++

Description:

	Creates a new registry key returning a handle to be used in subsequent accesses of 
	that key.

Parameters:

	phKey		
		Points to a returned handle if this call is successful.

	KeyName
		The registry key to create/open

Return Value:

	STATUS_SUCCESS,
	STATUS_INSUFFICIENT_RESOURCES, 
	STATUS_INVALID_HANDLE,
	STATUS_ACCESS_DENIED.

--*/
{
	NTSTATUS Status;
	OBJECT_ATTRIBUTES Attrib;
	ULONG Disposition;
	HANDLE hKey;
	UNICODE_STRING Name;

	RtlInitUnicodeString( &Name, KeyName );
	InitializeObjectAttributes( &Attrib, &Name, OBJ_CASE_INSENSITIVE, NULL, NULL );
	Status = ZwCreateKey( &hKey, KEY_ALL_ACCESS, &Attrib, 0, NULL, 0, &Disposition );

	if ( NT_SUCCESS( Status ) )
		*phKey = hKey;

	return Status;
}


NTSTATUS
DigiRegCreateSubkey(
	OUT PHANDLE phSubKey,
	IN HANDLE   hKeyParent,
	IN PCWSTR   SubKeyName )
/*++

Description:

	Create a registry key, returning a handle to be used in subsequent accesses of 
	that key.

Parameters:

	phSubKey
		Points to a returned handle if this call is successful.

	hKeyParent
		The parent key handle.

	SubKeyName
		The registry subkey to open

Return Value:

	STATUS_SUCCESS,
	STATUS_INSUFFICIENT_RESOURCES, 
	STATUS_INVALID_HANDLE,
	STATUS_ACCESS_DENIED.

--*/
{
	NTSTATUS Status;
	OBJECT_ATTRIBUTES Attrib;
	ULONG Disposition;
	HANDLE hSubKey;
	UNICODE_STRING Name;

	RtlInitUnicodeString( &Name, SubKeyName );
	InitializeObjectAttributes( &Attrib, &Name, OBJ_CASE_INSENSITIVE, hKeyParent, NULL );
	Status = ZwCreateKey( &hSubKey, KEY_ALL_ACCESS, &Attrib, 0, NULL, 0, &Disposition );

	if ( NT_SUCCESS( Status ) )
		*phSubKey = hSubKey;

	return Status;
}


STATIC NTSTATUS
DigiRegWriteValue( 
	IN HANDLE hKey,
	IN PCWSTR pValue,
	IN ULONG  Type,
	IN PVOID  pBuf,
	IN ULONG  cbBuf )
{
	UNICODE_STRING Value;
	RtlInitUnicodeString( &Value, pValue );
	return ZwSetValueKey( hKey, &Value, 0, Type, pBuf, cbBuf );
}


NTSTATUS 
DigiRegWriteBinary( 
	IN HANDLE hKey,
	IN PCWSTR pValue,
	IN PVOID  pBuf,
	IN ULONG  cbBuf )
{
	return DigiRegWriteValue( hKey, pValue, REG_BINARY, pBuf, cbBuf );
}


NTSTATUS 
DigiRegWriteDword( 
	IN HANDLE hKey,
	IN PCWSTR pValue,
	IN ULONG  dw )
{
	return DigiRegWriteValue( hKey, pValue, REG_DWORD, &dw, sizeof dw );
}


NTSTATUS 
DigiRegWriteSz( 
	IN HANDLE hKey,
	IN PCWSTR pValue,
	IN PCWSTR pSz )
{
	UNICODE_STRING Sz;
	RtlInitUnicodeString( &Sz, pSz );
	return DigiRegWriteValue( hKey, pValue, REG_SZ, Sz.Buffer, Sz.MaximumLength );
}


NTSTATUS 
DigiRegWriteMultiSz( 
	IN HANDLE hKey,
	IN PCWSTR pValue,
	IN PUNICODE_STRING pMulti )
{
	return DigiRegWriteValue( hKey, pValue, REG_MULTI_SZ, pMulti->Buffer, pMulti->Length+sizeof(WCHAR) );
}


STATIC NTSTATUS
DigiRegQuery( 
	IN HANDLE     hKey, 
	IN PCWSTR     pValue, 
	IN ULONG      Type, 
	OUT PVOID     pBuf,
	IN ULONG      Flags,
	IN VOID const *pDefault OPTIONAL,
	IN ULONG      cbDefault )
/*++

Description:

	Internal routine which performs actual registry reads.  Verifies registry value 
	type equals the requested type.  

Parameters:

	pKey	
		Internal struct defining the registry key to query

	pValue
		Wide string defining name of registry value to query

	Type
		REG_??? const defining expected registry value type

	pBuf
		Output buffer which will contain the value read

	Flags
		Flags define behavior when registry read fails.  Optional == use default,
		Required == return failure, ExceptionOnError == generate an exception.

	pDefault
		Optional default, if the registry read fails and Flags == Optional

	cbDefault
		Size of default value (in bytes).

Return Value:

	STATUS_SUCCESS, 
	STATUS_BUFFER_ALL_ZEROS,	// RegVal *not* present, default used.  NT_INFORMATION(Status)==TRUE.
	or error.

--*/
{
	ULONG Length;
	PKEY_VALUE_PARTIAL_INFORMATION pInfo = NULL;
	UNICODE_STRING ValueName;
	VOID const *pResult;
	ULONG cbResult;
	NTSTATUS Status;
	
	RtlInitUnicodeString( &ValueName, pValue );
#ifdef REGISTRY_DEBUG
  	DigiDump( DEBUG_DIGILIB, 
	          ( "DigiRegQuery:\n"
	          "   hKey      = 0x%x\n"
	          "   ValueName = %S\n"
	          "   pBuf      = 0x%x\n"
	          "   Type      = %s\n"
	          "   pDefault  = 0x%x\n"
	          "   cbDefault = 0x%x\n",
	          hKey, pValue, pBuf,
	          Type == REG_DWORD     ? "REG_DWORD"     : 
				 Type == REG_SZ        ? "REG_SZ"        :
				 Type == REG_EXPAND_SZ ? "REG_EXPAND_SZ" :
				 Type == REG_MULTI_SZ  ? "REG_MULTI_SZ"  : "Other",
				 pDefault, cbDefault ) );
#endif
	//
	// Get Info length, alloc buffer for info, read the info, check the value type
	//

	Status = ZwQueryValueKey( hKey, &ValueName, KeyValuePartialInformation, NULL, 0, &Length );
	if ( Status == STATUS_BUFFER_TOO_SMALL )
		pInfo = DigiAllocMem( PagedPool, Length );

	if ( pInfo != NULL )
		Status = ZwQueryValueKey( hKey, &ValueName, KeyValuePartialInformation, 
		                          pInfo, Length, &Length );

#ifdef REGISTRY_DEBUG
	DigiDump( DEBUG_DIGILIB, ("   QueryStatus = 0x%x\n", Status ) );
#endif

	if ( NT_SUCCESS( Status ) && pInfo->Type != Type )
		Status = STATUS_UNSUCCESSFUL;

	if ( NT_ERROR( Status ) )
	{
		// Error occurred.  If read is required, fail the call.  Otherwise, use default.
		if ( ( Flags & DigiRegFlagRequired ) != 0 )
		{
#ifndef REGISTRY_DEBUG
			DigiDump( DEBUG_ERROR, ( "*** DigiRegQuery: Required value %S not found! ***\n", pValue ) );
#endif
			goto Exit;
		}

		pResult = pDefault;
		cbResult = cbDefault;
		Status = STATUS_BUFFER_ALL_ZEROS;	// TODO: Find a better INFO return code!
	}
	else
	{
		// Read was successful.  Use data from registry.
		pResult = pInfo->Data;
		cbResult = pInfo->DataLength;
	}

	switch ( Type )
	{
		case REG_DWORD:
		{
			PULONG pDest = pBuf;
			ASSERT( cbResult == sizeof *pDest );
			*pDest = *(PULONG)pResult;
#ifdef REGISTRY_DEBUG
			DigiDump( DEBUG_DIGILIB, ( "   Dword       = 0x%x\n", *pDest ) );
#else
			DigiDump( DEBUG_DIGILIB, ( "DigiRegQuery: %S = 0x%x\n", pValue, *pDest ) );
#endif
			Status = STATUS_SUCCESS;
			break;
		}
		case REG_SZ:
		case REG_MULTI_SZ:
		case REG_EXPAND_SZ:
		{
			PUNICODE_STRING pDest = pBuf;
			UNICODE_STRING Source;

			Source.Buffer = (PVOID)pResult;
			Source.MaximumLength = (USHORT)cbResult;
			Source.Length = Type == REG_MULTI_SZ ? (USHORT)cbResult : 
			                        (USHORT)wcslen(pResult)*sizeof(WCHAR);

			pDest->MaximumLength = (USHORT)cbResult;
			pDest->Length = 0;
			pDest->Buffer = DigiAllocMem( PagedPool, pDest->MaximumLength );
			if ( pDest->Buffer == NULL )
			{
				Status = STATUS_INSUFFICIENT_RESOURCES;
				goto Exit;
			}
			RtlCopyUnicodeString( pDest, &Source );
#ifdef REGISTRY_DEBUG
			DigiDump( DEBUG_DIGILIB, ( "   String      = %S\n", pDest->Buffer ) );
#else
			DigiDump( DEBUG_DIGILIB, ( "DigiRegQuery: %S = %S\n", pValue, pDest->Buffer ) );
#endif
			Status = STATUS_SUCCESS;
			break;
		}
		default:
			Status = STATUS_UNSUCCESSFUL;
			goto Exit;
	}

Exit:

	if ( pInfo != NULL )
		DigiFreeMem( pInfo );

	if ( NT_ERROR( Status ) && Flags & DigiRegFlagExceptionOnError )
		ExRaiseStatus( Status );

#ifdef REGISTRY_DEBUG
	DigiDump( DEBUG_DIGILIB, ( "   FinalStatus = 0x%x\n", Status ) );
#endif
	return Status;
}


NTSTATUS
DigiRegQueryDword(
	IN HANDLE          hKey,
	IN PCWSTR          pValue,
	OUT PULONG         pBuf,
	IN ULONG           Flags,
	IN ULONG           Default OPTIONAL )
{
	ASSERT( MmIsAddressValid( pBuf ) );
	return DigiRegQuery( hKey, pValue, REG_DWORD, pBuf, Flags,
	                     &Default, sizeof Default );
}

NTSTATUS
DigiRegQueryWord(
	IN HANDLE          hKey,
	IN PCWSTR          pValue,
	OUT PUSHORT        pBuf,
	IN ULONG           Flags,
	IN USHORT          Default OPTIONAL )
{
	ULONG Buffer;
	NTSTATUS Status;
	if ( NT_SUCCESS( Status = DigiRegQueryDword( hKey, pValue, &Buffer, Flags, Default ) ) )
	{
		ASSERT( Buffer <= MAXUSHORT );
		*pBuf = (USHORT)Buffer;
	}

	return Status;
}


NTSTATUS
DigiRegQueryByte(
	IN HANDLE          hKey,
	IN PCWSTR          pValue,
	OUT PUCHAR         pBuf,
	IN ULONG           Flags,
	IN UCHAR           Default OPTIONAL )
{
	ULONG Buffer;
	NTSTATUS Status;
	if ( NT_SUCCESS( Status = DigiRegQueryDword( hKey, pValue, &Buffer, Flags, Default ) ) )
	{
		ASSERT( Buffer <= MAXUCHAR );
		*pBuf = (UCHAR)Buffer;
	}

	return Status;
}


STATIC NTSTATUS
DigiRegQueryString(
	IN HANDLE           hKey,
	IN PCWSTR           pValue,
	IN ULONG            Type,
	OUT PUNICODE_STRING pBuf,
	IN ULONG            Flags,
	IN PCWSTR           pDefault )
{
	UNICODE_STRING Default;
	ASSERT( Type == REG_SZ || Type == REG_MULTI_SZ || Type == REG_EXPAND_SZ );
	ASSERT( MmIsAddressValid( pBuf ) );

	ASSERT( Flags & DigiRegFlagRequired ? TRUE : pDefault != NULL );
	RtlInitUnicodeString( &Default, pDefault );
	return DigiRegQuery( hKey, pValue, Type, pBuf, Flags, pDefault, Default.MaximumLength );
}

NTSTATUS 
DigiRegQueryMultiSz( 
	IN HANDLE           hKey,
	IN PCWSTR           pValue,
	OUT PUNICODE_STRING pBuf,
	IN ULONG            Flags,
	IN PCWSTR           pDefault )
{
	return DigiRegQueryString( hKey, pValue, REG_MULTI_SZ, pBuf, Flags, pDefault );
}

NTSTATUS 
DigiRegQueryExpandSz( 
	IN HANDLE           hKey,
	IN PCWSTR           pValue,
	OUT PUNICODE_STRING pBuf,
	IN ULONG            Flags,
	IN PCWSTR           pDefault )
{
	// Not exported in digireg.h because not tested yet.  Does ZwQueryValueKey 
	// automatically expand the string or do I need to do it here?
	return DigiRegQueryString( hKey, pValue, REG_EXPAND_SZ, pBuf, Flags, pDefault );
}

NTSTATUS 
DigiRegQuerySz( 
	IN HANDLE           hKey,
	IN PCWSTR           pValue,
	OUT PUNICODE_STRING pBuf,
	IN ULONG            Flags,
	IN PCWSTR           pDefault )
{
	return DigiRegQueryString( hKey, pValue, REG_SZ, pBuf, Flags, pDefault );
}

VOID
DigiRegFreeString( 
	PUNICODE_STRING pString )
{
	if ( pString->Buffer )
	{
		DigiFreeMem( pString->Buffer );
		RtlInitUnicodeString( pString, NULL );
	}
}


NTSTATUS
DigiRegEnumSubkey( 
	IN HANDLE           hKey,
	IN ULONG            iSubkey,
	OUT PUNICODE_STRING pSubkey )
/*++

Description:

	Alloc's and returns a Unicode string defining the name of the i'th subkey.  To 
	enumerate all subkeys, call this routine repetitively, incrementing the iSubkey 
	value until the returncode is STATUS_NO_MORE_ENTRIES.  Free the string with 
	DigiRegFreeString().

Parameters:

	hKey
		Handle to previously opened parent key.

	iSubkey
		Zero-based index of the subkeys.

	pSubkey
		Output unicode string defines the name of the subkey, indexed by iSubkey.

Return Value:

	STATUS_SUCCESS, 
	STATUS_NO_MORE_ENTRIES,	// End of list 
	or error.

--*/
{
	NTSTATUS Status;
	PKEY_BASIC_INFORMATION pInfo = NULL;
	ULONG cbNeeded, cbWritten;

	// Get length of buffer needed
	Status = ZwEnumerateKey( hKey, iSubkey, KeyBasicInformation, NULL, 0, &cbNeeded );

	if ( Status != STATUS_BUFFER_TOO_SMALL )
		goto Exit;

	pInfo = DigiAllocMem( PagedPool, cbNeeded );
	if ( !pInfo )
	{
		Status = STATUS_INSUFFICIENT_RESOURCES;
		goto Exit;
	}

	Status = ZwEnumerateKey( hKey, iSubkey, KeyBasicInformation, pInfo, cbNeeded, &cbWritten );
	if ( NT_SUCCESS( Status ) )
	{
		ULONG cbBuf = pInfo->NameLength + sizeof(WCHAR);
		PWCHAR pBuf = DigiAllocMem( PagedPool, cbBuf );
		if ( !pBuf )
		{
			Status = STATUS_INSUFFICIENT_RESOURCES;
			goto Exit;
		}
		RtlZeroMemory( pBuf, cbBuf );
		RtlCopyMemory( pBuf, pInfo->Name, pInfo->NameLength );
		RtlInitUnicodeString( pSubkey, pBuf );
		Status = STATUS_SUCCESS;
	}
	
Exit:
	if ( pInfo )
		DigiFreeMem( pInfo );

	return Status;
}