No OneTemporary
Actions

Size

67 KB

Subscribers

None

View Options

	diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp
	index a0c5cca..22b4ddc 100644
	--- a/Scylla/MainGui.cpp
	+++ b/Scylla/MainGui.cpp
	@@ -1,1255 +1,1238 @@
	#include "MainGui.h"

	#include <atldlgs.h> // WTL common dialogs

	#include "definitions.h"
	#include "PluginLoader.h"
	#include "ConfigurationHolder.h"
	#include "PeDump.h"
	#include "PeRebuild.h"
	#include "DllInjectionPlugin.h"
	#include "DisassemblerGui.h"
	#include "PickApiGui.h"
	#include "NativeWinApi.h"
	#include "ImportRebuild.h"
	#include "SystemInformation.h"
	#include "AboutGui.h"
	#include "OptionsGui.h"

	extern CAppModule _Module; // o_O

	const WCHAR MainGui::filterExe[] = L"Executable (.exe)\0.exe\0All files\0.\0";
	const WCHAR MainGui::filterDll[] = L"Dynamic Link Library (.dll)\0.dll\0All files\0.\0";
	const WCHAR MainGui::filterExeDll[] = L"Executable (.exe)\0.exe\0Dynamic Link Library (.dll)\0.dll\0All files\0.\0";
	const WCHAR MainGui::filterTxt[] = L"Text file (.txt)\0.txt\0All files\0.\0";

	MainGui::MainGui() : selectedProcess(0), importsHandling(TreeImports), TreeImportsSubclass(this, IDC_TREE_IMPORTS)
	{
	Logger::getDebugLogFilePath();
	ConfigurationHolder::loadConfiguration();
	PluginLoader::findAllPlugins();
	NativeWinApi::initialize();
	SystemInformation::getSystemInformation();

	if(ConfigurationHolder::getConfigObject(DEBUG_PRIVILEGE)->isTrue())
	{
	processLister.setDebugPrivileges();
	}

	processAccessHelp.getProcessModules(GetCurrentProcessId(), processAccessHelp.ownModuleList);

	hIcon.LoadIcon(IDI_ICON_SCYLLA);
	hMenuImports.LoadMenu(IDR_MENU_IMPORTS);
	hMenuLog.LoadMenu(IDR_MENU_LOG);
	accelerators.LoadAccelerators(IDR_ACCELERATOR_MAIN);

	hIconCheck.LoadIcon(IDI_ICON_CHECK, 16, 16);
	hIconWarning.LoadIcon(IDI_ICON_WARNING, 16, 16);
	hIconError.LoadIcon(IDI_ICON_ERROR, 16, 16);

	if(hMenuImports)
	{
	appendPluginListToMenu(hMenuImports.GetSubMenu(0));
	}
	}

	BOOL MainGui::PreTranslateMessage(MSG* pMsg)
	{
	if(accelerators.TranslateAccelerator(m_hWnd, pMsg))
	{
	return TRUE;
	}
	else if(IsDialogMessage(pMsg))
	{
	return TRUE;
	}

	return FALSE;
	}

	BOOL MainGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam)
	{
	if (SystemInformation::currenOS == UNKNOWN_OS)
	{
	if(IDCANCEL == MessageBox(L"Operating System is not supported\r\nContinue anyway?", L"Scylla", MB_ICONWARNING \| MB_OKCANCEL))
	{
	SendMessage(WM_CLOSE);
	return FALSE;
	}
	}

	CMessageLoop* pLoop = _Module.GetMessageLoop();
	pLoop->AddMessageFilter(this);

	setupStatusBar();
	fillStatusBar();

	DoDataExchange(); // attach controls
	DlgResize_Init(true, true);

	EditOEPAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);
	EditIATAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);
	EditIATSize.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);

	appendPluginListToMenu(CMenuHandle(GetMenu()).GetSubMenu(MenuImportsOffsetTrace));

	enableDialogControls(FALSE);

	setIconAndDialogCaption();

	return TRUE;
	}

	void MainGui::OnDestroy()
	{
	PostQuitMessage(0);
	}

	void MainGui::OnSize(UINT nType, CSize size)
	{
	StatusBar.SendMessage(WM_SIZE);
	SetMsgHandled(FALSE);
	}

	void MainGui::OnLButtonDown(UINT nFlags, CPoint point)
	{
	SetMsgHandled(FALSE);
	}

	void MainGui::OnContextMenu(CWindow wnd, CPoint point)
	{
	// point = -1, -1 for keyboard invoked shortcut!
	switch(wnd.GetDlgCtrlID())
	{
	case IDC_TREE_IMPORTS:
	DisplayContextMenuImports(wnd, point);
	return;
	case IDC_LIST_LOG:
	DisplayContextMenuLog(wnd, point);
	return;
	//default: // wnd == m_hWnd?
	// DisplayContextMenu(wnd, point);
	// return;
	}

	SetMsgHandled(FALSE);
	}

	void MainGui::OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// Make sure it's a menu
	if(uNotifyCode == 0 && !wndCtl.IsWindow())
	{
	if ((nID >= PLUGIN_MENU_BASE_ID) && (nID <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID)))
	{
	pluginActionHandler(nID);
	return;
	}
	}
	SetMsgHandled(FALSE);
	}

	LRESULT MainGui::OnTreeImportsClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsDoubleClick(const NMHDR* pnmh)
	{
	if(TreeImports.GetCount() < 1)
	return 0;

	// Get item under cursor
	CPoint client = GetMessagePos();
	TreeImports.ScreenToClient(&client);
	UINT flags;
	CTreeItem over = TreeImports.HitTest(client, &flags);
	CTreeItem parent;
	if(over)
	{
	if(!(flags & TVHT_ONITEM))
	{
	over = NULL;
	}
	else
	{
	parent = over.GetParent();
	}
	}

	if(!over.IsNull() && !parent.IsNull())
	{
	pickApiActionHandler(over);
	}

	return 0;
	}

	LRESULT MainGui::OnTreeImportsRightClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsRightDoubleClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh)
	{
	const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh;
	switch(tkd->wVKey)
	{
	case VK_RETURN:
	{
	CTreeItem selected = TreeImports.GetFirstSelectedItem();
	if(!selected.IsNull() && !selected.GetParent().IsNull())
	{
	pickApiActionHandler(selected);
	}
	}
	return 1;
	case VK_DELETE:
	- {
	- CTreeItem selected = TreeImports.GetFirstSelectedItem();
	- while(!selected.IsNull())
	- {
	- if(selected.GetParent().IsNull())
	- {
	- importsHandling.deleteTreeNode(selected);
	- }
	- else
	- {
	- importsHandling.cutThunk(selected);
	- }
	- selected = TreeImports.GetNextSelectedItem(selected);
	- }
	- fillStatusBar();
	- }
	+ deleteSelectedImportsActionHandler();
	return 1;
	}

	SetMsgHandled(FALSE);
	return 0;
	}

	+void MainGui::deleteSelectedImportsActionHandler()
	+{
	+ CTreeItem selected = TreeImports.GetFirstSelectedItem();
	+ while(!selected.IsNull())
	+ {
	+ if(importsHandling.isModule(selected))
	+ {
	+ importsHandling.cutModule(selected);
	+ }
	+ else
	+ {
	+ importsHandling.cutImport(selected);
	+ }
	+ selected = TreeImports.GetNextSelectedItem(selected);
	+ }
	+ fillStatusBar();
	+}
	+
	+void MainGui::invalidateSelectedImportsActionHandler()
	+{
	+ CTreeItem selected = TreeImports.GetFirstSelectedItem();
	+ while(!selected.IsNull())
	+ {
	+ if(importsHandling.isImport(selected))
	+ {
	+ importsHandling.invalidateImport(selected);
	+ }
	+ selected = TreeImports.GetNextSelectedItem(selected);
	+ }
	+ fillStatusBar();
	+}
	+
	UINT MainGui::OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg)
	{
	if(lpMsg)
	{
	switch(lpMsg->wParam)
	{
	case VK_RETURN:
	return DLGC_WANTMESSAGE;
	}
	}

	SetMsgHandled(FALSE);
	return 0;
	}

	void MainGui::OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
	switch(nChar)
	{
	case VK_RETURN:
	break;
	default:
	SetMsgHandled(FALSE);
	break;
	}
	}

	void MainGui::OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	fillProcessListComboBox(ComboProcessList);
	}

	void MainGui::OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	processSelectedActionHandler(ComboProcessList.GetCurSel());
	}

	void MainGui::OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	pickDllActionHandler();
	}

	void MainGui::OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	optionsActionHandler();
	}

	void MainGui::OnDump(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dumpActionHandler();
	}

	void MainGui::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dumpFixActionHandler();
	}

	void MainGui::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	peRebuildActionHandler();
	}

	void MainGui::OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dllInjectActionHandler();
	}

	void MainGui::OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	iatAutosearchActionHandler();
	}

	void MainGui::OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	getImportsActionHandler();
	}

	void MainGui::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showInvalidImportsActionHandler();
	}

	void MainGui::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showSuspectImportsActionHandler();
	}

	void MainGui::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	clearImportsActionHandler();
	}

	void MainGui::OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	- // TODO
	+ invalidateSelectedImportsActionHandler();
	}

	void MainGui::OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	- // TODO
	+ deleteSelectedImportsActionHandler();
	}

	void MainGui::OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnAutotrace(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnExit(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	DestroyWindow();
	}

	void MainGui::OnAbout(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showAboutDialog();
	}

	void MainGui::setupStatusBar()
	{
	- StatusBar.Create(m_hWnd, NULL, L"", WS_CHILD \| WS_VISIBLE \| WS_CLIPCHILDREN \| WS_CLIPSIBLINGS \| SBARS_TOOLTIPS, NULL, ATL_IDW_STATUS_BAR);
	+ StatusBar.Create(m_hWnd, NULL, L"", WS_CHILD \| WS_VISIBLE \| WS_CLIPCHILDREN \| WS_CLIPSIBLINGS \| SBARS_TOOLTIPS, NULL, IDC_STATUS_BAR);

	CRect rcMain, rcStatus;
	GetClientRect(&rcMain);
	StatusBar.GetWindowRect(&rcStatus);

	const int PARTS = 4;
	int widths[PARTS];

	widths[PART_COUNT] = rcMain.Width() / 5;
	widths[PART_INVALID] = widths[PART_COUNT] + rcMain.Width() / 5;
	widths[PART_IMAGEBASE] = widths[PART_INVALID] + rcMain.Width() / 3;
	widths[PART_MODULE] = -1;

	StatusBar.SetParts(PARTS, widths);

	ResizeClient(rcMain.Width(), rcMain.Height() + rcStatus.Height(), FALSE);
	}

	void MainGui::fillStatusBar()
	{
	// Rewrite ImportsHandling so we get these easily
	- unsigned int totalImports = 0;
	- unsigned int invalidImports = 0;
	+ unsigned int totalImports = importsHandling.thunkCount();
	+ unsigned int invalidImports = importsHandling.invalidThunkCount();

	// \t = center, \t\t = right-align
	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tImports: %u"), totalImports);
	StatusBar.SetText(PART_COUNT, stringBuffer);

	if(invalidImports > 0)
	{
	StatusBar.SetIcon(PART_INVALID, hIconError);
	}
	else
	{
	StatusBar.SetIcon(PART_INVALID, hIconCheck);
	}

	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tInvalid: %u"), invalidImports);
	StatusBar.SetText(PART_INVALID, stringBuffer);

	if(selectedProcess)
	{
	DWORD_PTR imageBase = 0;
	const WCHAR * fileName = 0;

	if(processAccessHelp.selectedModule)
	{
	imageBase = processAccessHelp.selectedModule->modBaseAddr;
	fileName = processAccessHelp.selectedModule->getFilename();
	}
	else
	{
	imageBase = selectedProcess->imageBase;
	fileName = selectedProcess->filename;
	}

	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tImagebase: ")TEXT(PRINTF_DWORD_PTR_FULL), imageBase);
	StatusBar.SetText(PART_IMAGEBASE, stringBuffer);
	StatusBar.SetText(PART_MODULE, fileName);
	StatusBar.SetTipText(PART_MODULE, fileName);
	}
	else
	{
	StatusBar.SetText(PART_IMAGEBASE, L"");
	StatusBar.SetText(PART_MODULE, L"");
	}
	}

	bool MainGui::showFileDialog(WCHAR * selectedFile, bool save, const WCHAR * defFileName, const WCHAR * filter, const WCHAR * defExtension, const WCHAR * directory)
	{
	OPENFILENAME ofn = {0};

	// WTL doesn't support new explorer styles on Vista and up
	// This is because it uses a custom hook, we could remove it or derive
	// from CFileDialog but this solution is easier and allows more control anyway (e.g. initial dir)

	if(defFileName)
	{
	wcscpy_s(selectedFile, MAX_PATH, defFileName);
	}
	else
	{
	selectedFile[0] = _T('\0');
	}

	ofn.lStructSize = sizeof(ofn);
	ofn.hwndOwner = m_hWnd;
	ofn.lpstrFilter = filter;
	ofn.lpstrDefExt = defExtension; // only first 3 chars are used, no dots!
	ofn.lpstrFile = selectedFile;
	ofn.lpstrInitialDir = directory;
	ofn.nMaxFile = MAX_PATH;
	ofn.Flags = OFN_PATHMUSTEXIST \| OFN_HIDEREADONLY;

	/*
	*OFN_EXPLORER is automatically used, it only has to be specified
	*if using a custom hook
	*OFN_LONGNAMES is automatically used by explorer-style dialogs
	*/

	if(save)
	ofn.Flags \|= OFN_OVERWRITEPROMPT;
	else
	ofn.Flags \|= OFN_FILEMUSTEXIST;

	if(save)
	return 0 != GetSaveFileName(&ofn);
	else
	return 0 != GetOpenFileName(&ofn);
	}

	void MainGui::setIconAndDialogCaption()
	{
	SetIcon(hIcon, TRUE);
	SetIcon(hIcon, FALSE);

	SetWindowText(TEXT(APPNAME)TEXT(" ")TEXT(ARCHITECTURE)TEXT(" ")TEXT(APPVERSION));
	}

	void MainGui::pickDllActionHandler()
	{
	if(!selectedProcess)
	return;

	PickDllGui dlgPickDll(processAccessHelp.moduleList);
	if(dlgPickDll.DoModal())
	{
	//get selected module
	processAccessHelp.selectedModule = dlgPickDll.getSelectedModule();
	Logger::printfDialog(TEXT("->>> Module %s selected."), processAccessHelp.selectedModule->getFilename());
	Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),processAccessHelp.selectedModule->modBaseAddr,processAccessHelp.selectedModule->modBaseSize);
	}
	else
	{
	processAccessHelp.selectedModule = 0;
	}

	fillStatusBar();
	}

	void MainGui::pickApiActionHandler(CTreeItem item)
	{
	- CTreeItem parent = item.GetParent();
	- if(parent.IsNull())
	+ if(!importsHandling.isImport(item))
	return;

	// TODO: new node when user picked an API from another DLL?

	PickApiGui dlgPickApi(processAccessHelp.moduleList);
	if(dlgPickApi.DoModal())
	{
	const ApiInfo* api = dlgPickApi.getSelectedApi();
	if(api && api->module)
	{
	- std::map<DWORD_PTR, ImportModuleThunk>::iterator iterator1;
	- std::map<DWORD_PTR, ImportThunk>::iterator iterator2;
	-
	- iterator1 = importsHandling.moduleList.begin();
	- while(iterator1 != importsHandling.moduleList.end())
	- {
	- if(iterator1->second.hTreeItem == parent)
	- {
	- iterator2 = iterator1->second.thunkList.begin();
	- while(iterator2 != iterator1->second.thunkList.end())
	- {
	- if(iterator2->second.hTreeItem == item)
	- {
	- ImportThunk &imp = iterator2->second;
	- wcscpy_s(imp.moduleName, MAX_PATH, api->module->getFilename());
	- strcpy_s(imp.name, MAX_PATH, api->name);
	- imp.ordinal = api->ordinal;
	- //imp.apiAddressVA = api->va; //??
	- imp.hint = (WORD)api->hint;
	- imp.valid = true;
	- imp.suspect = api->isForwarded;
	-
	- importsHandling.updateImportInTreeView(&imp, item);
	- break;
	- }
	- iterator2++;
	- }
	- break;
	- }
	- iterator1++;
	- }
	+ importsHandling.setImport(item, api->module->getFilename(), api->name, api->ordinal, api->hint, true, api->isForwarded);
	}
	}

	fillStatusBar();
	}

	void MainGui::startDisassemblerGui(CTreeItem selectedTreeNode)
	{
	if(!selectedProcess)
	return;

	DWORD_PTR address = importsHandling.getApiAddressByNode(selectedTreeNode);
	if (address)
	{
	BYTE test;
	if(!ProcessAccessHelp::readMemoryFromProcess(address, sizeof(test), &test))
	{
	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("Can't read memory at ")TEXT(PRINTF_DWORD_PTR_FULL),address);
	MessageBox(stringBuffer, L"Failure", MB_ICONERROR);
	}
	else
	{
	DisassemblerGui dlgDisassembler(address);
	dlgDisassembler.DoModal();
	}
	}
	}

	void MainGui::processSelectedActionHandler(int index)
	{
	std::vector<Process>& processList = processLister.getProcessList();
	Process &process = processList.at(index);
	selectedProcess = 0;

	clearImportsActionHandler();

	Logger::printfDialog(TEXT("Analyzing %s"),process.fullPath);

	if (processAccessHelp.hProcess != 0)
	{
	processAccessHelp.closeProcessHandle();
	apiReader.clearAll();
	}

	if (!processAccessHelp.openProcessHandle(process.PID))
	{
	enableDialogControls(FALSE);
	Logger::printfDialog(TEXT("Error: Cannot open process handle."));
	fillStatusBar();
	return;
	}

	processAccessHelp.getProcessModules(process.PID, processAccessHelp.moduleList);

	apiReader.readApisFromModuleList();

	Logger::printfDialog(TEXT("Loading modules done."));

	//TODO improve
	processAccessHelp.selectedModule = 0;
	processAccessHelp.targetSizeOfImage = process.imageSize;
	processAccessHelp.targetImageBase = process.imageBase;

	ProcessAccessHelp::getSizeOfImageCurrentProcess();

	process.imageSize = (DWORD)processAccessHelp.targetSizeOfImage;


	Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),process.imageBase, process.imageSize);

	process.entryPoint = ProcessAccessHelp::getEntryPointFromFile(process.fullPath);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),process.entryPoint + process.imageBase);

	EditOEPAddress.SetWindowText(stringBuffer);

	selectedProcess = &process;
	enableDialogControls(TRUE);

	fillStatusBar();
	}

	void MainGui::fillProcessListComboBox(CComboBox& hCombo)
	{
	hCombo.ResetContent();

	std::vector<Process>& processList = processLister.getProcessListSnapshot();

	for (size_t i = 0; i < processList.size(); i++)
	{
	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("0x%04X - %s - %s"),processList[i].PID,processList[i].filename,processList[i].fullPath);
	hCombo.AddString(stringBuffer);
	}
	}

	void MainGui::addTextToOutputLog(const WCHAR * text)
	{
	if (m_hWnd)
	{
	ListLog.SetCurSel(ListLog.AddString(text));
	}
	}

	void MainGui::clearOutputLog()
	{
	if (m_hWnd)
	{
	ListLog.ResetContent();
	}
	}

	bool MainGui::saveLogToFile(const WCHAR * file)
	{
	const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian
	const WCHAR newLine[] = L"\r\n";
	bool success = true;

	HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
	if(hFile != INVALID_HANDLE_VALUE)
	{
	ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM);

	WCHAR * buffer = 0;
	int bufsize = 0;
	for(int i = 0; i < ListLog.GetCount(); i++)
	{
	int size = ListLog.GetTextLen(i);
	size += _countof(newLine)-1;
	if(size+1 > bufsize)
	{
	bufsize = size+1;
	delete[] buffer;
	try
	{
	buffer = new WCHAR[bufsize];
	}
	catch(std::bad_alloc&)
	{
	buffer = 0;
	success = false;
	break;
	}
	}

	ListLog.GetText(i, buffer);
	wcscat_s(buffer, bufsize, newLine);

	ProcessAccessHelp::writeMemoryToFileEnd(hFile, size * sizeof(WCHAR), buffer);
	}
	delete[] buffer;
	CloseHandle(hFile);
	}
	return success;
	}

	void MainGui::showInvalidImportsActionHandler()
	{
	- importsHandling.showImports(true, false);
	+ importsHandling.selectImports(true, false);
	}

	void MainGui::showSuspectImportsActionHandler()
	{
	- importsHandling.showImports(false, true);
	+ importsHandling.selectImports(false, true);
	}

	void MainGui::iatAutosearchActionHandler()
	{
	DWORD_PTR searchAddress = 0;
	DWORD_PTR addressIAT = 0;
	DWORD sizeIAT = 0;
	IATSearch iatSearch;

	if(!selectedProcess)
	return;

	if(EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 1)
	{
	searchAddress = stringToDwordPtr(stringBuffer);
	if (searchAddress)
	{
	if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT))
	{
	Logger::printfDialog(TEXT("IAT found at VA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" RVA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d)"),addressIAT, addressIAT - processAccessHelp.targetImageBase,sizeIAT,sizeIAT);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),addressIAT);
	EditIATAddress.SetWindowText(stringBuffer);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("%08X"),sizeIAT);
	EditIATSize.SetWindowText(stringBuffer);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("IAT found! Start Address ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d) "),addressIAT,sizeIAT,sizeIAT);
	MessageBox(stringBuffer, L"IAT found", MB_ICONINFORMATION);
	}
	else
	{
	Logger::printfDialog(TEXT("IAT not found at OEP ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("!"),searchAddress);
	}
	}
	}
	}

	void MainGui::getImportsActionHandler()
	{
	DWORD_PTR addressIAT = 0;
	DWORD sizeIAT = 0;

	if(!selectedProcess)
	return;

	if (EditIATAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	addressIAT = stringToDwordPtr(stringBuffer);
	}

	if (EditIATSize.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	sizeIAT = wcstoul(stringBuffer, NULL, 16);
	}

	if (addressIAT && sizeIAT)
	{
	apiReader.readAndParseIAT(addressIAT, sizeIAT,importsHandling.moduleList);
	importsHandling.displayAllImports();
	}

	fillStatusBar();
	}

	DWORD_PTR MainGui::stringToDwordPtr(const WCHAR * hexString)
	{
	DWORD_PTR address = 0;

	#ifdef _WIN64
	address = _wcstoui64(hexString, NULL, 16);
	#else
	address = wcstoul(hexString, NULL, 16);
	#endif

	if (address == 0)
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(L"stringToDwordPtr :: address == 0, %s",hexString);
	#endif
	return 0;
	}
	else
	{
	return address;
	}
	}

	void MainGui::SetupImportsMenuItems(bool isItem, bool isThunk)
	{
	// assert(!(!isItem && isThunk));

	CMenuHandle hSub = hMenuImports.GetSubMenu(0);

	UINT itemOnly = isItem ? MF_ENABLED : MF_GRAYED;
	UINT thunkOnly = isThunk ? MF_ENABLED : MF_GRAYED;

	hSub.EnableMenuItem(ID__INVALIDATEFUNCTION, thunkOnly);
	hSub.EnableMenuItem(ID__DISASSEMBLE, thunkOnly);
	hSub.EnableMenuItem(ID__CUTTHUNK, thunkOnly);

	hSub.EnableMenuItem(ID__DELETETREENODE, itemOnly);
	}

	void MainGui::DisplayContextMenuImports(CWindow hwnd, CPoint pt)
	{
	if(TreeImports.GetCount() < 1)
	return;

	CTreeItem over, parent;

	if(pt.x == -1 && pt.y == -1) // invoked by keyboard
	{
	CRect pos;
	over = TreeImports.GetFirstSelectedItem();
	if(over)
	{
	over.EnsureVisible();
	over.GetRect(&pos, TRUE);
	TreeImports.ClientToScreen(&pos);
	}
	else
	{
	TreeImports.GetWindowRect(&pos);
	}
	pt = pos.TopLeft();
	}
	else
	{
	// Get item under cursor
	CPoint client = pt;
	TreeImports.ScreenToClient(&client);
	UINT flags;
	over = TreeImports.HitTest(client, &flags);
	if(over && !(flags & TVHT_ONITEM))
	{
	over = NULL;
	}
	}

	- if(over)
	- {
	- parent = over.GetParent();
	- }
	-
	if (hMenuImports)
	{
	// Prepare hmenuImports
	- SetupImportsMenuItems(!over.IsNull(), !parent.IsNull());
	+ SetupImportsMenuItems(!over.IsNull(), importsHandling.isImport(over));

	CMenuHandle hSub = hMenuImports.GetSubMenu(0);

	BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN \| TPM_RIGHTBUTTON \| TPM_RETURNCMD, pt.x, pt.y, hwnd);
	if (menuItem)
	{
	if ((menuItem >= PLUGIN_MENU_BASE_ID) && (menuItem <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID)))
	{
	//wsprintf(stringBuffer, L"%d %s\n",menuItem,pluginList[menuItem - PLUGIN_MENU_BASE_ID].pluginName);
	//MessageBox(stringBuffer, L"plugin selection");

	pluginActionHandler(menuItem);
	return;
	}
	switch (menuItem)
	{
	case ID__INVALIDATEFUNCTION:
	- importsHandling.invalidateFunction(over);
	+ importsHandling.invalidateImport(over);
	break;
	case ID__DISASSEMBLE:
	startDisassemblerGui(over);
	break;
	case ID__EXPANDALLNODES:
	importsHandling.expandAllTreeNodes();
	break;
	case ID__COLLAPSEALLNODES:
	importsHandling.collapseAllTreeNodes();
	break;
	case ID__CUTTHUNK:
	- importsHandling.cutThunk(over);
	+ importsHandling.cutImport(over);
	break;
	case ID__DELETETREENODE:
	- importsHandling.deleteTreeNode(parent ? parent : over);
	+ importsHandling.cutModule(importsHandling.isImport(over) ? over.GetParent() : over);
	break;
	}
	}
	}

	fillStatusBar(); // ?
	}

	void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt)
	{
	if (hMenuLog)
	{
	if(pt.x == -1 && pt.y == -1) // invoked by keyboard
	{
	CRect pos;
	ListLog.GetWindowRect(&pos);
	pt = pos.TopLeft();
	}

	CMenuHandle hSub = hMenuLog.GetSubMenu(0);
	BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN \| TPM_RIGHTBUTTON \| TPM_RETURNCMD, pt.x, pt.y, hwnd);
	if (menuItem)
	{
	switch (menuItem)
	{
	case ID__SAVE:
	WCHAR selectedFilePath[MAX_PATH];
	if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt"))
	{
	saveLogToFile(selectedFilePath);
	}
	break;
	case ID__CLEAR:
	clearOutputLog();
	break;
	}
	}
	}
	}

	void MainGui::appendPluginListToMenu(CMenuHandle hMenu)
	{
	std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList();

	if (scyllaPluginList.size() > 0)
	{
	CMenuHandle newMenu;
	newMenu.CreatePopupMenu();

	for (size_t i = 0; i < scyllaPluginList.size(); i++)
	{
	newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName);
	}

	hMenu.AppendMenu(MF_MENUBARBREAK);
	hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins");
	}

	if (imprecPluginList.size() > 0)
	{
	CMenuHandle newMenu;
	newMenu.CreatePopupMenu();

	for (size_t i = 0; i < imprecPluginList.size(); i++)
	{
	newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName);
	}

	hMenu.AppendMenu(MF_MENUBARBREAK);
	hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins");
	}
	}

	void MainGui::dumpActionHandler()
	{
	if(!selectedProcess)
	return;

	WCHAR selectedFilePath[MAX_PATH];
	const WCHAR * fileFilter;
	const WCHAR * defExtension;
	PeDump peDump;

	if (processAccessHelp.selectedModule)
	{
	fileFilter = filterDll;
	defExtension = L"dll";
	}
	else
	{
	fileFilter = filterExe;
	defExtension = L"exe";
	}

	if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension))
	{
	if (processAccessHelp.selectedModule)
	{
	//dump DLL

	peDump.imageBase = processAccessHelp.selectedModule->modBaseAddr;
	peDump.sizeOfImage = processAccessHelp.selectedModule->modBaseSize;
	//get it from gui
	peDump.entryPoint = getOEPFromGui();
	wcscpy_s(peDump.fullpath, MAX_PATH, processAccessHelp.selectedModule->fullPath);
	}
	else
	{
	peDump.imageBase = ProcessAccessHelp::targetImageBase;
	peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage;
	//get it from gui
	peDump.entryPoint = getOEPFromGui();
	wcscpy_s(peDump.fullpath, MAX_PATH, selectedProcess->fullPath);
	}

	peDump.useHeaderFromDisk = ConfigurationHolder::getConfigObject(USE_PE_HEADER_FROM_DISK)->isTrue();
	if (peDump.dumpCompleteProcessToDisk(selectedFilePath))
	{
	Logger::printfDialog(TEXT("Dump success %s"),selectedFilePath);
	//MessageBox(L"Image dumped successfully.", L"Success");
	}
	else
	{
	Logger::printfDialog(TEXT("Error: Cannot dump image."));
	MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR);
	}
	}
	}

	DWORD_PTR MainGui::getOEPFromGui()
	{
	if (EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	return stringToDwordPtr(stringBuffer);
	}
	else
	{
	return 0;
	}
	}

	void MainGui::peRebuildActionHandler()
	{
	DWORD newSize = 0;
	WCHAR selectedFilePath[MAX_PATH];
	PeRebuild peRebuild;

	if(showFileDialog(selectedFilePath, false, NULL, filterExeDll))
	{
	if (ConfigurationHolder::getConfigObject(CREATE_BACKUP)->isTrue())
	{
	if (!ProcessAccessHelp::createBackupFile(selectedFilePath))
	{
	Logger::printfDialog(TEXT("Creating backup file failed %s"), selectedFilePath);
	}
	}

	LONGLONG fileSize = ProcessAccessHelp::getFileSize(selectedFilePath);
	LPVOID mapped = peRebuild.createFileMappingViewFull(selectedFilePath);

	newSize = peRebuild.realignPE(mapped, (DWORD)fileSize);
	peRebuild.closeAllMappingHandles();

	if (newSize < 10)
	{
	Logger::printfDialog(TEXT("Rebuild failed %s"), selectedFilePath);
	MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR);
	}
	else
	{
	peRebuild.truncateFile(selectedFilePath, newSize);

	Logger::printfDialog(TEXT("Rebuild success %s"), selectedFilePath);
	Logger::printfDialog(TEXT("-> Old file size 0x%08X new file size 0x%08X (%d %%)"), (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) );
	//MessageBox(L"Image rebuilded successfully.", L"Success", MB_ICONINFORMATION);
	}
	}
	}

	void MainGui::dumpFixActionHandler()
	{
	if(!selectedProcess)
	return;

	if (TreeImports.GetCount() < 2)
	{
	Logger::printfDialog(TEXT("Nothing to rebuild"));
	return;
	}

	WCHAR newFilePath[MAX_PATH];
	WCHAR selectedFilePath[MAX_PATH];
	const WCHAR * fileFilter;

	ImportRebuild importRebuild;

	if (processAccessHelp.selectedModule)
	{
	fileFilter = filterDll;
	}
	else
	{
	fileFilter = filterExe;
	}

	if (showFileDialog(selectedFilePath, false, NULL, fileFilter))
	{
	wcscpy_s(newFilePath,MAX_PATH,selectedFilePath);

	const WCHAR * extension = 0;

	WCHAR* dot = wcsrchr(newFilePath, L'.');
	if (dot)
	{
	*dot = L'\0';
	extension = selectedFilePath + (dot - newFilePath); //wcsrchr(selectedFilePath, L'.');
	}

	wcscat_s(newFilePath, MAX_PATH, L"_SCY");

	if(extension)
	{
	wcscat_s(newFilePath, MAX_PATH, extension);
	}

	if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList))
	{
	//MessageBox(L"Imports rebuilding successful", L"Success", MB_ICONINFORMATION);

	Logger::printfDialog(TEXT("Import Rebuild success %s"), newFilePath);
	}
	else
	{
	Logger::printfDialog(TEXT("Import Rebuild failed, target %s"), selectedFilePath);
	MessageBox(L"Imports rebuilding failed", L"Failure", MB_ICONERROR);
	}
	}
	}

	void MainGui::enableDialogControls(BOOL value)
	{
	- GetDlgItem(IDC_BTN_PICKDLL).EnableWindow(value);
	- GetDlgItem(IDC_BTN_DUMP).EnableWindow(value);
	- GetDlgItem(IDC_BTN_FIXDUMP).EnableWindow(value);
	- GetDlgItem(IDC_BTN_IATAUTOSEARCH).EnableWindow(value);
	- GetDlgItem(IDC_BTN_GETIMPORTS).EnableWindow(value);
	- GetDlgItem(IDC_BTN_SUSPECTIMPORTS).EnableWindow(value);
	- GetDlgItem(IDC_BTN_INVALIDIMPORTS).EnableWindow(value);
	- GetDlgItem(IDC_BTN_CLEARIMPORTS).EnableWindow(value);
	+ BOOL valButton = value ? TRUE : FALSE;
	+
	+ GetDlgItem(IDC_BTN_PICKDLL).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_DUMP).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_FIXDUMP).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_IATAUTOSEARCH).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_GETIMPORTS).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_SUSPECTIMPORTS).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_INVALIDIMPORTS).EnableWindow(valButton);
	+ GetDlgItem(IDC_BTN_CLEARIMPORTS).EnableWindow(valButton);

	CMenuHandle menu = GetMenu();

	- menu.EnableMenuItem(ID_FILE_DUMP, value ? MF_ENABLED : MF_GRAYED);
	- menu.EnableMenuItem(ID_FILE_FIXDUMP, value ? MF_ENABLED : MF_GRAYED);
	- menu.EnableMenuItem(ID_MISC_DLLINJECTION, value ? MF_ENABLED : MF_GRAYED);
	- menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetScylla, MF_BYPOSITION \| (value ? MF_ENABLED : MF_GRAYED));
	- menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetImpRec, MF_BYPOSITION \| (value ? MF_ENABLED : MF_GRAYED));
	+ UINT valMenu = value ? MF_ENABLED : MF_GRAYED;
	+
	+ menu.EnableMenuItem(ID_FILE_DUMP, valMenu);
	+ menu.EnableMenuItem(ID_FILE_FIXDUMP, valMenu);
	+ menu.EnableMenuItem(ID_IMPORTS_INVALIDATESELECTED, valMenu);
	+ menu.EnableMenuItem(ID_IMPORTS_CUTSELECTED, valMenu);
	+ menu.EnableMenuItem(ID_MISC_DLLINJECTION, valMenu);
	+ menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetScylla, MF_BYPOSITION \| valMenu);
	+ menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetImpRec, MF_BYPOSITION \| valMenu);

	//not yet implemented
	- GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE);
	GetDlgItem(IDC_BTN_SAVETREE).EnableWindow(FALSE);
	GetDlgItem(IDC_BTN_LOADTREE).EnableWindow(FALSE);
	+ GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE);

	- menu.EnableMenuItem(ID_IMPORTS_INVALIDATESELECTED, MF_GRAYED);
	- menu.EnableMenuItem(ID_IMPORTS_CUTSELECTED, MF_GRAYED);
	- menu.EnableMenuItem(ID_IMPORTS_SAVETREE, MF_GRAYED);
	menu.EnableMenuItem(ID_IMPORTS_SAVETREE, MF_GRAYED);
	menu.EnableMenuItem(ID_IMPORTS_LOADTREE, MF_GRAYED);
	menu.EnableMenuItem(ID_TRACE_AUTOTRACE, MF_GRAYED);
	}

	void MainGui::showAboutDialog()
	{
	AboutGui dlgAbout;
	dlgAbout.DoModal();
	}

	void MainGui::dllInjectActionHandler()
	{
	if(!selectedProcess)
	return;

	WCHAR selectedFilePath[MAX_PATH];
	HMODULE hMod = 0;
	DllInjection dllInjection;

	if (showFileDialog(selectedFilePath, false, NULL, filterDll))
	{
	hMod = dllInjection.dllInjection(ProcessAccessHelp::hProcess, selectedFilePath);
	if (hMod && ConfigurationHolder::getConfigObject(DLL_INJECTION_AUTO_UNLOAD)->isTrue())
	{
	if (!dllInjection.unloadDllInProcess(ProcessAccessHelp::hProcess, hMod))
	{
	Logger::printfDialog(TEXT("DLL unloading failed, target %s"), selectedFilePath);
	}
	}

	if (hMod)
	{
	Logger::printfDialog(TEXT("DLL Injection was successful, target %s"), selectedFilePath);
	}
	else
	{
	Logger::printfDialog(TEXT("DLL Injection failed, target %s"), selectedFilePath);
	}
	}
	}

	void MainGui::optionsActionHandler()
	{
	OptionsGui dlgOptions;
	dlgOptions.DoModal();
	}

	void MainGui::clearImportsActionHandler()
	{
	- TreeImports.DeleteAllItems();
	- importsHandling.moduleList.clear();
	+ importsHandling.clearAllImports();
	fillStatusBar();
	}

	void MainGui::pluginActionHandler( int menuItem )
	{
	if(!selectedProcess)
	return;

	DllInjectionPlugin dllInjectionPlugin;

	std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList();

	menuItem -= PLUGIN_MENU_BASE_ID;

	dllInjectionPlugin.hProcess = ProcessAccessHelp::hProcess;
	dllInjectionPlugin.apiReader = &apiReader;

	if (menuItem < (int)scyllaPluginList.size())
	{
	//scylla plugin
	dllInjectionPlugin.injectPlugin(scyllaPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize);
	}
	else
	{
	#ifndef _WIN64

	menuItem -= (int)scyllaPluginList.size();
	//imprec plugin
	dllInjectionPlugin.injectImprecPlugin(imprecPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize);

	#endif
	}

	importsHandling.scanAndFixModuleList();
	importsHandling.displayAllImports();
	fillStatusBar();
	}
	diff --git a/Scylla/MainGui.h b/Scylla/MainGui.h
	index 1d2a169..49a40bc 100644
	--- a/Scylla/MainGui.h
	+++ b/Scylla/MainGui.h
	@@ -1,282 +1,284 @@
	#pragma once

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

	// WTL
	#include <atlbase.h> // base ATL classes
	#include <atlapp.h> // base WTL classes
	#include <atlwin.h> // ATL GUI classes
	#include <atlframe.h> // WTL window frame helpers
	#include <atlmisc.h> // WTL utility classes
	#include <atlcrack.h> // WTL enhanced msg map macros
	#include <atlctrls.h> // WTL controls
	#include <atlddx.h> // WTL dialog data exchange
	#include "multitree.h"

	//#define _CRTDBG_MAP_ALLOC
	//#include <cstdlib>
	//#include <crtdbg.h>

	#include <cstdio>

	#include "Logger.h"
	#include "ProcessLister.h"
	#include "IATSearch.h"
	#include "PickDllGui.h"
	#include "ImportsHandling.h"

	class MainGui : public CDialogImpl<MainGui>, public CWinDataExchange<MainGui>, public CDialogResize<MainGui>, public CMessageFilter
	{
	public:
	enum { IDD = IDD_DLG_MAIN };

	BEGIN_DDX_MAP(MainGui)
	DDX_CONTROL(IDC_TREE_IMPORTS, TreeImportsSubclass) // subclass
	DDX_CONTROL(IDC_TREE_IMPORTS, TreeImports)
	DDX_CONTROL_HANDLE(IDC_CBO_PROCESSLIST, ComboProcessList) // attach
	DDX_CONTROL_HANDLE(IDC_LIST_LOG, ListLog)
	DDX_CONTROL_HANDLE(IDC_EDIT_OEPADDRESS, EditOEPAddress)
	DDX_CONTROL_HANDLE(IDC_EDIT_IATADDRESS, EditIATAddress)
	DDX_CONTROL_HANDLE(IDC_EDIT_IATSIZE, EditIATSize)
	END_DDX_MAP()

	BEGIN_MSG_MAP_EX(MainGui)

	MSG_WM_INITDIALOG(OnInitDialog)
	MSG_WM_DESTROY(OnDestroy)
	MSG_WM_SIZE(OnSize)
	MSG_WM_CONTEXTMENU(OnContextMenu)
	//MSG_WM_LBUTTONDOWN(OnLButtonDown)
	MSG_WM_COMMAND(OnCommand)

	//NOTIFY_HANDLER_EX(IDC_TREE_IMPORTS, NM_CLICK, OnTreeImportsClick)
	NOTIFY_HANDLER_EX(IDC_TREE_IMPORTS, NM_DBLCLK, OnTreeImportsDoubleClick)
	//NOTIFY_HANDLER_EX(IDC_TREE_IMPORTS, NM_RCLICK, OnTreeImportsRightClick)
	//NOTIFY_HANDLER_EX(IDC_TREE_IMPORTS, NM_RDBLCLK, OnTreeImportsRightDoubleClick)
	NOTIFY_HANDLER_EX(IDC_TREE_IMPORTS, TVN_KEYDOWN, OnTreeImportsKeyDown)

	COMMAND_HANDLER_EX(IDC_CBO_PROCESSLIST, CBN_DROPDOWN, OnProcessListDrop)
	COMMAND_HANDLER_EX(IDC_CBO_PROCESSLIST, CBN_SELENDOK, OnProcessListSelected)
	COMMAND_ID_HANDLER_EX(IDC_BTN_PICKDLL, OnPickDLL)
	COMMAND_ID_HANDLER_EX(IDC_BTN_OPTIONS, OnOptions)
	COMMAND_ID_HANDLER_EX(IDC_BTN_DUMP, OnDump)
	COMMAND_ID_HANDLER_EX(IDC_BTN_FIXDUMP, OnFixDump)
	COMMAND_ID_HANDLER_EX(IDC_BTN_PEREBUILD, OnPERebuild)
	COMMAND_ID_HANDLER_EX(IDC_BTN_IATAUTOSEARCH, OnIATAutoSearch)
	COMMAND_ID_HANDLER_EX(IDC_BTN_GETIMPORTS, OnGetImports)
	COMMAND_ID_HANDLER_EX(IDC_BTN_INVALIDIMPORTS, OnInvalidImports)
	COMMAND_ID_HANDLER_EX(IDC_BTN_SUSPECTIMPORTS, OnSuspectImports)
	COMMAND_ID_HANDLER_EX(IDC_BTN_CLEARIMPORTS, OnClearImports)
	COMMAND_ID_HANDLER_EX(ID_FILE_DUMP, OnDump)
	COMMAND_ID_HANDLER_EX(ID_FILE_PEREBUILD, OnPERebuild)
	COMMAND_ID_HANDLER_EX(ID_FILE_FIXDUMP, OnFixDump)
	COMMAND_ID_HANDLER_EX(ID_FILE_EXIT, OnExit)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_SHOWINVALID, OnInvalidImports)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_SHOWSUSPECT, OnSuspectImports)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_INVALIDATESELECTED, OnInvalidateSelected)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_CUTSELECTED, OnCutSelected)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_CLEARIMPORTS, OnClearImports)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_SAVETREE, OnSaveTree)
	COMMAND_ID_HANDLER_EX(ID_IMPORTS_LOADTREE, OnLoadTree)
	COMMAND_ID_HANDLER_EX(ID_TRACE_AUTOTRACE, OnAutotrace)
	COMMAND_ID_HANDLER_EX(ID_MISC_DLLINJECTION, OnDLLInject)
	COMMAND_ID_HANDLER_EX(ID_MISC_OPTIONS, OnOptions)
	COMMAND_ID_HANDLER_EX(ID_HELP_ABOUT, OnAbout)
	COMMAND_ID_HANDLER_EX(IDCANCEL, OnExit)

	REFLECT_NOTIFY_ID(IDC_TREE_IMPORTS)
	CHAIN_MSG_MAP(CDialogResize<MainGui>)

	ALT_MSG_MAP(IDC_TREE_IMPORTS) // message map for subclassed treeview

	MSG_WM_GETDLGCODE(OnTreeImportsSubclassGetDlgCode)
	MSG_WM_CHAR(OnTreeImportsSubclassChar)

	END_MSG_MAP()

	BEGIN_DLGRESIZE_MAP(MainGui)
	DLGRESIZE_CONTROL(IDC_GROUP_ATTACH, DLSZ_SIZE_X)
	DLGRESIZE_CONTROL(IDC_CBO_PROCESSLIST, DLSZ_SIZE_X)
	DLGRESIZE_CONTROL(IDC_BTN_PICKDLL, DLSZ_MOVE_X)

	DLGRESIZE_CONTROL(IDC_GROUP_IMPORTS, DLSZ_SIZE_X \| DLSZ_SIZE_Y)
	DLGRESIZE_CONTROL(IDC_TREE_IMPORTS, DLSZ_SIZE_X \| DLSZ_SIZE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_INVALIDIMPORTS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_SUSPECTIMPORTS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_SAVETREE, DLSZ_MOVE_X \| DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_LOADTREE, DLSZ_MOVE_X \| DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_CLEARIMPORTS, DLSZ_MOVE_X \| DLSZ_MOVE_Y)

	DLGRESIZE_CONTROL(IDC_GROUP_IATINFO, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_STATIC_OEPADDRESS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_STATIC_IATADDRESS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_STATIC_IATSIZE, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_EDIT_OEPADDRESS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_EDIT_IATADDRESS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_EDIT_IATSIZE, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_IATAUTOSEARCH, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_GETIMPORTS, DLSZ_MOVE_Y)

	DLGRESIZE_CONTROL(IDC_GROUP_ACTIONS, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_AUTOTRACE, DLSZ_MOVE_Y)

	DLGRESIZE_CONTROL(IDC_GROUP_DUMP, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_DUMP, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_PEREBUILD, DLSZ_MOVE_Y)
	DLGRESIZE_CONTROL(IDC_BTN_FIXDUMP, DLSZ_MOVE_Y)

	DLGRESIZE_CONTROL(IDC_GROUP_LOG, DLSZ_MOVE_Y \| DLSZ_SIZE_X)
	DLGRESIZE_CONTROL(IDC_LIST_LOG, DLSZ_MOVE_Y \| DLSZ_SIZE_X)
	END_DLGRESIZE_MAP()

	MainGui();

	void addTextToOutputLog(const WCHAR * text);

	protected:

	// Variables

	ProcessLister processLister;
	WCHAR stringBuffer[600];

	ImportsHandling importsHandling;
	ProcessAccessHelp processAccessHelp;
	ApiReader apiReader;

	Process * selectedProcess;

	// File selection stuff

	static const WCHAR filterExe[];
	static const WCHAR filterDll[];
	static const WCHAR filterExeDll[];
	static const WCHAR filterTxt[];

	// Controls

	CMultiSelectTreeViewCtrl TreeImports;
	CComboBox ComboProcessList;
	CEdit EditOEPAddress;
	CEdit EditIATAddress;
	CEdit EditIATSize;
	CListBox ListLog;
	CStatusBarCtrl StatusBar;

	enum StatusParts {
	PART_COUNT = 0,
	PART_INVALID,
	PART_IMAGEBASE,
	PART_MODULE
	};

	CContainedWindow TreeImportsSubclass;

	// Handles

	CIcon hIcon;
	CMenu hMenuImports;
	CMenu hMenuLog;
	CAccelerator accelerators;
	CIcon hIconCheck;
	CIcon hIconWarning;
	CIcon hIconError;

	static const int MenuImportsOffsetTrace = 2;
	static const int MenuImportsTraceOffsetScylla = 2;
	static const int MenuImportsTraceOffsetImpRec = 4;

	protected:

	virtual BOOL PreTranslateMessage(MSG* pMsg);

	// Message handlers

	BOOL OnInitDialog(CWindow wndFocus, LPARAM lInitParam);
	void OnDestroy();
	void OnSize(UINT nType, CSize size);
	void OnLButtonDown(UINT nFlags, CPoint point);
	void OnContextMenu(CWindow wnd, CPoint point);
	void OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl);

	LRESULT OnTreeImportsClick(const NMHDR* pnmh);
	LRESULT OnTreeImportsDoubleClick(const NMHDR* pnmh);
	LRESULT OnTreeImportsRightClick(const NMHDR* pnmh);
	LRESULT OnTreeImportsRightDoubleClick(const NMHDR* pnmh);
	LRESULT OnTreeImportsKeyDown(const NMHDR* pnmh);

	UINT OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg);
	void OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags);

	void OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl);

	void OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnDump(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl);

	void OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl);

	void OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnAutotrace(UINT uNotifyCode, int nID, CWindow wndCtl);

	void OnExit(UINT uNotifyCode, int nID, CWindow wndCtl);
	void OnAbout(UINT uNotifyCode, int nID, CWindow wndCtl);

	// GUI functions

	bool showFileDialog(WCHAR * selectedFile, bool save, const WCHAR * defFileName, const WCHAR * filter = NULL, const WCHAR * defExtension = NULL, const WCHAR * directory = NULL);

	void setupStatusBar();
	void fillStatusBar();
	void fillProcessListComboBox(CComboBox& hCombo);
	void setIconAndDialogCaption();
	void enableDialogControls(BOOL value);

	// Actions

	void pickDllActionHandler();
	void pickApiActionHandler(CTreeItem item);
	void processSelectedActionHandler(int index);
	void showInvalidImportsActionHandler();
	void showSuspectImportsActionHandler();
	+ void deleteSelectedImportsActionHandler();
	+ void invalidateSelectedImportsActionHandler();
	void iatAutosearchActionHandler();
	void getImportsActionHandler();
	void dumpActionHandler();
	void peRebuildActionHandler();
	void startDisassemblerGui(CTreeItem selectedTreeNode);
	void dumpFixActionHandler();
	void showAboutDialog();
	void dllInjectActionHandler();
	void optionsActionHandler();
	void clearImportsActionHandler();
	void pluginActionHandler(int menuItem);

	// Popup menu functions

	void SetupImportsMenuItems(bool isItem, bool isThunk);
	void appendPluginListToMenu(CMenuHandle hMenuTrackPopup);
	void DisplayContextMenuImports(CWindow, CPoint);
	void DisplayContextMenuLog(CWindow, CPoint);

	// Log

	void clearOutputLog();
	bool saveLogToFile(const WCHAR * file);

	// Misc

	DWORD_PTR getOEPFromGui();
	static DWORD_PTR stringToDwordPtr(const WCHAR * hexString);
	};
	diff --git a/Scylla/ProcessAccessHelp.h b/Scylla/ProcessAccessHelp.h
	index fccd537..2d3e63d 100644
	--- a/Scylla/ProcessAccessHelp.h
	+++ b/Scylla/ProcessAccessHelp.h
	@@ -1,202 +1,202 @@
	#pragma once

	#include <stdio.h>
	#include <windows.h>
	#include <tlhelp32.h>
	#include <hash_map>
	#include <map>

	/************************************************************************/
	/* distorm */
	/************************************************************************/
	#include "distorm.h"

	// The number of the array of instructions the decoder function will use to return the disassembled instructions.
	// Play with this value for performance...
	#define MAX_INSTRUCTIONS (200)

	/************************************************************************/

	class ApiInfo;

	class ModuleInfo {
	public:
	WCHAR fullPath[MAX_PATH];
	DWORD_PTR modBaseAddr;
	DWORD modBaseSize;

	bool isAlreadyParsed;
	bool parsing;

	/*
	for iat rebuilding with duplicate entries:

	ntdll = low priority
	kernelbase = low priority
	SHLWAPI = low priority

	kernel32 = high priority

	priority = 1 -> normal/high priority
	priority = 0 -> low priority
	*/
	int priority;

	std::vector<ApiInfo *> apiList;

	ModuleInfo()
	{
	modBaseAddr = 0;
	modBaseSize = 0;
	priority = 1;
	isAlreadyParsed = false;
	parsing = false;
	}

	const WCHAR * getFilename() const
	{
	const WCHAR* slash = wcsrchr(fullPath, L'\\');
	if(slash)
	{
	return slash+1;
	}
	return fullPath;
	}
	};

	class ApiInfo {
	public:
	char name[MAX_PATH];
	- DWORD hint;
	+ WORD hint;
	DWORD_PTR va;
	DWORD_PTR rva;
	- DWORD_PTR ordinal;
	+ WORD ordinal;
	bool isForwarded;
	ModuleInfo * module;
	};

	class ProcessAccessHelp {
	public:
	static HANDLE hProcess; //OpenProcess handle to target process

	static DWORD_PTR targetImageBase;
	static DWORD_PTR targetSizeOfImage;
	static DWORD_PTR maxValidAddress;

	static ModuleInfo * selectedModule;

	static std::vector<ModuleInfo> moduleList; //target process module list
	static std::vector<ModuleInfo> ownModuleList; //own module list

	static const int PE_HEADER_BYTES_COUNT = 2000;

	static BYTE fileHeaderFromDisk[PE_HEADER_BYTES_COUNT];


	//for decomposer
	static _DInst decomposerResult[MAX_INSTRUCTIONS];
	static unsigned int decomposerInstructionsCount;
	static _CodeInfo decomposerCi;

	//distorm :: Decoded instruction information.
	static _DecodedInst decodedInstructions[MAX_INSTRUCTIONS];
	static unsigned int decodedInstructionsCount;
	#ifdef _WIN64
	static const _DecodeType dt = Decode64Bits;
	#else
	static const _DecodeType dt = Decode32Bits;
	#endif

	/*
	* Open a new process handle
	*/
	static bool openProcessHandle(DWORD dwPID);

	static HANDLE NativeOpenProcess(DWORD dwDesiredAccess, DWORD dwProcessId);

	static void closeProcessHandle();

	/*
	* Get all modules from a process
	*/
	static bool getProcessModules(DWORD dwPID, std::vector<ModuleInfo> &moduleList);


	/*
	* file mapping view with different access level
	*/
	static LPVOID createFileMappingViewRead(const WCHAR * filePath);
	static LPVOID createFileMappingViewFull(const WCHAR * filePath);

	/*
	* Create a file mapping view of a file
	*/
	static LPVOID createFileMappingView(const WCHAR * filePath, DWORD accessFile, DWORD flProtect, DWORD accessMap);

	/*
	* Read memory from target process
	*/
	static bool readMemoryFromProcess(DWORD_PTR address, SIZE_T size, LPVOID dataBuffer);

	/*
	* Read memory from file
	*/
	static bool readMemoryFromFile(HANDLE hFile, LONG offset, DWORD size, LPVOID dataBuffer);

	/*
	* Write memory to file
	*/
	static bool writeMemoryToFile(HANDLE hFile, LONG offset, DWORD size, LPCVOID dataBuffer);

	/*
	* Write memory to file end
	*/
	static bool writeMemoryToFileEnd(HANDLE hFile, DWORD size, LPCVOID dataBuffer);

	/*
	* Disassemble Memory
	*/
	static bool disassembleMemory(BYTE * dataBuffer, SIZE_T bufferSize, DWORD_PTR startOffset);

	static bool decomposeMemory(BYTE * dataBuffer, SIZE_T bufferSize, DWORD_PTR startAddress);

	/*
	* Search for pattern
	*/
	static DWORD_PTR findPattern(DWORD_PTR startOffset, DWORD size, BYTE * pattern, const char * mask);

	/*
	* Get process ID by process name
	*/
	static DWORD getProcessByName(const WCHAR * processName);

	/*
	* Get memory region from address
	*/
	bool getMemoryRegionFromAddress(DWORD_PTR address, DWORD_PTR * memoryRegionBase, SIZE_T * memoryRegionSize);


	/*
	* Read PE Header from file
	*/
	static bool readHeaderFromFile(BYTE * buffer, DWORD bufferSize, const WCHAR * filePath);

	static bool readHeaderFromCurrentFile(const WCHAR * filePath);

	/*
	* Get real sizeOfImage value
	*/
	static SIZE_T getSizeOfImageProcess(HANDLE processHandle, DWORD_PTR moduleBase);

	/*
	* Get real sizeOfImage value current process
	*/
	static bool getSizeOfImageCurrentProcess();

	static LONGLONG getFileSize(HANDLE hFile);
	static LONGLONG getFileSize(const WCHAR * filePath);

	static DWORD getEntryPointFromFile(const WCHAR * filePath);

	static bool createBackupFile(const WCHAR * filePath);
	};
	diff --git a/Scylla/Thunks.h b/Scylla/Thunks.h
	index 58967bf..93cf06b 100644
	--- a/Scylla/Thunks.h
	+++ b/Scylla/Thunks.h
	@@ -1,38 +1,42 @@
	#pragma once

	#include <windows.h>
	#include <Commctrl.h>
	#include <map>

	// WTL
	#include <atlbase.h>
	#include <atlapp.h>
	#include <atlctrls.h> //CTreeItem

	-class ImportThunk {
	+class ImportThunk
	+{
	public:
	WCHAR moduleName[MAX_PATH];
	char name[MAX_PATH];
	DWORD_PTR va;
	DWORD_PTR rva;
	- DWORD_PTR ordinal;
	+ WORD ordinal;
	DWORD_PTR apiAddressVA;
	WORD hint;
	bool valid;
	bool suspect;

	CTreeItem hTreeItem;
	+
	+ void invalidate();
	};

	-class ImportModuleThunk {
	+class ImportModuleThunk
	+{
	public:
	WCHAR moduleName[MAX_PATH];
	std::map<DWORD_PTR, ImportThunk> thunkList;

	DWORD_PTR firstThunk;

	CTreeItem hTreeItem;

	DWORD_PTR getFirstThunk() const;
	bool isValid() const;
	};
	diff --git a/Scylla/multitree.h b/Scylla/multitree.h
	index 60ee5f3..59b8dc1 100644
	--- a/Scylla/multitree.h
	+++ b/Scylla/multitree.h
	@@ -1,580 +1,600 @@
	#pragma once

	/////////////////////////////////////////////////////////////////////////////
	// MultiSelectTree - A tree control with multi-select capabilities
	//
	// Written by Bjarke Viksoe ([email protected])
	// Copyright (c) 2005 Bjarke Viksoe.
	//
	// Add the following macro to the parent's message map:
	// REFLECT_NOTIFICATIONS()
	//
	// This code may be used in compiled form in any way you desire. This
	// source file may be redistributed by any means PROVIDING it is
	// not sold for profit without the authors written consent, and
	// providing that this notice and the authors name is included.
	//
	// This file is provided "as is" with no expressed or implied warranty.
	// The author accepts no liability if it causes any damage to you or your
	// computer whatsoever. It's free, so don't hassle me about it.
	//
	// Beware of bugs.
	//

	#ifndef __cplusplus
	#error WTL requires C++ compilation (use a .cpp suffix)
	#endif

	#ifndef __ATLMISC_H__
	#error multitree.h requires atlmisc.h to be included first
	#endif

	#ifndef __ATLCRACK_H__
	#error multitree.h requires atlcrack.h to be included first
	#endif

	#ifndef __ATLCTRLS_H__
	#error multitree.h requires atlctrls.h to be included first
	#endif

	// Extended MultiSelectTree styles
	static const DWORD MTVS_EX_NOMARQUEE = 0x00000001;

	// New control notifications
	static const UINT TVN_ITEMSELECTING = 0x0001;
	static const UINT TVN_ITEMSELECTED = 0x0002;

	-static bool operator==(CTreeItem ti1, CTreeItem ti2)
	+static bool operator==(const CTreeItem& ti1, const CTreeItem& ti2)
	{
	- return ti1.m_hTreeItem == ti2.m_hTreeItem && ti1.m_pTreeView == ti2.m_pTreeView;
	+ return ti1.m_hTreeItem == ti2.m_hTreeItem;
	}

	template< class T, class TBase = CTreeViewCtrlEx, class TWinTraits = CControlWinTraits >
	class ATL_NO_VTABLE CMultiSelectTreeViewImpl :
	public CWindowImpl< T, TBase, TWinTraits >,
	public CCustomDraw< T >
	{
	public:
	DECLARE_WND_SUPERCLASS(NULL, TBase::GetWndClassName())

	DWORD m_dwExStyle; // Additional styles
	CTreeItem m_hExtSelStart; // Item where SHIFT was last pressed
	bool m_bMarquee; // Are we drawing rubberband?
	CPoint m_ptDragStart; // Point where rubberband started
	CPoint m_ptDragOld; // Last mousepos of rubberband

	CSimpleMap<CTreeItem, bool> m_aData;

	CMultiSelectTreeViewImpl() : m_dwExStyle(0), m_bMarquee(false)
	{
	}

	+ CTreeItem _MakeItem(HTREEITEM hItem) const
	+ {
	+ return CTreeItem(hItem, (CTreeViewCtrlEx*)this);
	+ }
	+
	// Operations

	BOOL SubclassWindow(HWND hWnd)
	{
	ATLASSERT(m_hWnd == NULL);
	ATLASSERT(::IsWindow(hWnd));
	+
	BOOL bRet = CWindowImpl< T, TBase, TWinTraits >::SubclassWindow(hWnd);
	if( bRet )
	_Init();
	return bRet;
	}

	DWORD SetExtendedTreeStyle(DWORD dwStyle)
	{
	ATLASSERT(!m_ctrlTree.IsWindow()); // Before control is created, please!
	+
	DWORD dwOldStyle = m_dwTreeStyle;
	m_dwTreeStyle = dwStyle;
	return dwOldStyle;
	}

	- void SelectItem(HTREEITEM hItem, BOOL bSelect)
	+ BOOL SelectItem(HTREEITEM hItem, BOOL bSelect)
	{
	ATLASSERT(::IsWindow(m_hWnd));
	+
	_SelectItem(hItem, bSelect == TRUE);
	if( bSelect )
	TBase::SelectItem(hItem);
	+ return TRUE;
	}

	- void SelectAllItems(BOOL bSelect)
	+ BOOL SelectAllItems(BOOL bSelect)
	{
	ATLASSERT(::IsWindow(m_hWnd));

	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	_SelectItem(i, bSelect == TRUE);
	}
	+ return TRUE;
	+ }
	+
	+ BOOL IsItemSelected(HTREEITEM hItem)
	+ {
	+ ATLASSERT(::IsWindow(m_hWnd));
	+
	+ int iIndex = m_aData.FindKey(_MakeItem(hItem));
	+ if( iIndex >= 0 )
	+ {
	+ return m_aData.GetValueAt(iIndex) ? TRUE : FALSE;
	+ }
	+ return FALSE;
	}

	CTreeItem GetSelectedItem() const
	{
	ATLASSERT(false); // Not usable!
	return GetFirstSelectedItem();
	}

	CTreeItem GetFocusItem() const
	{
	return TBase::GetSelectedItem();
	}

	UINT GetItemState(HTREEITEM hItem, UINT nStateMask) const
	{
	UINT nRes = TBase::GetItemState(hItem, nStateMask);
	if( (nStateMask & TVIS_SELECTED) != 0 )
	{
	- int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
	+ int iIndex = m_aData.FindKey(_MakeItem(hItem));
	if( iIndex >= 0 )
	{
	nRes &= ~TVIS_SELECTED;
	if( m_aData.GetValueAt(iIndex) )
	nRes \|= TVIS_SELECTED;
	}
	}
	return nRes;
	}

	CTreeItem GetFirstSelectedItem() const
	{
	- if( m_aData.GetSize() == 0 )
	- return NULL;
	-
	- for( int i = 0; i < m_aData.GetSize(); i++ )
	+ if( m_aData.GetSize() > 0 )
	{
	- if( m_aData.GetValueAt(i) )
	- return m_aData.GetKeyAt(i);
	+ for( int i = 0; i < m_aData.GetSize(); i++ )
	+ {
	+ if( m_aData.GetValueAt(i) )
	+ return m_aData.GetKeyAt(i);
	+ }
	}
	- return NULL;
	+ return _MakeItem(NULL);
	}

	CTreeItem GetNextSelectedItem(HTREEITEM hItem) const
	{
	- int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
	- if( iIndex < 0 )
	- return NULL;
	-
	- for( int i = iIndex + 1; i < m_aData.GetSize(); i++ )
	+ int iIndex = m_aData.FindKey(_MakeItem(hItem));
	+ if( iIndex >= 0 )
	{
	- if( m_aData.GetValueAt(i) )
	- return m_aData.GetKeyAt(i);
	+ for( int i = iIndex + 1; i < m_aData.GetSize(); i++ )
	+ {
	+ if( m_aData.GetValueAt(i) )
	+ return m_aData.GetKeyAt(i);
	+ }
	}
	- return NULL;
	+ return _MakeItem(NULL);
	}

	int GetSelectedCount() const
	{
	int nCount = 0;
	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	if( m_aData.GetValueAt(i) )
	nCount++;
	}
	return nCount;
	}

	// Implementation

	void _Init()
	{
	ATLASSERT(::IsWindow(m_hWnd));

	ModifyStyle(TVS_SHOWSELALWAYS, 0);
	}

	void _SelectItem(int iIndex, bool bSelect, int action = TVC_UNKNOWN)
	{
	if( iIndex < 0 )
	return;
	bool bSelected = m_aData.GetValueAt(iIndex);
	// Don't change if state is already updated (avoids flicker)
	if( bSelected == bSelect )
	return;
	CTreeItem cItem = m_aData.GetKeyAt(iIndex);
	CWindow parent = GetParent();
	// Send notifications
	NMTREEVIEW nmtv = { 0 };
	nmtv.hdr.code = TVN_ITEMSELECTING;
	nmtv.hdr.hwndFrom = m_hWnd;
	nmtv.hdr.idFrom = GetDlgCtrlID();
	nmtv.action = action;
	nmtv.itemNew.hItem = cItem;
	nmtv.itemNew.lParam = GetItemData(cItem);
	nmtv.itemNew.state = bSelect ? TVIS_SELECTED : 0;
	nmtv.itemNew.stateMask = TVIS_SELECTED;
	if( parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv) != 0 )
	return;
	// Change state
	m_aData.SetAtIndex(iIndex, cItem, bSelect);
	// Repaint item
	CRect rcItem;
	if( GetItemRect(cItem, &rcItem, FALSE) )
	InvalidateRect(&rcItem, TRUE);
	// More notifications
	nmtv.hdr.code = TVN_ITEMSELECTED;
	parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv);
	}

	void _SelectItem(HTREEITEM hItem, bool bSelect, int action = TVC_UNKNOWN)
	{
	- _SelectItem(m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)), bSelect, action);
	+ _SelectItem(m_aData.FindKey(_MakeItem(hItem)), bSelect, action);
	}

	void _SelectTree(HTREEITEM hItem, HTREEITEM hGoal, int action)
	{
	if( !_SelectTreeSub(hItem, hGoal, action) )
	return;
	hItem = GetParentItem(hItem);
	while( (hItem = GetNextSiblingItem(hItem)) != NULL )
	{
	if( !_SelectTreeSub(hItem, hGoal, action) )
	return;
	}
	}

	bool _SelectTreeSub(HTREEITEM hItem, HTREEITEM hGoal, int action)
	{
	while( hItem != NULL )
	{
	_SelectItem(hItem, true, action);
	if( hItem == hGoal )
	return false;
	if( (TBase::GetItemState(hItem, TVIS_EXPANDED) & TVIS_EXPANDED) != 0 )
	{
	if( !_SelectTreeSub(GetChildItem(hItem), hGoal, action) )
	return false;
	}
	hItem = GetNextSiblingItem(hItem);
	}
	return true;
	}

	/*
	TODO: keep list of 'temporary' selected items
	if you select an item with the mouse and it gets
	out of the selection rectangle it stays selected.
	it should unselect it but keep old selected items
	(when holding CTRL)
	*/
	void _SelectBox(CRect rc)
	{
	- CTreeItem hItem = GetFirstVisibleItem();
	+ HTREEITEM hItem = GetFirstVisibleItem();
	while( hItem != NULL )
	{
	- CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this);
	- int i = m_aData.FindKey(cItem);
	+ int i = m_aData.FindKey(_MakeItem(hItem));
	if(i >= 0 && !m_aData.GetValueAt(i)) // ignore already selected
	{
	CRect rcItem, rcTemp;
	GetItemRect(hItem, &rcItem, TRUE);
	_SelectItem(hItem, rcTemp.IntersectRect(&rcItem, &rc) == TRUE, TVC_BYMOUSE);
	}
	hItem = GetNextVisibleItem(hItem);
	}
	}

	void _DrawDragRect(CPoint pt)
	{
	CClientDC dc(m_hWnd);
	CSize szFrame(1, 1);
	CRect rect(m_ptDragStart, pt);
	rect.NormalizeRect();
	//CRect rectOld(m_ptDragStart, m_ptDragOld);
	//rectOld.NormalizeRect();
	dc.DrawDragRect(&rect, szFrame, NULL/&rectOld/, szFrame);
	}

	// Message map and handlers

	BEGIN_MSG_MAP_EX(CMultiSelectTreeViewImpl)
	MESSAGE_HANDLER_EX(WM_CREATE, OnCreate)
	MSG_WM_DESTROY(OnDestroy)
	MSG_WM_KEYDOWN(OnKeyDown)
	MSG_WM_KEYUP(OnKeyUp)
	MSG_WM_CHAR(OnChar)
	MSG_WM_SETFOCUS(OnSetFocus)
	MSG_WM_LBUTTONDOWN(OnLButtonDown)
	MSG_WM_LBUTTONUP(OnLButtonUp)
	MSG_WM_MOUSEMOVE(OnMouseMove)
	MSG_WM_CAPTURECHANGED(OnCaptureChanged)
	MESSAGE_HANDLER_EX(TVM_INSERTITEM, OnInsertItem)
	REFLECTED_NOTIFY_CODE_HANDLER_EX(TVN_DELETEITEM, OnDeleteItem)
	CHAIN_MSG_MAP_ALT( CCustomDraw< T >, 1 )
	END_MSG_MAP()

	LRESULT OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
	LRESULT lRes = DefWindowProc();
	_Init();
	return lRes;
	}

	void OnDestroy()
	{
	m_aData.RemoveAll();
	SetMsgHandled(FALSE);
	}

	void OnKeyDown(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
	if( nChar == VK_SHIFT )
	m_hExtSelStart = GetFocusItem();

	if( ::GetKeyState(VK_SHIFT) < 0 && m_hExtSelStart == GetFocusItem() )
	{
	switch( nChar )
	{
	case VK_UP:
	case VK_DOWN:
	case VK_HOME:
	case VK_END:
	case VK_NEXT:
	case VK_PRIOR:
	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	_SelectItem(i, false, TVC_BYKEYBOARD);
	}
	}
	}
	SetMsgHandled(FALSE);
	}

	void OnKeyUp(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
	if( ::GetKeyState(VK_SHIFT) < 0 )
	{
	switch( nChar )
	{
	case VK_UP:
	case VK_DOWN:
	case VK_HOME:
	case VK_END:
	case VK_NEXT:
	case VK_PRIOR:
	HTREEITEM hItem = GetFocusItem();
	// Is current or first-shift-item the upper item?
	CRect rcItem1, rcItem2;
	GetItemRect(m_hExtSelStart, &rcItem1, TRUE);
	GetItemRect(hItem, &rcItem2, TRUE);
	// Select from current item to item where SHIFT was pressed
	if( rcItem1.top > rcItem2.top )
	_SelectTree(hItem, m_hExtSelStart, TVC_BYKEYBOARD);
	else
	_SelectTree(m_hExtSelStart, hItem, TVC_BYKEYBOARD);
	_SelectItem(hItem, true, TVC_BYKEYBOARD);
	}
	}
	SetMsgHandled(FALSE);
	}

	void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
	if( nChar == VK_SPACE )
	{
	HTREEITEM hItem = GetFocusItem();
	_SelectItem(hItem, (GetItemState(hItem, TVIS_SELECTED) & TVIS_SELECTED) == 0, TVC_BYKEYBOARD);
	return;
	}
	SetMsgHandled(FALSE);
	}

	void OnSetFocus(CWindow wndOld)
	{
	DefWindowProc();
	// FIX: We really need the focus-rectangle in this control since it
	// improves the navigation a lot. So let's ask Windows to display it.
	SendMessage(WM_UPDATEUISTATE, MAKEWPARAM(UIS_CLEAR, UISF_HIDEFOCUS));
	}

	void OnLButtonDown(UINT nFlags, CPoint point)
	{
	SetMsgHandled(FALSE);

	// Hit-test and figure out where we're clicking...
	TVHITTESTINFO hti = { 0 };
	hti.pt = point;
	HTREEITEM hItem = HitTest(&hti);
	if( (hItem == NULL \|\| (hti.flags & TVHT_ONITEMRIGHT) != 0) )
	{
	if( (m_dwExStyle & MTVS_EX_NOMARQUEE) == 0 && ::DragDetect(m_hWnd, point) )
	{
	// Great we're dragging a rubber-band
	// Clear selection of CTRL is not down
	if( ::GetKeyState(VK_CONTROL) >= 0 )
	{
	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	_SelectItem(i, false, TVC_BYMOUSE);
	}
	UpdateWindow();
	}
	// Now start drawing the rubber-band...
	SetCapture();
	m_ptDragStart = m_ptDragOld = point;
	_DrawDragRect(point);
	m_bMarquee = true;
	SetMsgHandled(TRUE);
	return;
	}
	}

	if( hItem == NULL )
	return;

	if( (hti.flags & TVHT_ONITEMBUTTON) != 0 )
	return;

	// Great, let's do an advanced selection
	if( (hti.flags & TVHT_ONITEMRIGHT) != 0 )
	{
	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	_SelectItem(i, false, TVC_BYMOUSE);
	}
	return;
	}
	- int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
	+ int iIndex = m_aData.FindKey(_MakeItem(hItem));
	if( iIndex < 0 )
	return;
	// Simulate drag'n'drop?
	if( m_aData.GetValueAt(iIndex) && (GetStyle() & TVS_DISABLEDRAGDROP) == 0 && ::DragDetect(m_hWnd, point) )
	{
	NMTREEVIEW nmtv = { 0 };
	nmtv.hdr.code = TVN_BEGINDRAG;
	nmtv.hdr.hwndFrom = m_hWnd;
	nmtv.hdr.idFrom = GetDlgCtrlID();
	nmtv.itemNew.hItem = hItem;
	nmtv.itemNew.lParam = GetItemData(hItem);
	CWindow parent = GetParent();
	parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv);
	}
	bool bSelected = m_aData.GetValueAt(iIndex);
	if( ::GetKeyState(VK_SHIFT) < 0 )
	{
	// Is current or first-shift-item the upper item?
	CRect rcItem1, rcItem2;
	GetItemRect(m_hExtSelStart, &rcItem1, TRUE);
	GetItemRect(hItem, &rcItem2, TRUE);
	// Select from current item to item where SHIFT was pressed
	if( rcItem1.top > rcItem2.top )
	_SelectTree(hItem, m_hExtSelStart, TVC_BYMOUSE);
	else
	_SelectTree(m_hExtSelStart, hItem, TVC_BYMOUSE);
	}
	else if( ::GetKeyState(VK_CONTROL) < 0 )
	{
	// Just toggle item
	_SelectItem(iIndex, !bSelected, TVC_BYMOUSE);
	}
	else
	{
	// Remove current selection and replace it with clicked item
	for( int i = 0; i < m_aData.GetSize(); i++ )
	{
	_SelectItem(i, i == iIndex, TVC_BYMOUSE);
	}
	}
	}

	void OnLButtonUp(UINT nFlags, CPoint point)
	{
	if( m_bMarquee )
	ReleaseCapture();
	SetMsgHandled(FALSE);
	}

	void OnMouseMove(UINT nFlags, CPoint point)
	{
	if( m_bMarquee )
	{
	CRect rc(m_ptDragStart, point);
	_DrawDragRect(m_ptDragOld);
	rc.NormalizeRect();
	_SelectBox(rc);
	UpdateWindow();
	_DrawDragRect(point);
	m_ptDragOld = point;
	}
	SetMsgHandled(FALSE);
	}

	void OnCaptureChanged(CWindow wnd)
	{
	if( m_bMarquee )
	{
	_DrawDragRect(m_ptDragOld);
	m_bMarquee = false;
	}
	SetMsgHandled(FALSE);
	}

	LRESULT OnInsertItem(UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
	HTREEITEM hItem = (HTREEITEM) DefWindowProc(uMsg, wParam, lParam);
	if( hItem == NULL )
	return (LRESULT) hItem;
	- CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this);
	// We manage a bit of extra information for each item. We'll store
	// this in an ATL::CSimpleMap. Not a particular speedy structure for lookups.
	// Don't keep too many items around in the tree!
	- m_aData.Add(cItem, false);
	+ m_aData.Add(_MakeItem(hItem), false);
	return (LRESULT) hItem;
	}

	LRESULT OnDeleteItem(NMHDR* pnmh)
	{
	const NMTREEVIEW* lpNMTV = (NMTREEVIEW*) pnmh;
	- m_aData.Remove(CTreeItem(lpNMTV->itemNew.hItem, (CTreeViewCtrlEx*)this));
	+ m_aData.Remove(_MakeItem(lpNMTV->itemNew.hItem));
	return 0;
	}

	// Custom Draw

	DWORD OnPrePaint(int /idCtrl/, NMCUSTOMDRAW* /lpNMCustomDraw/)
	{
	return CDRF_NOTIFYITEMDRAW; // We need per-item notifications
	}

	DWORD OnItemPrePaint(int /idCtrl/, NMCUSTOMDRAW* lpNMCustomDraw)
	{
	NMTVCUSTOMDRAW* lpTVCD = (NMTVCUSTOMDRAW*) lpNMCustomDraw;
	HTREEITEM hItem = (HTREEITEM) lpTVCD->nmcd.dwItemSpec;
	- int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
	+ int iIndex = m_aData.FindKey(_MakeItem(hItem));
	if( iIndex >= 0 )
	{
	bool bSelected = m_aData.GetValueAt(iIndex);
	// Trick TreeView into displaying correct selection colors
	if( bSelected )
	{
	lpTVCD->clrText = ::GetSysColor(COLOR_HIGHLIGHTTEXT);
	lpTVCD->clrTextBk = ::GetSysColor(COLOR_HIGHLIGHT);
	}
	else
	{
	// Special case of tree-item actually have selection, but our
	// state says it is currently not selected (CTRL+click on same item twice).
	if( (lpTVCD->nmcd.uItemState & CDIS_SELECTED) != 0 )
	{
	COLORREF clrText = GetTextColor();
	if( clrText == CLR_NONE )
	clrText = ::GetSysColor(COLOR_WINDOWTEXT);
	COLORREF clrBack = GetBkColor();
	if( clrBack == CLR_NONE )
	clrBack = ::GetSysColor(COLOR_WINDOW);
	//CDCHandle dc = lpTVCD->nmcd.hdc;
	//dc.SetTextColor(clrText);
	//dc.SetBkColor(clrBack);
	lpTVCD->clrText = clrText;
	lpTVCD->clrTextBk = clrBack;
	}
	}
	return CDRF_NEWFONT;
	}
	return CDRF_DODEFAULT;
	}
	};

	class CMultiSelectTreeViewCtrl : public CMultiSelectTreeViewImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS> >
	{
	public:
	DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName())
	};
	diff --git a/Scylla/resource.h b/Scylla/resource.h
	index ec931be..6819568 100644
	Binary files a/Scylla/resource.h and b/Scylla/resource.h differ

File Metadata

Mime Type: text/x-diff
Expires: Tue, Oct 28, 10:13 PM (1 d, 4 h)
Storage Engine: local-disk
Storage Format: Raw Data
Storage Handle: 0c/bf/9490ee97f19649dbca17119ed7ac

No OneTemporaryActions

View Options

File Metadata

Event Timeline

No OneTemporary
Actions