#include "stdafx.h"
#include "rnode.h"


CRegNode::CRegNode()
{
	m_pPrevSibling = m_pParent = m_pFirstChild = m_pNextSibling = NULL;
	m_dwType = REGNODE_TYPE_UNKNOWN;
	m_hItem = NULL;
	m_bValuesLoaded = m_bLoaded = FALSE;
	m_lpData = NULL;
}

CRegNode::CRegNode(CString strName, DWORD dwType)
{
	m_pPrevSibling = m_pParent = m_pFirstChild = m_pNextSibling = NULL;
	m_strName = strName;
	m_dwType = dwType;
	m_hItem = NULL;
	m_bValuesLoaded = m_bLoaded = FALSE;
	m_lpData = NULL;
}

CRegNode::~CRegNode()
{
	CRegNode *pTemp = m_pFirstChild;
	CRegNode *pPrev = NULL;
	while (pTemp)
	{
		pPrev = pTemp;
		pTemp = pTemp->m_pNextSibling;
		if (pPrev->m_lpData)
		{
			delete pPrev->m_lpData;
			pPrev->m_lpData = NULL;
		}
		delete pPrev;
	}
	if (m_lpData)
	{
		delete m_lpData;
		m_lpData = NULL;
	}
}


CRegNode::AddChild(CRegNode *pChild)
{
	CRegNode *pTemp, *pNext;
	pTemp = pNext = m_pFirstChild;
	if (pNext)
	do
	{
		pTemp = pNext;
		pNext = pNext->m_pNextSibling;
	} while (pNext);
	if (pTemp)
	{
		pTemp->m_pNextSibling = pChild;
		pChild->m_pPrevSibling = pTemp;
	}
	else
	{
		m_pFirstChild = pChild;
	}
}

HKEY CRegNode::GetKeyHandle()
{
	HKEY hKey = NULL;
	LONG lRet;
	if (m_strName == _T("HKEY_CLASSES_ROOT"))
		hKey = HKEY_CLASSES_ROOT;
	if (m_strName == _T("HKEY_CURRENT_USER"))
		hKey = HKEY_CURRENT_USER;
	if (m_strName == _T("HKEY_LOCAL_MACHINE"))
		hKey = HKEY_LOCAL_MACHINE;
	if (hKey)
		return hKey;
	HKEY hRoot = GetKeyRoot();
	if (hRoot)
	{
		lRet = RegOpenKeyEx(hRoot, GetKeyPath(), 0, KEY_ALL_ACCESS, &hKey);
		ASSERT(lRet == ERROR_SUCCESS);
	}
	return hKey;
}

HKEY CRegNode::GetKeyRoot()
{
	CRegNode *pTemp = m_pParent;

	while (pTemp)
	{
		if (pTemp->m_dwType & REGNODE_TYPE_ROOT)
		{
			return pTemp->GetKeyHandle();
		}
		pTemp = pTemp->m_pParent;
	}
	return NULL;
}

CString CRegNode::GetKeyPath()
{
	CRegNode *pTemp = m_pParent;
	CString strTemp1, strTemp2;
	strTemp1 = m_strName;
	while (pTemp && !(pTemp->m_dwType & REGNODE_TYPE_ROOT))
	{
		strTemp1 = pTemp->m_strName + _T("\\") + strTemp1;
		pTemp = pTemp->m_pParent;
	}
	return strTemp1;
}

BOOL LoadKey(CRegNode *pParentNode)
{
	HKEY hParent;
	LONG lRet = ERROR_SUCCESS, lRet2 = ERROR_SUCCESS;
	DWORD dwIndex = 0;
	TCHAR szName[512];
	DWORD dwSizeName = sizeof(szName)/sizeof(TCHAR);
	FILETIME ftLastWriteTime;
	CRegNode *pNode;

	if (pParentNode->m_bLoaded)
		return TRUE;
	hParent = pParentNode->GetKeyHandle();
	while (lRet != ERROR_NO_MORE_ITEMS)
	{
		lRet = RegEnumKeyEx(hParent, dwIndex, szName, &dwSizeName, NULL, NULL, NULL, &ftLastWriteTime);
		if (lRet == ERROR_NO_MORE_ITEMS)
			break;
		pNode = new CRegNode(szName, REGNODE_TYPE_KEY);
		AddNode(pNode, pParentNode);
		dwIndex++;
		dwSizeName = sizeof(szName)/sizeof(TCHAR);
	}
	pParentNode->m_bLoaded = TRUE;
	if (!(hParent == HKEY_LOCAL_MACHINE 
		||hParent == HKEY_CURRENT_USER 
		|| hParent == HKEY_CLASSES_ROOT 
		|| hParent == HKEY_USERS 
#ifndef _WIN32_WCE
		|| hParent == HKEY_CURRENT_CONFIG
#endif
		))
		RegCloseKey(hParent);
	return dwIndex > 0;
}


void LoadValues(CRegNode *pNode)
{
	HKEY hParent;
	LONG lRet = ERROR_SUCCESS;
	TCHAR szName[512];
	DWORD dwDataSize, dwSizeName = sizeof(szName)/sizeof(TCHAR);
	FILETIME ftLastWriteTime;
	DWORD dwType, dwSubKeys, cbMaxSubKeyLen, cValues, cbMaxValueNameLen, cbMaxValueLen, cbSecurityDescriptor;
	LPBYTE lpData = 0;
	CRegNode *pValueNode = NULL;

	if (pNode->m_bValuesLoaded)
		return;
	hParent = pNode->GetKeyHandle();
	lRet = RegQueryInfoKey(hParent,NULL, NULL, NULL, &dwSubKeys, &cbMaxSubKeyLen, NULL, &cValues, &cbMaxValueNameLen, &cbMaxValueLen, &cbSecurityDescriptor, &ftLastWriteTime);
	if (lRet != ERROR_SUCCESS)
	{
	if (!(hParent == HKEY_LOCAL_MACHINE 
		||hParent == HKEY_CURRENT_USER 
		|| hParent == HKEY_CLASSES_ROOT 
		|| hParent == HKEY_USERS 
#ifndef _WIN32_WCE
		|| hParent == HKEY_CURRENT_CONFIG
#endif
		))
			RegCloseKey(hParent);
		return;
	}
	for (UINT i = 0; i < cValues; i++)
	{
		dwDataSize = cbMaxValueLen;
		dwSizeName = sizeof(szName)/sizeof(TCHAR);
		lpData = new BYTE[cbMaxValueLen+2];
		lRet = RegEnumValue(hParent, i, szName, &dwSizeName, NULL, &dwType, lpData, &dwDataSize);
		if (lRet != ERROR_SUCCESS)
		{
			delete [] lpData;
			continue;
		}
		pValueNode = new CRegNode;
		pValueNode->m_strName = szName;		
		pValueNode->m_dwType = REGNODE_TYPE_VALUE;
		pValueNode->m_dwDataType = dwType;
		pValueNode->m_lpData = lpData;
		pValueNode->m_dwDataSize = dwDataSize;
		AddNode(pValueNode, pNode);
	}
	pNode->m_bValuesLoaded = TRUE;
	if (!(hParent == HKEY_LOCAL_MACHINE 
		||hParent == HKEY_CURRENT_USER 
		|| hParent == HKEY_CLASSES_ROOT 
		|| hParent == HKEY_USERS 
#ifndef _WIN32_WCE
		|| hParent == HKEY_CURRENT_CONFIG
#endif
		))
		RegCloseKey(hParent);
}

void AddNode(CRegNode *pNode, CRegNode *pParentNode)
{
	pNode->m_pParent = pParentNode;
	CRegNode *pTemp;	
	CRegNode *pPrev = NULL;

	if (pParentNode)
	{
		pTemp = pParentNode->m_pFirstChild;
		while (pTemp)
		{
			pPrev = pTemp;
			pTemp = pTemp->m_pNextSibling;
		}
		if (pPrev)
		{
			pPrev->m_pNextSibling = pNode;
			pNode->m_pPrevSibling = pPrev;
		}
		else
		{
			pParentNode->m_pFirstChild = pNode;
		}
	}
}

void DeleteNode(CRegNode *pNode)
{
	if (pNode->m_pPrevSibling)
		pNode->m_pPrevSibling->m_pNextSibling = pNode->m_pNextSibling;
	else
		pNode->m_pParent->m_pFirstChild = pNode->m_pNextSibling;

	delete pNode;
}

BOOL HasSubKeys(CRegNode *pNode)
{
	CRegNode *pTemp;
	if (pNode->m_pFirstChild) // has sub keys
	{
		pTemp = pNode->m_pFirstChild;
		while (pTemp)
		{
			if (pTemp->m_dwType & REGNODE_TYPE_KEY)
				return TRUE;
			pTemp = pTemp->m_pNextSibling;
		}
	}
	return FALSE;
}

BOOL DeleteKeyRecursive(CRegNode *pNode)
{
	CRegNode *pTemp;

	if (HasSubKeys(pNode)) // has sub keys
	{
		pTemp = pNode->m_pFirstChild;
		while (pTemp)
		{
			if (pTemp->m_dwType & REGNODE_TYPE_KEY)
				DeleteKeyRecursive(pTemp);
			pTemp = pTemp->m_pNextSibling;
		}
	}
	{
		// delete the key
		CString strTemp;
		strTemp = pNode->GetKeyPath();
		LONG lErr = RegDeleteKey(pNode->GetKeyRoot(), (LPCTSTR)strTemp);
		if (lErr != ERROR_SUCCESS)
		{
			strTemp.Format(_T("Delete failed (%s), Err %ld"), (LPCTSTR)pNode->m_strName, lErr);
			AfxMessageBox(strTemp);
		}
	}
	return TRUE;
}

void LoadKeysRecursive(CRegNode *pNode, BOOL bLoadValues)
{
	// load all subkeys and their subkeys ...
	CRegNode *pTemp;
	if (pNode->m_dwType & REGNODE_TYPE_KEY)
	{
		LoadKey(pNode);
		if (bLoadValues)
			LoadValues(pNode);
	}
	pTemp = pNode->m_pFirstChild;
	while (pTemp)
	{
		if (pTemp->m_dwType & REGNODE_TYPE_KEY)
			LoadKeysRecursive(pTemp, bLoadValues);
		pTemp = pTemp->m_pNextSibling;
	}
}

void CreateKeyRecursive(CRegNode *pNode, BOOL bCreateValues)
{
	CRegNode *pTemp;
	if (pNode->m_dwType & REGNODE_TYPE_KEY)
	{
		CreateKey(pNode);
	}
	pTemp = pNode->m_pFirstChild;
	while (pTemp)
	{
		if (pTemp->m_dwType & REGNODE_TYPE_KEY)
			CreateKeyRecursive(pTemp, bCreateValues);
		if (pTemp->m_dwType & REGNODE_TYPE_VALUE && bCreateValues)
			CreateValue(pTemp);
		pTemp = pTemp->m_pNextSibling;
	}
}

BOOL CreateKey(CRegNode *pNode)
{
	if (!pNode->m_pParent)
		return FALSE;
	HKEY hKeyParent = pNode->m_pParent->GetKeyHandle(), hNewKey;
	DWORD dwDisposition;
	LONG lErr;

	lErr = RegCreateKeyEx(hKeyParent, pNode->m_strName, 0, _T(""), REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, NULL, &hNewKey, &dwDisposition);
	if (lErr == ERROR_SUCCESS)
	{
		RegCloseKey(hNewKey);
		RegCloseKey(hKeyParent);
		return TRUE;
	}
	RegCloseKey(hKeyParent);
	return FALSE;
}

BOOL CreateValue(CRegNode *pNode)
{
	LONG lRet;

	if (!pNode  || !pNode->m_pParent)
		return FALSE;

	HKEY hKey = pNode->m_pParent->GetKeyHandle();

	if (hKey)
	{
		lRet = RegSetValueEx(hKey, (LPCTSTR)pNode->m_strName, 0, pNode->m_dwDataType, (LPBYTE)pNode->m_lpData, pNode->m_dwDataSize);
		if (lRet == ERROR_SUCCESS)
		{
			RegCloseKey(hKey);
			return TRUE;
		}
		RegCloseKey(hKey);
	}

	return FALSE;
}