// EncryptActions.cpp: implementation of the CEncryptActions class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "EncryptActions.h"

#include "resource.h"
#include <appmisc.h>


//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

bool CEncryptActions::EncryptPrepare::Do(CActionEngine* pEngine, 
	  								     HWND hwndUpdates)
{
	ASSERT(pEngine != NULL);

	Update(hwndUpdates, _T("Listing files to encrypt..."), NULL);

	// Iterate through source folders
	string_array::const_iterator iterSrc;
	for(iterSrc = m_aFiles.begin(); iterSrc != m_aFiles.end(); iterSrc++)
	{
		Update(hwndUpdates, NULL, (LPCTSTR)*iterSrc);

		file_iterator fileIter((file_path)*iterSrc, file_iterator::sub_folders | file_iterator::full_paths), end;
		for( ; fileIter != end; fileIter++)
		{
			file_path srcFile(fileIter->cFileName);
			if(!(srcFile.attributes() & FILE_ATTRIBUTE_DIRECTORY))
			{
				if(m_asIgnore.find(srcFile.ext().make_lower()) == m_asIgnore.end())
				{
					// Construct File
					Encrypt* pEncrypt = new Encrypt(srcFile);
					pEngine->Add(pEncrypt);
				}
			}
		}
		
		// Error from file_iterator (listing folders)
		if(fileIter.state() != ERROR_SUCCESS)
			throw CActionError(HRESULT_FROM_WIN32(fileIter.state()), IDS_ENCRYPT_ERR_LIST, (file_path)*iterSrc);
	}

	return false;
}

HRESULT CEncryptActions::EncryptPrepare::Initialize(const CPropertyBag& data)
{
	if(!data.GetStringSet(NS_ENCRYPT_REG_FILES, m_aFiles))
		return S_FALSE;

	CEncryptActions::Encrypt::m_bEncryptReadOnly = data.GetInt(NS_ENCRYPT_REG_READONLY, false) ? true : false;

	string_array asTemp;
	data.GetStringSet(NS_IGNORE_REG_EXT, asTemp);

	// Make them all lower
	for_each(asTemp.begin(), asTemp.end(), Lower);

	// And copy them to the set (removes doubles)
	copy(asTemp.begin(), asTemp.end(), inserter(m_asIgnore, m_asIgnore.begin()));

	// Remove any blanks from ignore
	m_asIgnore.erase(_T(""));

	return S_OK;
}


IXpyExPtr CEncryptActions::Encrypt::m_pXpy;
bool CEncryptActions::Encrypt::m_bEncryptReadOnly = false;
UINT CEncryptActions::Encrypt::m_nXpyRefs = 0;

bool CEncryptActions::Encrypt::Do(CActionEngine* pEngine, 
	  								 HWND hwndUpdates)
{
	// Update any monitors
	Update(hwndUpdates, _T("Encrypting..."), m_file);

	// Initialize XPY
	HRESULT hr;
	if(FAILED(hr = InitXpy()))
	{
		// If it doesn't work then trash all the encryption actions 
		// in the engine
		if(pEngine)
			pEngine->RemoveAll(this);

		// And spit out an error
		throw CActionError(hr, IDS_ENCRYPT_ERR_XPYINIT);
	}


	// Use this in case we need to reset the attributes
	// in order to encrypt the file. XPY won't encrypt
	// read-only files. There's an option here 
	// to let you. So basically what we have to do is 
	// reset the file to normal attributes and then set 
	// it back

	// If it's 0 that means we haven't messed with the 
	// attributes yet. If it has a value that means we 
	// need to set the value back before returning
	DWORD dwAttributes = 0;

encrypt:

	try
	{
		// Here's the actual close code
		// rest is basically error checking
		m_pXpy->Close((LPCTSTR)m_file);

	}
	catch(_com_error& e)
	{
		// Here we check for the read-only error that Xpy 
		// will spit up. If the flag is set then go ahead
		// and reset the attributes

		if(e.Error() == XPY_E_FILE_ATTRIBUTES && // Read Only Error
		   m_bEncryptReadOnly)					 // flag to encrypt read-only
		{

			// If we haven't messed with the attributes 
			// yet then go ahead and reset them
			if(!dwAttributes)
			{
				// Get Attributes
				dwAttributes = GetFileAttributes(m_file);

				// Let System files blow! (with an error)
				if(!(dwAttributes & FILE_ATTRIBUTE_SYSTEM))
				{
					SetFileAttributes(m_file, FILE_ATTRIBUTE_NORMAL);
				
					// Reencrypt
					goto encrypt;
				}

			}

			// Otherwise we won't try again
		}

		// Set the Attributes back
		if(dwAttributes)
			SetFileAttributes(m_file, dwAttributes);

		// Get the proper error message
		_bstr_t bsError = e.Description();

		if(bsError.length())
			throw CActionError((LPCTSTR)bsError, e.Error(), IDS_ENCRYPT_ERR, m_file);
		else
			throw CActionError(e.Error(), IDS_ENCRYPT_ERR, m_file);
	}

	// Set the Attributes back
	if(dwAttributes)
		SetFileAttributes(m_file, dwAttributes);

	return false;
}

bool CEncryptActions::Encrypt::IsFixable(HRESULT hr) const
{
	switch(hr)
	{
	case XPY_E_FILE_ATTRIBUTES:
	case HRESULT_FROM_WIN32(ERROR_SHARING_VIOLATION):
	case HRESULT_FROM_WIN32(ERROR_LOCK_VIOLATION):
		return true;
	default:
		return false;
	}
}

void CEncryptActions::Encrypt::Fix(HRESULT hr)
{
	if(hr == XPY_E_FILE_ATTRIBUTES)
	{
		if(SetFileAttributes(m_file, FILE_ATTRIBUTE_NORMAL))
			Do(NULL, NULL);
		else
			throw CActionError(HRESULT_FROM_WIN32(::GetLastError()), IDS_ENCRYPT_ERR, m_file);
	}
	else if(hr == HRESULT_FROM_WIN32(ERROR_SHARING_VIOLATION) || 
		    hr == HRESULT_FROM_WIN32(ERROR_LOCK_VIOLATION))
	{
		bool bFirst = true;
start:
		try
		{
			// Try again
			Do(NULL, NULL);
			return;
		}
		catch(CActionError& err)
		{
			// If it's the first around and the error is the same...
			if(bFirst && IsFixable(err.m_hRes))
			{
				HCURSOR hCurs = ::GetCursor();

				// ... Prompt user to close programs
				MessageBox(NULL, _T("Files cannot be encrypted while they are opened in a program. Close all other programs and then click the OK button."), _T("Encrypt Files"), MB_OK | MB_ICONINFORMATION);

				if(hCurs)
					::SetCursor(hCurs);

				bFirst = false;
			}
			// Otherwise out with the error
			else
				throw;
		}

		goto start;
	}
}

HRESULT CEncryptActions::Encrypt::InitXpy()
{
	if(m_pXpy)
		return S_OK;
	else 
		return m_pXpy.CreateInstance(__uuidof(XpyEx));
}