/****************************************************************************/
// keynode.cpp
//
// Copyright (C) 1997-1999 Microsoft Corp.
/****************************************************************************/


#include <stdio.h>
#include "KeyNode.h"

extern ULONG   g_length_TERMSRV_USERREGISTRY_DEFAULT;
extern ULONG   g_length_TERMSRV_INSTALL;
extern WCHAR   g_debugFileName[MAX_PATH];
extern FILE    *g_debugFilePointer;
extern BOOLEAN g_debugIO;

KeyBasicInfo::KeyBasicInfo():
    pNameSz(NULL)
{
    size = sizeof(KEY_BASIC_INFORMATION) + MAX_PATH*sizeof(WCHAR);
    pInfo = ( KEY_BASIC_INFORMATION *)RtlAllocateHeap(RtlProcessHeap(), 0, size );

    if (!pInfo) {
        status = STATUS_NO_MEMORY;
        pInfo=NULL;
    }
    else
        status = STATUS_SUCCESS;

}

KeyBasicInfo::~KeyBasicInfo()
{
    if (pInfo)
    {
        RtlFreeHeap( RtlProcessHeap(), 0, pInfo);
    }

    if (pNameSz)
    {
        delete pNameSz;
    }
}

PCWSTR KeyBasicInfo::NameSz()
{
    if (Ptr()->NameLength < 2 * MAX_PATH )
    {
        if (!pNameSz)
        {
            pNameSz = new WCHAR [ MAX_PATH  + 1 ];
        }

        // the reason we re do this every call of NameSz() is because
        // Ptr() might changes, since KeyBasicInfo is being used as a
        // scratch pad and passed around for storing pointers to some
        // basic set of info on any key.

        // see if allocation was successful
        if ( pNameSz )
        {
            for ( ULONG i=0; i < Ptr()->NameLength / sizeof(WCHAR) ; i++)
            {
                pNameSz[i] = ( (USHORT)Ptr()->Name[i] );
            }
            pNameSz[i]=L'\0';
        }
    }
    else
    {
        status = STATUS_NO_MEMORY;
        pNameSz[0]=L'\0';
    }

    return pNameSz;

}

#if 0 // NOT USED yet!
KeyNodeInfo::KeyNodeInfo()
{
    size = sizeof(KEY_NODE_INFORMATION) + MAX_PATH*sizeof(WCHAR);
    pInfo = ( KEY_NODE_INFORMATION *)RtlAllocateHeap(RtlProcessHeap(), 0, size );

    if (!pInfo) {
        status = STATUS_NO_MEMORY;
        pInfo=NULL;
    }
    else
        status = STATUS_SUCCESS;


}
KeyNodeInfo::~KeyNodeInfo()
{
    if (pInfo)
    {
        RtlFreeHeap( RtlProcessHeap(), 0, pInfo);
    }
}

#endif

KeyFullInfo::KeyFullInfo() :
    pInfo(NULL)
{
    size = sizeof(KEY_FULL_INFORMATION) + MAX_PATH*sizeof(WCHAR);
    pInfo = ( KEY_FULL_INFORMATION *)RtlAllocateHeap(RtlProcessHeap(), 0, size );

    if (!pInfo) {
        status = STATUS_NO_MEMORY;
        pInfo=NULL;
    }
    else
        status = STATUS_SUCCESS;


}

KeyFullInfo::~KeyFullInfo()
{
    if (pInfo)
    {
        RtlFreeHeap( RtlProcessHeap(), 0, pInfo);
    }

}

KeyNode::KeyNode(HANDLE root, ACCESS_MASK access, PCWSTR name ) :
    root(NULL), hKey(NULL),
    accessMask(NULL),basic(NULL),full(NULL),
    pFullPath(NULL), pNameSz(NULL)
{
    hKey = NULL;
    PCWSTR n = name;
    accessMask = access;

    RtlInitUnicodeString(&uniName, n);

    InitializeObjectAttributes(&ObjAttr,
                           &uniName,
                           OBJ_CASE_INSENSITIVE,
                           root,
                           NULL);
    status=STATUS_SUCCESS;
}

KeyNode::KeyNode(KeyNode *pParent, KeyBasicInfo   *pInfo ) :
    root(NULL), hKey(NULL),
    accessMask(NULL),basic(NULL), full(NULL),
    pFullPath(NULL), pNameSz(NULL)
{
    hKey = NULL;
    PCWSTR n = pInfo->NameSz();
    accessMask = pParent->Masks();

    RtlInitUnicodeString(&uniName, n);

    InitializeObjectAttributes(&ObjAttr,
                           &uniName,
                           OBJ_CASE_INSENSITIVE,
                           pParent->Key(),
                           NULL);
    status=STATUS_SUCCESS;

}


KeyNode::~KeyNode()
{
    Close();

    if (basic)
    {
        delete basic;
    }

    if (full)
    {
        delete full;
    }

    if (pFullPath)
    {
        RtlFreeHeap(RtlProcessHeap(), 0, pFullPath);
    }

    if( pNameSz )
    {
        delete pNameSz;
    }

}

NTSTATUS KeyNode::Open()
{
    status = NtOpenKey(&hKey,
                    accessMask,
                    &ObjAttr);


    if ( !NT_SUCCESS( status))
    {
       hKey=NULL;
       // Debug(DBG_OPEN_FAILED );
    }

    return status;

}

NTSTATUS KeyNode::Close()
{

    if ( hKey )
    {
        status = NtClose( hKey );
        hKey = 0;
    }

    return status;
}

NTSTATUS KeyNode::Create( UNICODE_STRING *uClass)
{
    ULONG   ultmp;
    status = NtCreateKey(&hKey,
                         accessMask,
                         &ObjAttr,
                         0,
                         uClass,
                         REG_OPTION_NON_VOLATILE,
                         &ultmp);
    // Debug(DBG_CREATE);

    return status;
}



// Recursively create the reg path given by the uniName member variable
// Upon completion, open the reg key for access.

NTSTATUS KeyNode::CreateEx( UNICODE_STRING *uClass)
{
    ULONG   wsize = uniName.Length/sizeof(WCHAR);
    PWCHAR pTmpFullPath = new WCHAR[  uniName.Length + sizeof( WCHAR ) ];

    if(!pTmpFullPath)
    {
        status = STATUS_NO_MEMORY;
        return status;
    }

    wcsncpy(pTmpFullPath, uniName.Buffer , wsize);
    pTmpFullPath[ wsize ] = L'\0';

    PWCHAR    p;
    WCHAR     sep[]= {L"\\"};
    p = wcstok( pTmpFullPath, sep);

    // we know how many keys to create now.
    // start over again
    wcsncpy(pTmpFullPath, uniName.Buffer , wsize );
    pTmpFullPath[ wsize ] = L'\0';

    KeyNode *pKN1=NULL, *pKN2=NULL;
    p = wcstok( pTmpFullPath, sep);

    // the first item is "Registry", make it "\Registry" since we are opening
    // from the root.
    PWCHAR pTmpName = new WCHAR[  wcslen(p) + 2 ];

    if(!pTmpName)
    {
        DELETE_AND_NULL(pTmpFullPath);
        status = STATUS_NO_MEMORY;
        return status;
    }

    wcscpy(pTmpName, L"\\");
    wcscat( pTmpName , p );

    NTSTATUS st = STATUS_SUCCESS;
    while( p != NULL )
    {
        // @@@
        // ADD error handling, else you will create keys in the wrong places instead of bailing out.
        // @@@

        if ( pKN2 )
        {
            // ---- STEP 3 ---

            // NOT-first time around

            p = wcstok( NULL, sep);

            if ( p )    // we have more sub keys
            {
                pKN1 = new KeyNode( pKN2->Key(),  accessMask,  p );
                if (pKN1)
                {
                    st = pKN1->Open();

                    // if Open fails, then key does not exist, so create it
                    if ( !NT_SUCCESS( st ))
                    {
                        st = pKN1->Create();
                    }
                }
                else
                {
                    status = STATUS_NO_MEMORY;
                    break;
                }
            }
        }
        else
        {
            // ---- STEP 1 ---

            // First time around, we are opening \Registry node, use
            // pTmpName instead of "p"
            pKN1 = new KeyNode( NULL, accessMask , pTmpName );
            if (pKN1)
            {
                st = pKN1->Open();
            }
            else
            {
                status = STATUS_NO_MEMORY ;
                break;
            }

        }

        p = wcstok( NULL, sep);

        if (p)  // we have more sub keys
        {

            // ---- STEP 2 ---

            pKN2 = new KeyNode( pKN1->Key(), accessMask, p );
            if (pKN2 )
            {
                st = pKN2->Open();
                if ( !NT_SUCCESS( pKN2->Status() ))
                {
                    st = pKN2->Create();
                }
            }
            else
            {
                status = STATUS_NO_MEMORY;
                DELETE_AND_NULL (pKN1);
                break;
            }

            DELETE_AND_NULL (pKN1);
            pKN1 = pKN2;
        }
    }

    DELETE_AND_NULL( pKN2 );

    // since the last node was created above, now we can open ourselfs incase
    // caller wants to use us.
    if ( NT_SUCCESS(status) )
    {
        Open();
    }

    DELETE_AND_NULL(pTmpName);

    DELETE_AND_NULL(pTmpFullPath);

    return status;

}

NTSTATUS KeyNode::Delete()
{
    if (hKey)
    {
        status = NtDeleteKey( hKey );
        // Debug(DBG_DELETE);
    }

    return status;
}


NTSTATUS KeyNode::DeleteSubKeys()
{
    if (hKey && NT_SUCCESS( status ))
    {
        KeyBasicInfo    basicInfo;
        status = basicInfo.Status();

        if (NT_SUCCESS( status )) 
        {
            status = EnumerateAndDeleteSubKeys( this, &basicInfo );
        }
    }
    return status;
}

NTSTATUS KeyNode::EnumerateAndDeleteSubKeys(
    IN KeyNode      *pSource,
    IN KeyBasicInfo *pBasicInfo )
{
    NTSTATUS  st = STATUS_SUCCESS;

    ULONG   ulCount=0;
    ULONG   ultemp;

    while (NT_SUCCESS(st) && st != STATUS_NO_MORE_ENTRIES )
    {
        ULONG       ultemp;
        NTSTATUS    st2;

        st = NtEnumerateKey(    pSource->Key(),
                                    ulCount,
                                    pBasicInfo->Type(),
                                    pBasicInfo->Ptr(),
                                    pBasicInfo->Size(),
                                    &ultemp);

        if (NT_SUCCESS(st) && st != STATUS_NO_MORE_ENTRIES )
        {
            pBasicInfo->Ptr()->Name[ pBasicInfo->Ptr()->NameLength/sizeof(WCHAR) ] = L'\0';

            KeyNode SourcesubKey(pSource, pBasicInfo);

            if (NT_SUCCESS( SourcesubKey.Open() )  )
            {

                // enumerate sub key down.
                st2 = EnumerateAndDeleteSubKeys(
                            &SourcesubKey,
                            pBasicInfo );

            }

            st = SourcesubKey.Delete();
        }

    }

    return st;
}

#if 0
NTSTATUS KeyNode::Query( KEY_NODE_INFORMATION **result , ULONG   *resultSize)
{

    if ( hKey )
    {
        // first time around we allocate memory and keep using it
        // as our scratch pad
        if (!node )
        {
            node = new KeyNodeInfo();
        }

        status = NtQueryKey(hKey,
            node->Type(),     //        Keynode,
            node->Ptr(),
            node->Size(),
            resultSize);

        *result = node->Ptr();
    }
    else
        status = STATUS_OBJECT_NAME_NOT_FOUND; // need to call open or key is not found

    return status;

}
#endif

NTSTATUS KeyNode::Query( KEY_FULL_INFORMATION **result , ULONG   *pResultSize)
{

    if ( hKey )
    {
        // first time around we allocate memory and keep using it
        // as our scratch pad
        if (!full )
        {
            full = new KeyFullInfo();
        }

        if (full)
        {
            status = NtQueryKey(hKey,
                full->Type(),     //        KeyFullInformation,
                full->Ptr(),
                full->Size(),
                pResultSize);
    
            *result = full->Ptr();
        }
        else
            status = STATUS_NO_MEMORY ;
    }
    else
        status = STATUS_OBJECT_NAME_NOT_FOUND; // need to call open or key is not found

    return status;

}

//This will allocate and set pFullPath. pFullPath is a global variable which is deallocated
//by the destructor, so this function should only be called if pFullPath has not already
//been allocated for this object.
NTSTATUS KeyNode::GenerateFullPath()
{
    status = STATUS_SUCCESS;

    // A key handle or root directory was specified, so get its path
    if (hKey)
    {
        ULONG ultemp = 0;
        ultemp = sizeof(UNICODE_STRING) + (sizeof(WCHAR) * MAX_PATH * 2);
        pFullPath = RtlAllocateHeap(RtlProcessHeap(),
                                  0,
                                  ultemp);

        // Got the buffer OK, query the path
        if (pFullPath)
        {
            // Get the path for key or root directory
            status = NtQueryObject(hKey ,
                                   ObjectNameInformation,
                                   (PVOID)pFullPath,
                                   ultemp,
                                   NULL);

            if (!NT_SUCCESS(status))
                RtlFreeHeap(RtlProcessHeap(), 0, pFullPath);
        }
        else
             status = STATUS_NO_MEMORY;
    }
    else
        status = STATUS_OBJECT_NAME_NOT_FOUND; // need to call open or key is not found

    return (status);
}


NTSTATUS KeyNode::GetPath(PWCHAR *pwch)
{
    status = STATUS_SUCCESS;

    if (pFullPath == NULL)
        status = GenerateFullPath();

    if (NT_SUCCESS(status))
    {
        // Build the full path to the key to be created
        *pwch = ((PUNICODE_STRING)pFullPath)->Buffer;

        // Make sure the string is zero terminated
        ULONG ulWcharLength = 0;
        ulWcharLength = ((PUNICODE_STRING)pFullPath)->Length / sizeof(WCHAR);
        (*pwch)[ulWcharLength] = L'\0';
    }

    return (status);
}


void KeyNode::Debug( DebugType type )
{
    if ( debug )
    {
        ULONG i;

        switch( type )
        {
        case DBG_DELETE :
            fwprintf( g_debugFilePointer ,
                    L"Deleted key=%lx; status=%lx, name=", status, hKey );
            DbgPrint("Deleted key=%lx; status=%lx, name=", status, hKey );
            break;

        case DBG_OPEN_FAILED:
            fwprintf( g_debugFilePointer,
                    L"Unable to Open, status=%lx, name=", hKey, status );
            DbgPrint("Unable to Open, status=%lx, name=", hKey, status );
            break;

        case DBG_KEY_NAME:
            fwprintf( g_debugFilePointer,
                    L"hKey=%lx, name=", hKey);
            DbgPrint("hKey=%lx, name=", hKey);
            break;

        case DBG_CREATE:
            fwprintf( g_debugFilePointer,
                    L"Created hKey=%lx, status=%lx,name=", hKey, status);
            DbgPrint("Created hKey=%lx, status=%lx,name=", hKey, status );
            break;
        }

        fwprintf( g_debugFilePointer, L"%s\n",NameSz() );
        fflush( g_debugFilePointer );
        DbgPrint("%s\n",(char *)NameSz() );
    }
}


PCWSTR KeyNode::NameSz()
{
    if (!pNameSz)
    {
        pNameSz = new WCHAR [ uniName.Length / sizeof(WCHAR) + 1 ];

        if (pNameSz)
        {
            for ( ULONG i=0; i < uniName.Length / sizeof(WCHAR) ; i++)
            {
                pNameSz[i] = ( (USHORT)uniName.Buffer[i] );
            }
            pNameSz[i]=L'\0';
        }
        else
        {
            status = STATUS_NO_MEMORY;
        }
    }

    return pNameSz;
}

NTSTATUS KeyNode::GetFullInfo( KeyFullInfo   **p)
{
    // do a self query
    if ( !full )
    {
        ULONG   size;
        KEY_FULL_INFORMATION    *tmp;
        Query( &tmp , &size );
    }

    *p = full;

    return status;

}