
#include "stdafx.h"
#include "CustomUninstallKey.h"
#include "RegistryKey.h"
#include "RegDataSerializer.h"
#include <strutil.h>

#define CustomUninstallKey_ExecCA L"CustomUninstallKey_deferred"
#define CustomUninstallKey_RollbackCA L"CustomUninstallKey_rollback"
#define CustomUninstallKeyQuery L"SELECT `Id`, `Name`, `Data`, `DataType`, `Attributes`, `Condition` FROM `PSW_CustomUninstallKey`"
enum eCustomUninstallKeyQuery { Id = 1, Name, Data, DataType, Attributes, Condition };

extern "C" __declspec( dllexport ) UINT CustomUninstallKey_Immediate(MSIHANDLE hInstall)
{
	HRESULT hr = S_OK;
	UINT er = ERROR_SUCCESS;
	CCustomUninstallKey data(hInstall);

	hr = WcaInitialize(hInstall, "CustomUninstallKey_Immediate");
	BreakExitOnFailure(hr, "Failed to initialize");
	WcaLog(LOGMSG_STANDARD, "Initialized.");
	
	// Ensure table PSW_CustomUninstallKey exists.
	hr = WcaTableExists(L"PSW_CustomUninstallKey");
	BreakExitOnFailure(hr, "Table does not exist 'PSW_CustomUninstallKey'. Have you authored 'PanelSw:CustomUninstallKey' entries in WiX code?");

	// TODO: Add your custom action code here.
	hr = data.CreateCustomActionData();
	BreakExitOnFailure(hr, "Failed");

LExit:
	er = SUCCEEDED(hr) ? ERROR_SUCCESS : ERROR_INSTALL_FAILURE;
	return WcaFinalize(er);
}

extern "C" __declspec( dllexport ) UINT CustomUninstallKey_deferred(MSIHANDLE hInstall)
{
	HRESULT hr = S_OK;
	UINT er = ERROR_SUCCESS;
	CCustomUninstallKey data(hInstall);

	hr = WcaInitialize(hInstall, "CustomUninstallKey_deferred");
	BreakExitOnFailure(hr, "Failed to initialize");
	WcaLog(LOGMSG_STANDARD, "Initialized.");

	hr = data.Execute();
	BreakExitOnFailure(hr, "Failed");

LExit:
	er = SUCCEEDED(hr) ? ERROR_SUCCESS : ERROR_INSTALL_FAILURE;
	return WcaFinalize(er);
}

CCustomUninstallKey::CCustomUninstallKey( MSIHANDLE hInstall)
	: _hInstall( hInstall)
{
}

CCustomUninstallKey::~CCustomUninstallKey(void)
{
}

HRESULT CCustomUninstallKey::Execute()
{
	CRegistryXmlParser parser;
	WCHAR* customActionData = NULL;
	HRESULT hr = S_OK;

	hr = WcaGetProperty( L"CustomActionData", &customActionData);
	BreakExitOnFailure( hr, "Failed to get CustomActionData");
	WcaLog( LOGLEVEL::LOGMSG_STANDARD, "CustomActionData= '%ls'", customActionData);

	hr = parser.Execute( customActionData);
	BreakExitOnFailure( hr, "Failed to get parse-execute CustomActionData");

LExit:
	return hr;
}

HRESULT CCustomUninstallKey::CreateCustomActionData()
{
	PMSIHANDLE hView;
    PMSIHANDLE hRec;
	CRegistryXmlParser xmlParser, xmlRollbackParser;
	WCHAR uninstallKey[ MAX_PATH];
	CComBSTR xmlString = L"";
	HRESULT hr = S_OK;
	bool bRemoving = false;

	hr = GetUninstallKey( uninstallKey);
	BreakExitOnFailure(hr, "Failed to get uninstall registry key");
	
	hr = WcaOpenExecuteView( CustomUninstallKeyQuery, &hView);
	BreakExitOnFailure(hr, "Failed to execute view");

	if (WcaIsPropertySet("REMOVE"))
	{
		bRemoving = true;
	}

	while ( E_NOMOREITEMS != (hr = WcaFetchRecord(hView, &hRec)))
    {
		BreakExitOnFailure(hr, "Failed to fetch record");
		
		// Get record.
		WCHAR* pId = NULL;
		WCHAR* pName = NULL;
		WCHAR* pData = NULL;
		WCHAR* pDataType = NULL;
		WCHAR* pCondition = NULL;
		int nAttrib;
		
		hr = WcaGetRecordString( hRec, eCustomUninstallKeyQuery::Id, &pId);
		BreakExitOnFailure(hr, "Failed to get Id");
		hr = WcaGetRecordString( hRec, eCustomUninstallKeyQuery::Name, &pName);
		BreakExitOnFailure(hr, "Failed to get Name");
		hr = WcaGetRecordString( hRec, eCustomUninstallKeyQuery::Data, &pData);
		BreakExitOnFailure(hr, "Failed to get Data");
		hr = WcaGetRecordString( hRec, eCustomUninstallKeyQuery::DataType, &pDataType);
		BreakExitOnFailure(hr, "Failed to get DataType");
		hr = WcaGetRecordInteger( hRec, eCustomUninstallKeyQuery::Attributes, &nAttrib);
		BreakExitOnFailure(hr, "Failed to get Attributes");
		hr = WcaGetRecordString( hRec, eCustomUninstallKeyQuery::Condition, &pCondition);
		BreakExitOnFailure(hr, "Failed to get Condition");

		MSICONDITION cond = ::MsiEvaluateCondition( _hInstall, pCondition);
		switch( cond)
		{
		case MSICONDITION::MSICONDITION_NONE:
		case MSICONDITION::MSICONDITION_TRUE:
			WcaLog( LOGLEVEL::LOGMSG_STANDARD, "Condition evaluated true/none for %ls", pId);
			break;

		case MSICONDITION::MSICONDITION_FALSE:
			WcaLog( LOGLEVEL::LOGMSG_STANDARD, "Condition evaluated false for %ls", pId);
			continue;

		case MSICONDITION::MSICONDITION_ERROR:
			hr = E_FAIL;
			BreakExitOnFailure(hr, "Failed to evaluate condition");
		}

		// Install / UnInstall ?
		if (bRemoving)
		{
			hr = xmlParser.AddDeleteKey(pId, CRegistryKey::RegRoot::LocalMachine, uninstallKey, CRegistryKey::RegArea::Default);
			BreakExitOnFailure(hr, "Failed to create XML element");
		}
		else
		{
			CRegDataSerializer dataSer;
			WCHAR* pDataStr = NULL;

			hr = dataSer.Set(pData, pDataType);
			BreakExitOnFailure(hr, "Failed to create parse registry data");

			hr = dataSer.Serialize(&pDataStr);
			BreakExitOnFailure(hr, "Failed to serialize registry data");

			hr = xmlParser.AddCreateValue(pId, CRegistryKey::RegRoot::LocalMachine, uninstallKey, CRegistryKey::RegArea::Default, pName, (CRegistryKey::RegValueType)dataSer.DataType(), pDataStr);
			BreakExitOnFailure(hr, "Failed to create XML element");
		}

		hr = CreateRollbackCustomActionData( &xmlRollbackParser, pId, pName);
		BreakExitOnFailure(hr, "Failed to create rollback XML element");
	}
	
	hr = xmlRollbackParser.GetXmlString( &xmlString);
	BreakExitOnFailure(hr, "Failed to read XML as text");
	hr = WcaDoDeferredAction( CustomUninstallKey_RollbackCA, xmlString, 0);
	BreakExitOnFailure(hr, "Failed to set property");	

	hr = xmlParser.GetXmlString( &xmlString);
	BreakExitOnFailure(hr, "Failed to read XML as text");
	hr = WcaDoDeferredAction( CustomUninstallKey_ExecCA, xmlString, 0);
	BreakExitOnFailure(hr, "Failed to set property");
	xmlString = std::nullptr_t(NULL);

LExit:

	return hr;
}

HRESULT CCustomUninstallKey::CreateRollbackCustomActionData( CRegistryXmlParser *pRollbackParser, WCHAR* pId, WCHAR* pName)
{
	HRESULT hr = S_OK;
	CRegistryKey key;
	CRegistryKey::RegValueType valType;
	BYTE* pDataBytes = NULL;
	WCHAR* pDataStr = NULL;
	DWORD dwSize = 0;
	WCHAR keyName[ MAX_PATH];
	CRegDataSerializer dataSer;

	hr = GetUninstallKey( keyName);
	BreakExitOnFailure( hr, "Failed to get Uninstall key");

	hr = key.Open( CRegistryKey::RegRoot::LocalMachine, keyName, CRegistryKey::RegArea::Default, CRegistryKey::RegAccess::ReadOnly);
	if( hr == E_FILENOTFOUND)
	{
		// Delete key on rollback
		hr = pRollbackParser->AddDeleteKey(pId, CRegistryKey::RegRoot::LocalMachine, keyName, CRegistryKey::RegArea::Default);
		BreakExitOnFailure( hr, "Failed to create rollback XML element");
		ExitFunction();
	}

	hr = key.GetValue(pName, &pDataBytes, &valType, &dwSize);
	if( hr == E_FILENOTFOUND)
	{
		// Delete value on rollback
		hr = pRollbackParser->AddDeleteValue(pId, CRegistryKey::RegRoot::LocalMachine, keyName, CRegistryKey::RegArea::Default, pName);
		BreakExitOnFailure( hr, "Failed to create rollback XML element");
		ExitFunction();
	}
	
	hr = dataSer.Set(pDataBytes, valType, dwSize);
	BreakExitOnFailure(hr, "Failed to analyze registry data");

	hr = dataSer.Serialize(&pDataStr);
	BreakExitOnFailure(hr, "Failed to serialize registry data");

	hr = pRollbackParser->AddCreateValue(pId, CRegistryKey::RegRoot::LocalMachine, keyName, CRegistryKey::RegArea::Default, pName, valType, pDataStr);
	BreakExitOnFailure( hr, "Failed to create rollback XML element");

LExit:
	if (pDataBytes != NULL)
	{
		delete[] pDataBytes;
	}

	return hr;
}

HRESULT CCustomUninstallKey::GetUninstallKey( WCHAR* keyName)
{
	WCHAR *prodCode = NULL;
	HRESULT hr = S_OK;

	hr = WcaGetProperty( L"ProductCode", &prodCode);
	BreakExitOnFailure( hr, "Failed to get ProductCode");

	wsprintfW( keyName, L"Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall\\%ls", prodCode);

LExit:
	return hr;
}
