//---------------------------------
// Name:			MainDialog.cpp
// Purpose:			Firmware packager main dialog class
// Toolchain(s):	Microsoft Visual C++
// Creation date:	02-Jun-06
// Created by:		JJo
//---------------------------------

//---------------------------------
// Includes

#include "MainDialog.h"

//---------------------------------
// Message map
BEGIN_MESSAGE_MAP(MainDialog, CDialog)
	ON_BN_CLICKED(IDC_EXIT, OnExitClicked)
	ON_BN_CLICKED(IDC_INPUTBROWSE, OnInputBrowseClicked)
	ON_BN_CLICKED(IDC_OUTPUTBROWSE, OnOutputBrowseClicked)
	ON_BN_CLICKED(IDC_PACKAGE, OnPackageClicked)
	ON_BN_CLICKED(IDC_AES, OnAESClicked)
	ON_BN_CLICKED(IDC_3DES, On3DESClicked)
	ON_BN_CLICKED(IDC_ECB, OnECBClicked)
	ON_BN_CLICKED(IDC_CBC, OnCBCClicked)
	ON_BN_CLICKED(IDC_CTR, OnCBCClicked)
	ON_BN_CLICKED(IDC_EAX, OnCBCClicked)
	ON_EN_UPDATE(IDC_ENCRYPTIONKEY, OnKeyUpdated)
	ON_EN_UPDATE(IDC_IV, OnIVUpdated)
	ON_WM_CLOSE()
END_MESSAGE_MAP()

//---------------------------------
// Private methods

/**
 * Name:		IsECBSelected
 * Purpose:		Indicates whether ECB mode is selected or not
 * Output:		TRUE if ECB mode is selected, false otherwise
 */
BOOL MainDialog::IsECBSelected() {

	if (((CButton *) this->GetDlgItem(IDC_ECB))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsCBCSelected
 * Purpose:		Indicates whether CBC mode is selected or not
 * Output:		TRUE if CBC mode is selected, false otherwise
 */
BOOL MainDialog::IsCBCSelected() {

	if (((CButton *) this->GetDlgItem(IDC_CBC))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsCTRSelected
 * Purpose:		Indicates whether CTR mode is selected or not
 * Output:		TRUE if CTR mode is selected, false otherwise
 */
BOOL MainDialog::IsCTRSelected() {

	if (((CButton *) this->GetDlgItem(IDC_CTR))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsEAXSelected
 * Purpose:		Indicates whether EAX mode is selected or not
 * Output:		TRUE if EAX mode is selected, false otherwise
 */
BOOL MainDialog::IsEAXSelected() {

	if (((CButton *) this->GetDlgItem(IDC_EAX))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsAESSelected
 * Purpose:		Indicates whether or not AES is currently selected
 * Output:		TRUE if AES is selected, FALSE otherwise
 */
BOOL MainDialog::IsAESSelected() {

	if (((CButton *) this->GetDlgItem(IDC_AES))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		Is3DESSelected
 * Purpose:		Indicates whether or not 3DES is currently selected
 * Output:		TRUE if 3DES is selected, FALSE otherwise
 */
BOOL MainDialog::Is3DESSelected() {

	if (((CButton *) this->GetDlgItem(IDC_3DES))->GetCheck() == 1) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsInputFileSelected
 * Purpose:		Indicates whether or not an input file is currently selected
 * Output:		TRUE if an input file is selected, FALSE otherwise
 */
BOOL MainDialog::IsInputFileSelected() {

	if (((CComboBox *) this->GetDlgItem(IDC_INPUTFILE))->GetCount() > 0) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		IsOutputFileSelected
 * Purpose:		Indicates whether or not an output file is currently selected
 * Output:		TRUE if an output file is selected, FALSE otherwise
 */
BOOL MainDialog::IsOutputFileSelected() {

	if (((CComboBox *) this->GetDlgItem(IDC_OUTPUTFILE))->GetCount() > 0) {

		return TRUE;
	}
	else {

		return FALSE;
	}
}

/**
 * Name:		GetKey
 * Purpose:		Returns the current key
 * Output:		The current key
 */
CString MainDialog::GetKey() {

	// Local variables
	CString key;

	((CEdit *) this->GetDlgItem(IDC_ENCRYPTIONKEY))->GetWindowText(key);
	return key;
}

/**
 * Name:		GetIV
 * Purpose:		Returns the current IV
 * Output:		The current IV
 */
CString MainDialog::GetIV() {

	// Local variables
	CString key;

	((CEdit *) this->GetDlgItem(IDC_IV))->GetWindowText(key);
	return key;
}

/**
 * Name:		GetInputFilePath
 * Purpose:		Returns the currently selected input file path
 * Output:		Path of currently selected input file
 */
CString MainDialog::GetInputFilePath() {

	// Local variables
	CComboBox * combo;

	combo = (CComboBox *) this->GetDlgItem(IDC_INPUTFILE);
	return CString((char *) combo->GetItemDataPtr(combo->GetCurSel()));
}

/**
 * Name:		GetOutputFilePath
 * Purpose:		Returns the currently selected output file path
 * Output:		Path of currently selected output file
 */
CString MainDialog::GetOutputFilePath() {

	// Local variables
	CComboBox * combo;

	combo = (CComboBox *) this->GetDlgItem(IDC_OUTPUTFILE);
	return CString((char *) combo->GetItemDataPtr(combo->GetCurSel()));
}

/**
 * Name:		GetOutputFileName
 * Purpose:		Returns the currently selected output file name
 * Output:		Name of currently selected file
 */
CString MainDialog::GetOutputFileName() {

	// Local variables
	CString name;

	((CComboBox *) this->GetDlgItem(IDC_OUTPUTFILE))->GetWindowText(name);
	return name;
}

/**
 * Name:		GetKeySel
 * Purpose:		Returns the current selection in the key edit box
 * Output:		Current key selection
 */
DWORD MainDialog::GetKeySel() {

	return ((CEdit *) this->GetDlgItem(IDC_ENCRYPTIONKEY))->GetSel();
}

/**
 * Name:		GetIVSel
 * Purpose:		Returns the current selection in the IV edit box
 * Output:		Current IV selection
 */
DWORD MainDialog::GetIVSel() {

	return ((CEdit *) this->GetDlgItem(IDC_IV))->GetSel();
}

/**
 * Name:		SetKey
 * Purpose:		Sets the key value
 * Input(s):
 *  - Key value
 */
void MainDialog::SetKey(const char * value) {

	((CEdit *) this->GetDlgItem(IDC_ENCRYPTIONKEY))->SetWindowText(value);
}

/**
 * Name:		SetIV
 * Purpose:		Sets the IV value
 * Input(s):
 *  - IV value
 */
void MainDialog::SetIV(const char * value) {

	((CEdit *) this->GetDlgItem(IDC_IV))->SetWindowText(value);
}

/**
 * Name:		SetKeySel
 * Purpose:		Sets the selection for the key edit field
 * Input(s):
 *  - Start and end of selection
 */
void MainDialog::SetKeySel(DWORD value) {

	((CEdit *) this->GetDlgItem(IDC_ENCRYPTIONKEY))->SetSel(value);
}

/**
 * Name:		SetIVSel
 * Purpose:		Sets the selection for the IV edit field
 * Input(s):
 *  - Start and end of selection
 */
void MainDialog::SetIVSel(DWORD value) {

	((CEdit *) this->GetDlgItem(IDC_IV))->SetSel(value);
}

/**
 * Name:		AddInputFile
 * Purpose:		Adds an input file name & path to the input file combobox
 * Input(s):
 *	- File name
 *  - Path to file
 */
void MainDialog::AddInputFile(const char * fileName, const char * filePath) {

	// Local variables
	CComboBox * combo;

	combo = (CComboBox *) this->GetDlgItem(IDC_INPUTFILE);
	combo->InsertString(0, fileName);
	combo->SetCurSel(0);
	combo->SetItemDataPtr(0, (void *) (LPCTSTR (* new CString(filePath))));

}

/**
 * Name:		AddOutputFile
 * Purpose:		Adds an output file name & path to the output file combobox
 * Input(s):
 *	- File name
 *  - Path to file
 */
void MainDialog::AddOutputFile(const char * fileName, const char * filePath) {

	// Local variables
	CComboBox * combo;

	combo = (CComboBox *) this->GetDlgItem(IDC_OUTPUTFILE);
	combo->InsertString(0, fileName);
	combo->SetCurSel(0);
	combo->SetItemDataPtr(0, (void *) (LPCTSTR (* new CString(filePath))));
}

/**
 * Name:		CheckAES
 * Purpose:		Selects the AES radio button
 */
void MainDialog::CheckAES() {

	((CButton *) this->GetDlgItem(IDC_AES))->SetCheck(1);
}

/**
 * Name:		CheckECB
 * Purpose:		Selects the ECB radio button
 */
void MainDialog::CheckECB() {

	((CButton *) this->GetDlgItem(IDC_ECB))->SetCheck(1);
}

/**
 * Name:		EnablePackageButton
 * Purpose:		Enable package button if settings are entered
 */
void MainDialog::EnablePackageButton() {

	// Local variables
	CString key;
	CString IV;

	// Check values
	key = this->GetKey();
	IV = this->GetIV();

	if (this->IsInputFileSelected() &&
		this->IsOutputFileSelected() &&
		(key.GetLength() == (this->m_keyLength*2)) &&
		((IV.GetLength() == (this->m_blockLength*2) || this->IsECBSelected()))) {

		((CButton *) this->GetDlgItem(IDC_PACKAGE))->EnableWindow(TRUE);
	}
	else {

		((CButton *) this->GetDlgItem(IDC_PACKAGE))->EnableWindow(FALSE);
	}
}

/**
 * Name:		EnableIV
 * Purpose:		Enables the IV edit field
 */
void MainDialog::EnableIV() {

	((CEdit *) this->GetDlgItem(IDC_IV))->EnableWindow(TRUE);
}

/**
 * Name:		DisableIV
 * Purpose:		Disables the IV edit field
 */
void MainDialog::DisableIV() {

	((CEdit *) this->GetDlgItem(IDC_IV))->EnableWindow(FALSE);
}

/**
 * Name:		ASCII2Hex
 * Purpose:		Converts an ASCII string to an hexadecimal value
 * Input(s):
 *  - ASCII string to convert
 *  - Buffer to store conversion
 *  - Length of buffer
 */
void MainDialog::ASCII2Hex(const char * ascii, unsigned char * hex, unsigned int length) {

	for (unsigned int i=0; i < length; i++, ascii++, hex++) {

		if (*ascii >= 'A' && *ascii <= 'F') {

			*hex = *ascii - 'A' + 10;
		}
		else {

			*hex = *ascii - '0';
		}

		*hex <<= 4;
		ascii++;

		if (*ascii >= 'A' && *ascii <= 'F') {

			*hex += *ascii - 'A' + 10;
		}
		else {

			*hex += *ascii - '0';
		}
	}
}

//---------------------------------
// Public methods

/**
 * Name:		MainDialog
 * Purpose:		Constructor for the MainDialog class
 */
MainDialog::MainDialog() {

	this->Create(IDD_MAINDIALOG, NULL);

	// Register all potentially needed ciphers
	if (register_cipher(&rijndael_desc) == -1) {

		throw new CException();
	}

	if (register_cipher(&des3_desc) == -1) {

		throw new CException();
	}
}

/**
 * Name:		~MainDialog
 * Purpose:		Destructor for the MainDialog class
 */
MainDialog::~MainDialog() {}

/**
 * Name:		OnInitDialog
 * Purpose:		Handles dialog initialization
 */
BOOL MainDialog::OnInitDialog() {

	// Call superclass method
	CDialog::OnInitDialog();

	// Setup controls
	this->CheckAES();
	this->CheckECB();
	this->DisableIV();
	this->m_blockLength = 16;
	this->m_keyLength = 16;

	this->EnablePackageButton();

	return TRUE;
}

/**
 * Name:		OnClose
 * Purpose:		Handles window closing
 */
void MainDialog::OnClose() {

	this->DestroyWindow();
}

/**
 * Name:		OnExitClicked
 * Purpose:		Handles clicks on exit button
 */
void MainDialog::OnExitClicked() {

	this->DestroyWindow();
}

/**
 * Name:		OnInputBrowseClicked
 * Purpose:		Handles clicks on input file browse button
 */
void MainDialog::OnInputBrowseClicked() {

	// Open "open file" dialog
	CFileDialog fileDialog(TRUE,
						   NULL,
						   NULL,
						   OFN_PATHMUSTEXIST | OFN_HIDEREADONLY,
						   "Binary files (*.bin)|*.bin|All files (*.*)|*.*||",
						   this);

	// Show dialog and get selection
	if (fileDialog.DoModal() == IDOK) {

		// Add filename in input combobox
		this->AddInputFile(fileDialog.GetFileName(), fileDialog.GetPathName());

		// Add filename in output combobox
		if (this->IsOutputFileSelected() == 0) {
		
			CString name(fileDialog.GetFileName());
			CString path(fileDialog.GetPathName());
			UINT count = fileDialog.GetFileName().GetLength() - fileDialog.GetFileTitle().GetLength();
			
			name.Delete(name.GetLength() - count, count);
			name += ".frw";

			path.Delete(path.GetLength() - count, count);
			path += ".frw";
			
			this->AddOutputFile(name, path);
		}

		this->EnablePackageButton();
	}
}

/**
 * Name:		OnOutputBrowseClicked
 * Purpose:		Handles clicks on output file browse button
 */
void MainDialog::OnOutputBrowseClicked() {

	// Display "Save as ..." dialog
	CFileDialog fileDialog(FALSE,
						   ".frw",
						   this->GetOutputFileName(),
						   OFN_HIDEREADONLY | OFN_OVERWRITEPROMPT,
						   "Firmware files (*.frw)|*.frw|All files (*.*)|*.*||",
						   this);

	// Add filename in output box
	if (fileDialog.DoModal() == IDOK) {

		this->AddOutputFile(fileDialog.GetFileName(), fileDialog.GetPathName());
		this->EnablePackageButton();
	}
}

/**
 * Name:		Encrypt
 * Purpose:		Asynchronous file encryption
 * Input(s):
 *	- Pointer to a MainDialog object
 */
UINT MainDialog::Encrypt(LPVOID pParam) {

	// Local variables
	CString cipherName;
	int cipherID;
	unsigned char * key = NULL;
	unsigned char * IV = NULL;
	unsigned char * block = NULL;
	unsigned char * tag = NULL;
	unsigned int bytes;
	unsigned int rate;
	unsigned int progress;
	MainDialog * pMainDialog = (MainDialog *) pParam;
	symmetric_ECB sECB;
	symmetric_CTR sCTR;
	symmetric_CBC sCBC;
	eax_state sEAX;

	try {

		// Get cipher
		if (pMainDialog->IsAESSelected()) {

			cipherName = "rijndael";
		}
		else if (pMainDialog->Is3DESSelected()) {

			cipherName = "3des";
		}

		// Setup progress window
		pMainDialog->m_pProgressDialog->SetTitle("Packaging ...");
		pMainDialog->m_pProgressDialog->SetRange(0, 20);
		pMainDialog->m_pProgressDialog->SetStep(1);

		// Open files
		CFile inputFile(pMainDialog->GetInputFilePath(),
						CFile::modeRead | CFile::shareExclusive);
		CFile outputFile(pMainDialog->GetOutputFilePath(),
						 CFile::modeWrite | CFile::modeCreate);

		// Calculate progress refresh rate
		rate = inputFile.GetLength() / 20;

		// Convert key
		key = new unsigned char[pMainDialog->m_keyLength];
		pMainDialog->ASCII2Hex((LPCTSTR) pMainDialog->GetKey(),
							   key, pMainDialog->m_keyLength);

		// Convert IV
		if (pMainDialog->IsCTRSelected() || pMainDialog->IsCBCSelected() || pMainDialog->IsEAXSelected()) {

			// Load IV
			IV = new unsigned char[pMainDialog->m_blockLength];
			pMainDialog->ASCII2Hex((LPCTSTR) pMainDialog->GetIV(),
								   IV, pMainDialog->m_blockLength);
		}
		
		// Find cipher
		if ((cipherID = find_cipher(cipherName)) < 0) {

			AfxMessageBox("Cannot find cipher");
			throw new CException();
		}

		// Initialize mode
		// ECB
		if (pMainDialog->IsECBSelected()) {

			if (ecb_start(cipherID, key, pMainDialog->m_keyLength, 0, &sECB) != CRYPT_OK) {

				AfxMessageBox("Cannot start cipher");
				throw new CException();
			}
		}
		// CBC
		else if (pMainDialog->IsCBCSelected()) {

			// Start CBC encryption
			if (cbc_start(cipherID, IV, key, pMainDialog->m_keyLength,
						  0, &sCBC) != CRYPT_OK) {

				AfxMessageBox("Cannot start cipher");
				throw new CException();
			}
		}
		// CTR
		else if (pMainDialog->IsCTRSelected()) {

			// Start CTR encryption
			if (ctr_start(cipherID, IV, key, pMainDialog->m_keyLength,
						  0, CTR_COUNTER_BIG_ENDIAN, &sCTR) != CRYPT_OK) {

				AfxMessageBox("Cannot start cipher");
				throw new CException();
			}
		}
		// EAX
		else if (pMainDialog->IsEAXSelected()) {

			// Initializes EAX encryption
			if (eax_init(&sEAX, cipherID, key, pMainDialog->m_keyLength,
				IV, pMainDialog->m_blockLength, NULL, 0) != CRYPT_OK) {

				AfxMessageBox("Cannot start EAX mode");
				throw new CException();
			}

			// Reserve file space for tag
			outputFile.Write(IV, pMainDialog->m_blockLength);
		}

		// Encrypt file
		progress = 0;
		block = new unsigned char[pMainDialog->m_blockLength];

		while ((bytes = inputFile.Read(block, pMainDialog->m_blockLength)) > 0) {

			// Stop encryption is requested
			if (pMainDialog->m_closeThread) {

				throw 0;
			}

			if (pMainDialog->IsECBSelected() || pMainDialog->IsCBCSelected()) {

				// Pad block if necessary
				while (bytes < (unsigned int) pMainDialog->m_blockLength) {

					block[bytes] = 0xFF;
					bytes++;
				}
			}

			// Encrypt block
			// ECB
			if (pMainDialog->IsECBSelected()) {
				if (ecb_encrypt(block, block, bytes, &sECB) != CRYPT_OK) {

					throw new CException();
				}
			}
			// CBC
			if (pMainDialog->IsCBCSelected()) {
				if (cbc_encrypt(block, block, bytes, &sCBC) != CRYPT_OK) {

					throw new CException();
				}
			}
			// CTR
			else if (pMainDialog->IsCTRSelected()) {
				if (ctr_encrypt(block, block, bytes, &sCTR) != CRYPT_OK) {

					throw new CException();
				}
			}
			// EAX
			else if (pMainDialog->IsEAXSelected()) {
				if (eax_encrypt(&sEAX, block, block, bytes) != CRYPT_OK) {

					throw new CException();
				}
			}

			// Write output
			outputFile.Write(block, bytes);

			// Refresh progress window
			progress += bytes;

			while (progress >= rate) {

				pMainDialog->m_pProgressDialog->StepIt();
				progress -= rate;
			}
		}

		pMainDialog->m_pProgressDialog->Finish();

		// Mode cleanup
		if (pMainDialog->IsECBSelected()) {

			ecb_done(&sECB);
		}
		if (pMainDialog->IsCBCSelected()) {

			cbc_done(&sCBC);
		}
		else if (pMainDialog->IsCTRSelected()) {

			ctr_done(&sCTR);
		}
		else if (pMainDialog->IsEAXSelected()) {

			unsigned long tagLength = pMainDialog->m_blockLength;
			tag = new unsigned char[pMainDialog->m_blockLength];
			eax_done(&sEAX, tag, &tagLength);

			// Prepend tag to output file
			outputFile.SeekToBegin();
			outputFile.Write(tag, tagLength);
		}
	}
	catch (int) {

		// Delete output file
		CFile::Remove(pMainDialog->GetOutputFileName());
	}
	catch (CException * e) {

		e->ReportError();
		e->Delete();
		pMainDialog->m_pProgressDialog->EndDialog(0);
	}

	// Cleanup
	if (key != NULL) {
		
		delete key;
	}
	if (block != NULL) {

		delete block;
	}
	if (IV != NULL) {

		delete IV;
	}
	if (tag != NULL) {

		delete tag;
	}

	return 0;
}

/**
 * Name:		OnPackageClicked
 * Purpose:		Handles clicks on the package button
 */
void MainDialog::OnPackageClicked() {

	// Local variables
	CString outputFilename;
	CFile file;
	BOOL encrypt = TRUE;

	// Check if output file exists
	if (file.Open(this->GetOutputFilePath(), CFile::modeRead) != 0) {

		file.Close();

		if (this->MessageBox("Filename " + this->GetOutputFilePath() + " already exists.\nDo you want to replace it?",
							 "File already exists",
							 MB_YESNO | MB_ICONEXCLAMATION) == IDNO) {

			encrypt = FALSE;
		}
	}
	
	if (encrypt) {

		// Display progress
		this->m_pProgressDialog = new ProgressDialog(this);

		this->m_closeThread = FALSE;
		this->m_thread = AfxBeginThread(MainDialog::Encrypt, (LPVOID) this);
		
		if (this->m_pProgressDialog->DoModal() == IDCANCEL) {

			// Stop thread
			this->m_closeThread = TRUE;
			WaitForSingleObject(this->m_thread->m_hThread, INFINITE);
		}

		delete this->m_pProgressDialog;
	}
}

/**
 * Name:		OnKeyUpdated
 * Purpose:		Handles updates of the key editbox
 */
void MainDialog::OnKeyUpdated() {

	// Local variables
	CString key;
	BOOL modified;
	DWORD sel;

	// Key must be in hexadecimal
	key = this->GetKey();
	sel = this->GetKeySel();

	modified = FALSE;
	for (int i=0; i < key.GetLength(); i++) {

		if (!((key[i] >= '0' && key[i] <= '9') || (key[i] >= 'A' && key[i] <= 'F'))) {

			key.Delete(i);
			modified = TRUE;
			i--;
		}
	}

	if (modified) {

		this->SetKey(key);
		this->SetKeySel(sel);
		MessageBeep(MB_ICONEXCLAMATION);
	}

	this->EnablePackageButton();
}

/**
 * Name:		OnIVUpdated
 * Purpose:		Handles updates of the IV editbox
 */
void MainDialog::OnIVUpdated() {

	// Local variables
	CString IV;
	BOOL modified;
	DWORD sel;

	// IV must be in hexadecimal
	IV = this->GetIV();
	sel = this->GetIVSel();

	modified = FALSE;
	for (int i=0; i < IV.GetLength(); i++) {

		if (!((IV[i] >= '0' && IV[i] <= '9') || (IV[i] >= 'A' && IV[i] <= 'F'))) {

			IV.Delete(i);
			modified = TRUE;
			i--;
		}
	}

	if (modified) {

		this->SetIV(IV);
		this->SetIVSel(sel);
		MessageBeep(MB_ICONEXCLAMATION);
	}

	this->EnablePackageButton();
}

/**
 * Name:		OnAESClicked
 * Purpose:		Handles clicks on AES radio button
 */
void MainDialog::OnAESClicked() {

	this->m_keyLength = 16;
	this->m_blockLength = 16;

	this->EnablePackageButton();
}

/**
 * Name:		On3DESClicked
 * Purpose:		Handles clicks on 3DES radio button
 */
void MainDialog::On3DESClicked() {

	this->m_keyLength = 24;
	this->m_blockLength = 8;

	this->EnablePackageButton();
}

/**
 * Name:		OnECBClicked
 * Purpose:		Handles clicks on the ECB radio button
 */
void MainDialog::OnECBClicked() {

	this->DisableIV();
	this->EnablePackageButton();
}

/**
 * Name:		OnCBCClicked
 * Purpose:		Handles clicks on the CBC radio button
 */
void MainDialog::OnCBCClicked() {

	this->EnableIV();
	this->EnablePackageButton();
}