// ==========================================================================
// Author:  Yee Hsu
// Date:    9/5/2008
// File:    Registry.cpp
//
// Desc:    Registry class wrapper API that allows easier functions to 
//          manipulate the registry. Access, write, read, and query the
//          registry using this class.
// ==========================================================================

#include "stdafx.h"
#include "Registry.h"
#include "IntCommon.h"

// ==========================================================================
// Identifier:  Registry()
//
// Description: Registry Constructor
// ==========================================================================

Registry::Registry()
{
    this->hKey = NULL;
    this->mRegistry.clear();

    this->mHKey[HKEY_CLASSES_ROOT]      = "HKEY_CLASSES_ROOT";
    this->mHKey[HKEY_CURRENT_USER]      = "HKEY_CURRENT_USER";
    this->mHKey[HKEY_LOCAL_MACHINE]     = "HKEY_LOCAL_MACHINE";
    this->mHKey[HKEY_USERS]             = "HKEY_USERS";
    this->mHKey[HKEY_CURRENT_CONFIG]    = "HKEY_CURRENT_CONFIG";
}

// ==========================================================================
// Identifier:  Registry()
//
// Description: Registry Destructor
// ==========================================================================

Registry::~Registry()
{
    this->mRegistry.clear();
}

// ==========================================================================
// Identifier:  EnumRegistryKey()
//
// Description: Enumerates the registries base on keyname
// ==========================================================================

MapRegistry Registry::EnumRegistryKey(const HKEY hKey, const string sKeyName)
{
    this->mRegistry.clear();
    this->hKey = hKey;
    this->sKeyName = sKeyName;

    VectorEnumKey       ListEnumKey;
    VectorEnumKeyValue  ListEnumKeyValue;
    RegistryInfo        RegInfo;

    EnumRegistryKey(this->hKey, sKeyName, ListEnumKey);

    for (VectorEnumKey::iterator iKey = ListEnumKey.begin(); iKey != ListEnumKey.end(); iKey++)
    {
        RegInfo.sFullKeyName = this->sKeyName + iKey->c_str();
        EnumRegistryKeyValue(this->hKey, RegInfo.sFullKeyName, ListEnumKeyValue);
        RegInfo.ListEnumKeyValue = ListEnumKeyValue;

        this->mRegistry[RegInfo.sFullKeyName] = RegInfo;
    }
    return this->mRegistry;
}

// ==========================================================================
// Identifier:  GetRegistryType()
//
// Description: Gets registry type
// ==========================================================================

std::string Registry::GetRegistryType(const DWORD dwRegType)
{
    string sRegType = "";

    switch (dwRegType)
    {
        case REG_SZ:                            sRegType = "REG_SZ";                            break;
        case REG_EXPAND_SZ:                     sRegType = "REG_EXPAND_SZ";                     break;
        case REG_MULTI_SZ:                      sRegType = "REG_MULTI_SZ";                      break;
        case REG_BINARY:                        sRegType = "REG_BINARY";                        break;
        case REG_RESOURCE_LIST:                 sRegType = "REG_RESOURCE_LIST";                 break;
        case REG_FULL_RESOURCE_DESCRIPTOR:      sRegType = "REG_FULL_RESOURCE_DESCRIPTOR";      break;
        case REG_RESOURCE_REQUIREMENTS_LIST:    sRegType = "REG_RESOURCE_REQUIREMENTS_LIST";    break;
        case REG_DWORD:                         sRegType = "REG_DWORD";                         break;
        default:                                sRegType = "REG_UNKNOWN";                       break;
    }
    return sRegType;
}

// ==========================================================================
// Identifier:  GetRegistryData()
//
// Description: Gets registry data type
// ==========================================================================

std::string Registry::GetRegistryData(const DWORD dwRegType, const RegData m_RegData)
{
    string sRegData = "";

    switch (dwRegType)
    {
        case REG_SZ:
        case REG_EXPAND_SZ:
        case REG_MULTI_SZ:
            sRegData = m_RegData.sData;
            break;

        case REG_BINARY:
        case REG_RESOURCE_LIST:
        case REG_FULL_RESOURCE_DESCRIPTOR:
        case REG_RESOURCE_REQUIREMENTS_LIST:
            if (m_RegData.pData) { sRegData = m_RegData.pData; }
            break;

        case REG_DWORD:
            char szValue[32];
            sprintf(szValue, "0x%08x", m_RegData.dwData);
            sRegData = szValue;
            break;
    }
    return sRegData;
}

// ==========================================================================
// Identifier:  GetHKeyString()
//
// Description: Gets registry HKEY string
// ==========================================================================

std::string Registry::GetHKeyString(const HKEY hKey)
{
    return this->mHKey[hKey];
}

// ==========================================================================
// Identifier:  GetHKeyString()
//
// Description: Gets registry HKEY string
// ==========================================================================

std::string Registry::GetHKeyString(void)
{
    return this->mHKey[this->hKey];
}

// ==========================================================================
// Identifier:  EnumRegistryTree()
//
// Description: Enumererates the entire registry tree
// ==========================================================================

MapRegistry Registry::EnumRegistryTree(HKEY hKey, const string sKeyName)
{
    this->mRegistry.clear();
    this->hKey = hKey;
    this->sKeyName = sKeyName;
    this->EnumRegistryKey(sKeyName);
    return this->mRegistry;
}

// ==========================================================================
// Identifier:  EnumRegistryKey()
//
// Description: Enumererates registry tree by keyname recursively
// ==========================================================================

void Registry::EnumRegistryKey(const string sKeyName)       // TOP DOG RECURSIVE FUNCTION!
{
    VectorEnumKey ListEnumKey;

    EnumRegistryKey(this->hKey, sKeyName, ListEnumKey);

    for (int i = 0; i < ListEnumKey.size(); i++)
    {
        RegistryInfo RegInfo;
        VectorEnumKeyValue ListEnumKeyValue;

        RegInfo.sFullKeyName = sKeyName + ListEnumKey[i].c_str();
        this->EnumRegistryKeyValue(this->hKey, RegInfo.sFullKeyName, ListEnumKeyValue);
        RegInfo.ListEnumKeyValue = ListEnumKeyValue;
        this->mRegistry[RegInfo.sFullKeyName] = RegInfo;

        //////////////////////////////////////////////////////////////////// FLUSH TO STDOUT for real time parsing
        fprintf(stderr, "%s\n", RegInfo.sFullKeyName.c_str());
        //////////////////////////////////////////////////////////////////////////////////////////////////////////  

        this->EnumRegistryKey(sKeyName + ListEnumKey[i] + "\\");
    }
}

// ==========================================================================
// Identifier:  FlushRegistryToTxt()
//
// Description: Writes registry information to tect file
// ==========================================================================

void Registry::FlushRegistryToTxt(FILE* fp)
{
    for (MapRegistry::iterator iRegistry = this->mRegistry.begin(); iRegistry != this->mRegistry.end(); iRegistry++)
    {
        fprintf(fp, "%s/%s\n", this->mHKey[this->hKey].c_str(), iRegistry->first.c_str());

        for (VectorEnumKeyValue::iterator iSubKey = iRegistry->second.ListEnumKeyValue.begin();
            iSubKey != iRegistry->second.ListEnumKeyValue.end(); iSubKey++)
        {
            // PRINT ONLY KEYS FOR DisplayName OR DisplayVersion
            if (iSubKey->sRegKeyName == "DisplayName" || iSubKey->sRegKeyName == "DisplayVersion")
            {
                fprintf(fp, "\t%s-%s-%s\n",
                    iSubKey->sRegKeyName.c_str(),
                    this->GetRegistryType(iSubKey->dwRegType).c_str(),
                    this->GetRegistryData(iSubKey->dwRegType, iSubKey->m_RegData).c_str());
            }
        }
    }
}

// ==========================================================================
// Identifier:  FlushRegistryToXml()
//
// Description: Writes registry information to XML file
// ==========================================================================

void Registry::FlushRegistryToXml(FILE* fp)
{
    fprintf(fp, "\t<Registry Path=\"%s\">\n", this->mHKey[this->hKey].c_str());
    for (MapRegistry::iterator iRegistry = this->mRegistry.begin(); iRegistry != this->mRegistry.end(); iRegistry++)
    {
        fprintf(fp, "\t\t<RegSub Key=\"%s\">\n", ConvertStringToXmlCompliant(iRegistry->second.sFullKeyName).c_str());

        for (VectorEnumKeyValue::iterator iSubKey = iRegistry->second.ListEnumKeyValue.begin();
            iSubKey != iRegistry->second.ListEnumKeyValue.end(); iSubKey++)
        {
            // PRINT ONLY KEYS FOR DisplayName OR DisplayVersion
            if (iSubKey->sRegKeyName == "DisplayName" || iSubKey->sRegKeyName == "DisplayVersion")
            {
                fprintf(fp, "\t\t\t<Reg Name=\"%s\" Type=\"%s\" Data=\"%s\" />\n",
                    iSubKey->sRegKeyName.c_str(),
                    this->GetRegistryType(iSubKey->dwRegType).c_str(),
                    ConvertStringToXmlCompliant(this->GetRegistryData(iSubKey->dwRegType, iSubKey->m_RegData)).c_str());
            }
        }
        fprintf(fp, "\t\t</RegSub>\n");
    }
    fprintf(fp, "\t</Registry>\n");
}


/* ****************************** DONT MODIFY BELOW OR DIE *********************************************** */

// ==========================================================================
// Identifier:  EnumRegistryKey()
//
// Description: Enumerates the regsitry hive
// ==========================================================================

void Registry::EnumRegistryKey(HKEY hKey, string sKeyName, VectorEnumKey& ListEnumKey)
{
    LONG retcode = ERROR_SUCCESS;
    HKEY hOpenKey = NULL;
    char str[MAX_PATH];
    memset( str, '\0', sizeof(str));
    ListEnumKey.clear();
    
    if(sKeyName.length()>0)
        retcode = RegOpenKey( hKey, sKeyName.c_str(), &hOpenKey);
    else
        hOpenKey = hKey;

    if( retcode != (DWORD)ERROR_SUCCESS )
        return;

    for (int i = 0, retCode = ERROR_SUCCESS; retCode == ERROR_SUCCESS; i++) 
    {
        retCode = RegEnumKey(hOpenKey, i, str, MAX_PATH); 

        if (retCode == (DWORD)ERROR_SUCCESS) 
        {
            string sNewKeyName;
            sNewKeyName = str;
            
            if(sNewKeyName.length()>0)
                ListEnumKey.push_back( sNewKeyName );
        }
    }

    if(hKey)
        RegCloseKey( hOpenKey );
}

// ==========================================================================
// Identifier:  EnumRegistryKeyValue()
//
// Description: Enumerates and gets the registry base on key
// ==========================================================================

void Registry::EnumRegistryKeyValue(HKEY hKey, string sKeyName, VectorEnumKeyValue& ListEnumKey)
{
    LONG retcode = ERROR_SUCCESS;
    HKEY hOpenKey = NULL;
    DWORD dwType = REG_SZ;
    char str[MAX_REG_KEY_NAME];
    memset( str, '\0', sizeof(str));
    
    if(sKeyName.length()>0)
        retcode = RegOpenKeyEx(hKey, sKeyName.c_str(), 0, KEY_READ, &hOpenKey);
    else
        hOpenKey = hKey;

    if( retcode != (DWORD)ERROR_SUCCESS )
        return;

    BYTE *data;
    data = new BYTE[MAX_REG_KEY_VALUE];
    memset( data, '\0', sizeof(data));

    ListEnumKey.clear();

    DWORD Size;
    DWORD dwNo = 0;

    RegKeyDetail KeyValue;

    for (int i = 0, retCode = ERROR_SUCCESS; retCode == ERROR_SUCCESS; i++) 
    {
        Size = MAX_REG_KEY_NAME;
        DWORD dwNo = MAX_REG_KEY_VALUE;
        
        retCode = RegEnumValue(hOpenKey, i, str, &Size, NULL, &dwType, data, &dwNo);

        if (retCode != (DWORD) ERROR_NO_MORE_ITEMS)// && retCode != ERROR_INSUFFICIENT_BUFFER)
        {
            KeyValue.sRegKeyName = str;
            KeyValue.dwRegType = dwType;

            if(dwType==REG_SZ)
            {
                if( dwNo >0 && dwNo < MAX_REG_KEY_VALUE )
                {
                    data[(dwNo-1)] = '\0';
                    KeyValue.m_RegData.sData = (char*)data;
                }
                else
                {
                    KeyValue.m_RegData.sData = "";
                }
            }
            else if(dwType==REG_EXPAND_SZ)
            {
                if( dwNo >0 && dwNo < MAX_REG_KEY_VALUE )
                {
                    data[(dwNo-1)] = '\0';
                    KeyValue.m_RegData.sData = (char*)data;
                }
                else
                {
                    KeyValue.m_RegData.sData = "";
                }
            }
            else if(dwType==REG_MULTI_SZ)
            {
                if( dwNo > 1 && dwNo < MAX_REG_KEY_VALUE )
                {
                    for(int i=0;i<dwNo;i++)
                    {
                        if( data[i]=='\0' )
                        {
                            data[i] = '\r';
                        }
                    }
                    data[(dwNo-1)] = '\0';
                    data[(dwNo-2)] = '\0';

                    KeyValue.m_RegData.sData = (char*)data;
                }
                else
                {
                    KeyValue.m_RegData.sData = "";
                }
            }
            else if( (dwType==REG_BINARY)|| (dwType==REG_RESOURCE_LIST) || (dwType==REG_FULL_RESOURCE_DESCRIPTOR) || (dwType==REG_RESOURCE_REQUIREMENTS_LIST) )
            {
                if( dwNo >0 && dwNo < MAX_REG_KEY_VALUE )
                {
                    char temp[4];
                    memset(temp, '\0', sizeof(temp));

                    KeyValue.m_RegData.pData = new char[dwNo*3 + 10];
                    memset( KeyValue.m_RegData.pData, '\0', (dwNo*3 + 10) );

                    for(int j=0; j<dwNo; j++)
                    {
                        sprintf(temp, "%02X ", (long int)data[j]);
                        lstrcat( KeyValue.m_RegData.pData, temp );
                    }
                }
                else
                {
                    KeyValue.m_RegData.pData = NULL;
                }
            }
            else if(dwType==REG_DWORD)
            {
                memcpy( &KeyValue.m_RegData.dwData, data, sizeof(DWORD));
            }
            else
            {
                int i = 0;
            }
            
            if(KeyValue.sRegKeyName.length()>0)
                ListEnumKey.push_back( KeyValue );

            retCode = ERROR_SUCCESS;
        }
    }

    if(hKey)
        RegCloseKey( hOpenKey );

    delete [] data;
}