// RegistryOps.cpp : Implementation of CRegistryOps

#include "stdafx.h"
#include "RegistryOps.h"
#include ".\registryops.h"
#include <string>
using namespace std ;

// CRegistryOps


HKEY CRegistryOps::getRootHandle(std::string rootString)
{
    if(rootString == "HKEY_LOCAL_MACHINE")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000002)) ;
	else if(rootString == "HKEY_CURRENT_USER")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000001)) ;
	else if(rootString == "HKEY_USERS")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000003)) ;
	else if(rootString == "HKEY_CLASSES_ROOT")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000000)) ;
	else if(rootString == "HKEY_CURRENT_CONFIG")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000005)) ;
	else if(rootString == "HKEY_DYN_DATA")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000006)) ;
	else if(rootString == "HKEY_PERFORMANCE_DATA")
		return ((HKEY)(ULONG_PTR)((LONG)0x80000004)) ;
	else
		return (HKEY)(NULL) ; ;
}

STDMETHODIMP CRegistryOps::saveHive(BSTR fileName, BSTR subKey, BSTR rootKey, BOOL * retCode)
{
	TOKEN_PRIVILEGES tp ;
	HANDLE hToken ;
	LUID luid ;
	//LONG rc ;
	HKEY hKeyRes ;
	LPTSTR MachineName = NULL ;
	DWORD dwSubKeyIndex = 0 ;
	//DWORD dwDisposition ;
	//char szSubKey[_MAX_FNAME] ;
	DWORD dwSubKeyLength = _MAX_FNAME ;

	try
	{
		if ( !OpenProcessToken (GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken) )
		{
			* retCode = FALSE ; 
			std::cout << "ERROR open process token:" << GetLastError () ;
			return S_OK ;
		}
		
		if ( !LookupPrivilegeValue (MachineName, SE_BACKUP_NAME, &luid) )
		{
			CloseHandle(hToken) ;
			* retCode = FALSE; 
			std::cout << "ERROR lookup privilege:" << GetLastError () ;
			return S_OK ;
		}

		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

		AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof (TOKEN_PRIVILEGES), 0, 0);

		if (GetLastError() != ERROR_SUCCESS)
		{
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR adj token privilege:" << GetLastError () ;
			return S_OK ;
		}

		USES_CONVERSION ;
		std::string subKeyName = OLE2A (subKey) ;
		std::string file = OLE2A (fileName) ;
		std::string rootString = OLE2A (rootKey) ;

		if(rootString == subKeyName)
		{
			if (RegOpenKeyEx (getRootHandle(rootString), NULL, 0, KEY_ALL_ACCESS, &hKeyRes) != ERROR_SUCCESS)			
			{
				CloseHandle(hToken);
				RegCloseKey (hKeyRes);
				* retCode = FALSE; 
				std::cout << "ERROR Open:" << GetLastError () ;
				return S_OK ;
			}
		}
		else if (RegOpenKeyEx (getRootHandle(rootString), subKeyName.c_str(), 0, KEY_ALL_ACCESS, &hKeyRes) != ERROR_SUCCESS)			
			{
				CloseHandle(hToken);
				RegCloseKey (hKeyRes);
				* retCode = FALSE; 
				std::cout << "ERROR Open:" << GetLastError () ;
				return S_OK ;
			}

		if (RegSaveKey (hKeyRes, file.c_str(), 0) != ERROR_SUCCESS)
		{
			CloseHandle(hToken);
			RegCloseKey (hKeyRes);
			* retCode = FALSE; 
			std::cout << "ERROR save:" << GetLastError () ;
			return S_OK ;
		}

		RegCloseKey (hKeyRes);
		AdjustTokenPrivileges(hToken, TRUE, 0, 0, 0, 0);
		CloseHandle(hToken);
	}
	catch (...)
	{
		* retCode = FALSE; 
		std::cout << "ERROR exception:" << GetLastError () ;
		return S_OK ;
	}

	* retCode = TRUE; // success
	return S_OK;
}

STDMETHODIMP CRegistryOps::restoreHive(BSTR fileName, BSTR subKey, BSTR rootKey, BOOL * retCode)
{
	HKEY hKeyRes;
	TOKEN_PRIVILEGES tp;
	HANDLE hToken;
	LUID luid;
//	LONG rc;
	LPTSTR MachineName = NULL;
	DWORD dwSubKeyIndex = 0;
//	char szSubKey [_MAX_FNAME];
	DWORD dwSubKeyLength = _MAX_FNAME;

	try
	{
	
		if(!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken))
		{
			* retCode = FALSE; 
			std::cout << "ERROR open process token:" << GetLastError () ;
			return S_OK ;
		}

		if(!LookupPrivilegeValue(MachineName, SE_BACKUP_NAME , &luid))
		{
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR privilege:" << GetLastError () ;
			return S_OK ;
		}

		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

		AdjustTokenPrivileges(hToken, false, &tp, sizeof (TOKEN_PRIVILEGES), 0, 0);

		if (GetLastError() != ERROR_SUCCESS)
		{
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR adj token privelege:" << GetLastError () ;
			return S_OK ;
		}

		if(!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken))
		{
			* retCode = FALSE; 
			std::cout << "ERROR adj open process token:" << GetLastError () ;
			return S_OK ;
		}


		if(!LookupPrivilegeValue(MachineName, SE_RESTORE_NAME, &luid))
		{
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR 2nd lookup privilege:" << GetLastError () ;
			return S_OK ;
		}

		tp.PrivilegeCount = 1;
		tp.Privileges[0].Luid = luid;
		tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

		AdjustTokenPrivileges(hToken, false, &tp, sizeof (TOKEN_PRIVILEGES), 0, 0);

		if (GetLastError() != ERROR_SUCCESS)
		{
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR 2nd adj token privilege:" << GetLastError () ;
			return S_OK ;
		}

		USES_CONVERSION ;
		std::string subKeyName = OLE2A (subKey) ;
		std::string file = OLE2A (fileName) ;
		std::string rootString = OLE2A (rootKey) ;

		if(rootString == subKeyName)
		{
			if (RegOpenKeyEx (getRootHandle(rootString), NULL, 0, KEY_ALL_ACCESS, &hKeyRes) != ERROR_SUCCESS)
			{
				RegCloseKey (hKeyRes);
				CloseHandle(hToken);
				* retCode = FALSE; 
				std::cout << "ERROR : OPEN :" << GetLastError () ;
				return S_OK ;
			}
		}
		else if (RegOpenKeyEx (getRootHandle(rootString), subKeyName.c_str(), 0, KEY_ALL_ACCESS, &hKeyRes) != ERROR_SUCCESS)
		{
			RegCloseKey (hKeyRes);
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR : OPEN :" << GetLastError () ;
			return S_OK ;
		}

		if (RegRestoreKey (hKeyRes, file.c_str(), 0) != ERROR_SUCCESS)
		{
			RegCloseKey (hKeyRes);
			CloseHandle(hToken);
			* retCode = FALSE; 
			std::cout << "ERROR : RESTORE :" << GetLastError () ;
			return S_OK ;
		}

		RegCloseKey (hKeyRes);
		CloseHandle(hToken);
	}
	catch (...)
	{
		* retCode = FALSE; 
		std::cout << "ERROR exception:" << GetLastError () ;
			return S_OK ;
	}

	* retCode = TRUE; 
	return S_OK ;
}