No OneTemporary
Actions

Size

47 KB

Subscribers

None

View Options

	diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp
	index 7ee228e..79deb91 100644
	--- a/Scylla/MainGui.cpp
	+++ b/Scylla/MainGui.cpp
	@@ -1,1248 +1,1244 @@
	#include "MainGui.h"

	#include "Architecture.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 "Scylla.h"
	#include "AboutGui.h"
	#include "OptionsGui.h"
	#include "TreeImportExport.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";
	const WCHAR MainGui::filterXml[] = L"XML file (.xml)\0.xml\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);
	*/

	Scylla::init();

	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);
	}

	BOOL MainGui::PreTranslateMessage(MSG* pMsg)
	{
	if(accelerators.TranslateAccelerator(m_hWnd, pMsg))
	{
	return TRUE; // handled keyboard shortcuts
	}
	else if(IsDialogMessage(pMsg))
	{
	return TRUE; // handled dialog messages
	}

	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;
	}
	}

	// register ourselves to receive PreTranslateMessage
	CMessageLoop* pLoop = _Module.GetMessageLoop();
	pLoop->AddMessageFilter(this);

	setupStatusBar();

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

	Scylla::windowLog.setWindow(ListLog);

	appendPluginListToMenu(hMenuImports.GetSubMenu(0));
	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::OnContextMenu(CWindow wnd, CPoint point)
	{
	switch(wnd.GetDlgCtrlID())
	{
	case IDC_TREE_IMPORTS:
	DisplayContextMenuImports(wnd, point);
	return;
	case IDC_LIST_LOG:
	DisplayContextMenuLog(wnd, point);
	return;
	}

	SetMsgHandled(FALSE);
	}

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

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

	// Get item under cursor
	CTreeItem over = findTreeItem(CPoint(GetMessagePos()), true);
	if(over && importsHandling.isImport(over))
	{
	pickApiActionHandler(over);
	}

	return 0;
	}

	LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh)
	{
	const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh;
	switch(tkd->wVKey)
	{
	case VK_RETURN:
	{
	CTreeItem selected = TreeImports.GetFocusItem();
	if(!selected.IsNull() && importsHandling.isImport(selected))
	{
	pickApiActionHandler(selected);
	}
	}
	return 1;
	case VK_DELETE:
	deleteSelectedImportsActionHandler();
	return 1;
	}

	SetMsgHandled(FALSE);
	return 0;
	}

	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)
	{
	invalidateSelectedImportsActionHandler();
	}

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

	void MainGui::OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	saveTreeActionHandler();
	}

	void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	loadTreeActionHandler();
	}

	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, 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::updateStatusBar()
	{
	// Rewrite ImportsHandling so we get these easily
	unsigned int totalImports = importsHandling.thunkCount();
	unsigned int invalidImports = importsHandling.invalidThunkCount();

	// \t = center, \t\t = right-align
	swprintf_s(stringBuffer, L"\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, L"\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, L"\tImagebase: " 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] = L'\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(APPNAME L" " ARCHITECTURE L" " APPVERSION);
	}

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

	PickDllGui dlgPickDll(ProcessAccessHelp::moduleList);
	if(dlgPickDll.DoModal())
	{
	//get selected module
	ProcessAccessHelp::selectedModule = dlgPickDll.getSelectedModule();

	ProcessAccessHelp::targetImageBase = ProcessAccessHelp::selectedModule->modBaseAddr;

	Scylla::windowLog.log(L"->>> Module %s selected.", ProcessAccessHelp::selectedModule->getFilename());
	Scylla::windowLog.log(L"Imagebase: " PRINTF_DWORD_PTR_FULL L" Size: %08X", ProcessAccessHelp::selectedModule->modBaseAddr, ProcessAccessHelp::selectedModule->modBaseSize);
	}
	else
	{
	ProcessAccessHelp::selectedModule = 0;
	}

	updateStatusBar();
	}

	void MainGui::pickApiActionHandler(CTreeItem item)
	{
	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)
	{
	importsHandling.setImport(item, api->module->getFilename(), api->name, api->ordinal, api->hint, true, api->isForwarded);
	}
	}

	updateStatusBar();
	}

	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, L"Can't read memory at " 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 = Scylla::processLister.getProcessList();
	Process &process = processList.at(index);
	selectedProcess = 0;

	clearImportsActionHandler();

	Scylla::windowLog.log(L"Analyzing %s", process.fullPath);

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

	if (!ProcessAccessHelp::openProcessHandle(process.PID))
	{
	enableDialogControls(FALSE);
	Scylla::windowLog.log(L"Error: Cannot open process handle.");
	updateStatusBar();
	return;
	}

	ProcessAccessHelp::getProcessModules(process.PID, ProcessAccessHelp::moduleList);

	apiReader.readApisFromModuleList();

	Scylla::windowLog.log(L"Loading modules done.");

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

	ProcessAccessHelp::getSizeOfImageCurrentProcess();

	process.imageSize = (DWORD)ProcessAccessHelp::targetSizeOfImage;


	Scylla::windowLog.log(L"Imagebase: " PRINTF_DWORD_PTR_FULL L" Size: %08X", process.imageBase, process.imageSize);

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

	EditOEPAddress.SetValue(process.entryPoint + process.imageBase);

	selectedProcess = &process;
	enableDialogControls(TRUE);

	updateStatusBar();
	}

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

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

	for (size_t i = 0; i < processList.size(); i++)
	{
	swprintf_s(stringBuffer, L"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;
	size_t bufsize = 0;
	for(int i = 0; i < ListLog.GetCount(); i++)
	{
	size_t 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, (DWORD)(size * sizeof(WCHAR)), buffer);
	}
	delete[] buffer;
	CloseHandle(hFile);
	}
	return success;
	}

	void MainGui::showInvalidImportsActionHandler()
	{
	importsHandling.selectImports(true, false);
	GotoDlgCtrl(TreeImports);
	}

	void MainGui::showSuspectImportsActionHandler()
	{
	importsHandling.selectImports(false, true);
	GotoDlgCtrl(TreeImports);
	}

	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);
	}
	updateStatusBar();
	}

	void MainGui::invalidateSelectedImportsActionHandler()
	{
	CTreeItem selected = TreeImports.GetFirstSelectedItem();
	while(!selected.IsNull())
	{
	if(importsHandling.isImport(selected))
	{
	importsHandling.invalidateImport(selected);
	}
	selected = TreeImports.GetNextSelectedItem(selected);
	}
	updateStatusBar();
	}

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

	WCHAR selectedFilePath[MAX_PATH];
	- TreeImportExport treeIO;
	- DWORD_PTR addrOEP = 0;
	- DWORD_PTR addrIAT = 0;
	- DWORD sizeIAT = 0;
	-
	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if(showFileDialog(selectedFilePath, false, NULL, filterXml, NULL, stringBuffer))
	{
	- if(!treeIO.importTreeList(selectedFilePath, importsHandling.moduleList, &addrOEP, &addrIAT, &sizeIAT))
	+ TreeImportExport treeIO(selectedFilePath);
	+ DWORD_PTR addrOEP = 0;
	+ DWORD_PTR addrIAT = 0;
	+ DWORD sizeIAT = 0;
	+
	+ if(!treeIO.importTreeList(importsHandling.moduleList, &addrOEP, &addrIAT, &sizeIAT))
	{
	Scylla::windowLog.log(L"Loading tree file failed %s", selectedFilePath);
	MessageBox(L"Loading tree file failed.", L"Failure", MB_ICONERROR);
	}
	else
	{
	EditOEPAddress.SetValue(addrOEP);
	EditIATAddress.SetValue(addrIAT);
	EditIATSize.SetValue(sizeIAT);

	importsHandling.displayAllImports();
	updateStatusBar();

	Scylla::windowLog.log(L"Loaded tree file %s", selectedFilePath);
	Scylla::windowLog.log(L"-> OEP: " PRINTF_DWORD_PTR_FULL, addrOEP);
	Scylla::windowLog.log(L"-> IAT: " PRINTF_DWORD_PTR_FULL L" Size: " PRINTF_DWORD_PTR, addrIAT, sizeIAT);
	}
	}
	}

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

	WCHAR selectedFilePath[MAX_PATH];
	- TreeImportExport treeIO;
	- DWORD_PTR addrOEP;
	- DWORD_PTR addrIAT;
	- DWORD sizeIAT;
	-
	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if(showFileDialog(selectedFilePath, true, NULL, filterXml, L"xml", stringBuffer))
	{
	- addrOEP = EditOEPAddress.GetValue();
	- addrIAT = EditIATAddress.GetValue();
	- sizeIAT = EditIATSize.GetValue();
	+ TreeImportExport treeIO(selectedFilePath);
	+ DWORD_PTR addrOEP = EditOEPAddress.GetValue();
	+ DWORD_PTR addrIAT = EditIATAddress.GetValue();
	+ DWORD sizeIAT = EditIATSize.GetValue();

	- if(!treeIO.exportTreeList(selectedFilePath, importsHandling.moduleList, selectedProcess, addrOEP, addrIAT, sizeIAT))
	+ if(!treeIO.exportTreeList(importsHandling.moduleList, selectedProcess, addrOEP, addrIAT, sizeIAT))
	{
	Scylla::windowLog.log(L"Saving tree file failed %s", selectedFilePath);
	MessageBox(L"Saving tree file failed.", L"Failure", MB_ICONERROR);
	}
	else
	{
	Scylla::windowLog.log(L"Saved tree file %s", selectedFilePath);
	}
	}
	}

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

	if(!selectedProcess)
	return;

	if(EditOEPAddress.GetWindowTextLength() > 0)
	{
	searchAddress = EditOEPAddress.GetValue();
	if (searchAddress)
	{
	if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT))
	{
	Scylla::windowLog.log(L"IAT found at VA " PRINTF_DWORD_PTR_FULL L" RVA " PRINTF_DWORD_PTR_FULL L" Size 0x%04X (%d)", addressIAT, addressIAT - ProcessAccessHelp::targetImageBase, sizeIAT, sizeIAT);

	EditIATAddress.SetValue(addressIAT);
	EditIATSize.SetValue(sizeIAT);

	swprintf_s(stringBuffer, L"IAT found:\r\n\r\nStart: " PRINTF_DWORD_PTR_FULL L"\r\nSize: 0x%04X (%d) ", addressIAT, sizeIAT, sizeIAT);
	MessageBox(stringBuffer, L"IAT found", MB_ICONINFORMATION);
	}
	else
	{
	Scylla::windowLog.log(L"IAT not found at OEP " PRINTF_DWORD_PTR_FULL L"!", searchAddress);
	}
	}
	}
	}

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

	DWORD_PTR addressIAT = EditIATAddress.GetValue();
	DWORD sizeIAT = EditIATSize.GetValue();

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

	updateStatusBar();
	}
	}

	void MainGui::SetupImportsMenuItems(CTreeItem item)
	{
	bool isItem, isImport = false;
	isItem = !item.IsNull();
	if(isItem)
	{
	isImport = importsHandling.isImport(item);
	}

	CMenuHandle hSub = hMenuImports.GetSubMenu(0);

	UINT itemOnly = isItem ? MF_ENABLED : MF_GRAYED;
	UINT importOnly = isImport ? MF_ENABLED : MF_GRAYED;

	hSub.EnableMenuItem(ID__INVALIDATE, itemOnly);
	hSub.EnableMenuItem(ID__DISASSEMBLE, importOnly);
	hSub.EnableMenuItem(ID__CUTTHUNK, importOnly);

	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.GetFocusItem();
	if(over)
	{
	over.EnsureVisible();
	over.GetRect(&pos, TRUE);
	TreeImports.ClientToScreen(&pos);
	}
	else
	{
	TreeImports.GetWindowRect(&pos);
	}
	pt = pos.TopLeft();
	}
	else
	{
	// Get item under cursor
	over = findTreeItem(pt, true);
	}

	SetupImportsMenuItems(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)(Scylla::plugins.getScyllaPluginList().size() + Scylla::plugins.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__INVALIDATE:
	if(importsHandling.isModule(over))
	importsHandling.invalidateModule(over);
	else
	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.cutImport(over);
	break;
	case ID__DELETETREENODE:
	importsHandling.cutModule(importsHandling.isImport(over) ? over.GetParent() : over);
	break;
	}
	}

	updateStatusBar();
	}

	void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt)
	{
	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];
	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt", stringBuffer))
	{
	saveLogToFile(selectedFilePath);
	}
	break;
	case ID__CLEAR:
	clearOutputLog();
	break;
	}
	}
	}

	void MainGui::appendPluginListToMenu(CMenuHandle hMenu)
	{
	std::vector<Plugin> &scyllaPluginList = Scylla::plugins.getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = Scylla::plugins.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";
	}

	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension, stringBuffer))
	{
	if (ProcessAccessHelp::selectedModule)
	{
	//dump DLL

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

	peDump.useHeaderFromDisk = Scylla::config[USE_PE_HEADER_FROM_DISK].isTrue();
	if (peDump.dumpCompleteProcessToDisk(selectedFilePath))
	{
	Scylla::windowLog.log(L"Dump success %s", selectedFilePath);
	}
	else
	{
	Scylla::windowLog.log(L"Error: Cannot dump image.");
	MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR);
	}
	}
	}

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

	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if(showFileDialog(selectedFilePath, false, NULL, filterExeDll, NULL, stringBuffer))
	{
	if (Scylla::config[CREATE_BACKUP].isTrue())
	{
	if (!ProcessAccessHelp::createBackupFile(selectedFilePath))
	{
	Scylla::windowLog.log(L"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)
	{
	Scylla::windowLog.log(L"Rebuild failed %s", selectedFilePath);
	MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR);
	}
	else
	{
	peRebuild.truncateFile(selectedFilePath, newSize);

	Scylla::windowLog.log(L"Rebuild success %s", selectedFilePath);
	Scylla::windowLog.log(L"-> Old file size 0x%08X new file size 0x%08X (%d %%)", (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) );
	}
	}
	}

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

	if (TreeImports.GetCount() < 2)
	{
	Scylla::windowLog.log(L"Nothing to rebuild");
	return;
	}

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

	if (ProcessAccessHelp::selectedModule)
	{
	fileFilter = filterDll;
	}
	else
	{
	fileFilter = filterExe;
	}

	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if (showFileDialog(selectedFilePath, false, NULL, fileFilter, NULL, stringBuffer))
	{
	wcscpy_s(newFilePath, 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, L"_SCY");

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

	ImportRebuild importRebuild;
	if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList))
	{
	Scylla::windowLog.log(L"Import Rebuild success %s", newFilePath);
	}
	else
	{
	Scylla::windowLog.log(L"Import Rebuild failed %s", selectedFilePath);
	MessageBox(L"Import Rebuild failed", L"Failure", MB_ICONERROR);
	}
	}
	}

	void MainGui::enableDialogControls(BOOL 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();

	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_IMPORTS_SAVETREE, valMenu);
	menu.EnableMenuItem(ID_IMPORTS_LOADTREE, 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);
	menu.EnableMenuItem(ID_TRACE_AUTOTRACE, MF_GRAYED);
	}

	CTreeItem MainGui::findTreeItem(CPoint pt, bool screenCoordinates)
	{
	if(screenCoordinates)
	{
	TreeImports.ScreenToClient(&pt);
	}

	UINT flags;
	CTreeItem over = TreeImports.HitTest(pt, &flags);
	if(over)
	{
	if(!(flags & TVHT_ONITEM))
	{
	over.m_hTreeItem = NULL;
	}
	}

	return over;
	}

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

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

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

	getCurrentModulePath(stringBuffer, _countof(stringBuffer));
	if (showFileDialog(selectedFilePath, false, NULL, filterDll, NULL, stringBuffer))
	{
	hMod = dllInjection.dllInjection(ProcessAccessHelp::hProcess, selectedFilePath);
	if (hMod && Scylla::config[DLL_INJECTION_AUTO_UNLOAD].isTrue())
	{
	if (!dllInjection.unloadDllInProcess(ProcessAccessHelp::hProcess, hMod))
	{
	Scylla::windowLog.log(L"DLL unloading failed, target %s", selectedFilePath);
	}
	}

	if (hMod)
	{
	Scylla::windowLog.log(L"DLL Injection was successful, target %s", selectedFilePath);
	}
	else
	{
	Scylla::windowLog.log(L"DLL Injection failed, target %s", selectedFilePath);
	}
	}
	}

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

	void MainGui::clearImportsActionHandler()
	{
	importsHandling.clearAllImports();
	updateStatusBar();
	}

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

	DllInjectionPlugin dllInjectionPlugin;

	std::vector<Plugin> &scyllaPluginList = Scylla::plugins.getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = Scylla::plugins.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();
	updateStatusBar();
	}

	bool MainGui::getCurrentModulePath(WCHAR * buffer, size_t bufferSize)
	{
	if(!selectedProcess)
	return false;

	if(ProcessAccessHelp::selectedModule)
	{
	wcscpy_s(buffer, bufferSize, ProcessAccessHelp::selectedModule->fullPath);
	}
	else
	{
	wcscpy_s(buffer, bufferSize, selectedProcess->fullPath);
	}

	WCHAR * slash = wcsrchr(buffer, L'\\');
	if(slash)
	{
	*(slash+1) = L'\0';
	}

	return true;
	}
	diff --git a/Scylla/TreeImportExport.cpp b/Scylla/TreeImportExport.cpp
	index e17c027..0dc54c2 100644
	--- a/Scylla/TreeImportExport.cpp
	+++ b/Scylla/TreeImportExport.cpp
	@@ -1,337 +1,296 @@
	#include "TreeImportExport.h"
	#include "Architecture.h"
	#include "Scylla.h"
	#include "StringConversion.h"

	-bool TreeImportExport::exportTreeList(const WCHAR * targetXmlFile, std::map<DWORD_PTR, ImportModuleThunk> & moduleList, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT)
	+TreeImportExport::TreeImportExport(const WCHAR * targetXmlFile)
	+{
	+ wcscpy_s(xmlPath, targetXmlFile);
	+}
	+
	+bool TreeImportExport::exportTreeList(const std::map<DWORD_PTR, ImportModuleThunk> & moduleList, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT)
	{
	TiXmlDocument doc;

	- TiXmlDeclaration * decl = new TiXmlDeclaration( "1.0", "", "");
	+ TiXmlDeclaration * decl = new TiXmlDeclaration("1.0", "", "");
	doc.LinkEndChild(decl);

	TiXmlElement * rootElement = new TiXmlElement("target");

	- setTargetInformation(rootElement, process,addressOEP,addressIAT,sizeIAT);
	+ setTargetInformation(rootElement, process, addressOEP, addressIAT, sizeIAT);

	addModuleListToRootElement(rootElement, moduleList);

	doc.LinkEndChild(rootElement);

	- return saveXmlToFile(doc,targetXmlFile);
	+ return saveXmlToFile(doc, xmlPath);
	}

	-bool TreeImportExport::importTreeList(const WCHAR * targetXmlFile, std::map<DWORD_PTR, ImportModuleThunk> & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT)
	+bool TreeImportExport::importTreeList(std::map<DWORD_PTR, ImportModuleThunk> & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT)
	{
	- TiXmlElement * targetElement;
	- TiXmlDocument doc;
	- char * buffer = readXmlFile(targetXmlFile);
	- int count = 0;
	-
	moduleList.clear();
	+ addressOEP = addressIAT = 0;
	+ *sizeIAT = 0;

	- if (buffer)
	+ TiXmlDocument doc;
	+ if(!readXmlFile(doc, xmlPath))
	{
	- doc.Parse(buffer);
	- if (doc.Error())
	- {
	- Scylla::windowLog.log(L"Load Tree :: Error parsing xml %S: %S\r\n", doc.Value(), doc.ErrorDesc());
	- delete [] buffer;
	- return false;
	- }
	-
	- targetElement = doc.FirstChildElement();
	-
	- *addressOEP = ConvertStringToDwordPtr(targetElement->Attribute("oep_va"));
	- *addressIAT = ConvertStringToDwordPtr(targetElement->Attribute("iat_va"));
	- *sizeIAT = (DWORD)ConvertStringToDwordPtr(targetElement->Attribute("iat_size"));
	+ Scylla::windowLog.log(L"Load Tree :: Error parsing xml %S: %S\r\n", doc.Value(), doc.ErrorDesc());
	+ return false;
	+ }

	- //test = targetElement->Attribute("filename");
	+ TiXmlElement * targetElement = doc.FirstChildElement();

	- parseAllElementModules(targetElement, moduleList);
	+ *addressOEP = ConvertStringToDwordPtr(targetElement->Attribute("oep_va"));
	+ *addressIAT = ConvertStringToDwordPtr(targetElement->Attribute("iat_va"));
	+ *sizeIAT = (DWORD)ConvertStringToDwordPtr(targetElement->Attribute("iat_size"));

	- delete [] buffer;
	- }
	+ parseAllElementModules(targetElement, moduleList);

	return true;
	}

	-void TreeImportExport::setTargetInformation(TiXmlElement * rootElement, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT)
	+void TreeImportExport::setTargetInformation(TiXmlElement * rootElement, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT)
	{
	StringConversion::ToASCII(process->filename, xmlStringBuffer, _countof(xmlStringBuffer));
	rootElement->SetAttribute("filename", xmlStringBuffer);

	ConvertDwordPtrToString(addressOEP);
	rootElement->SetAttribute("oep_va", xmlStringBuffer);

	ConvertDwordPtrToString(addressIAT);
	rootElement->SetAttribute("iat_va", xmlStringBuffer);

	ConvertDwordPtrToString(sizeIAT);
	rootElement->SetAttribute("iat_size", xmlStringBuffer);
	}

	-char * TreeImportExport::readXmlFile(const WCHAR * xmlFilePath)
	+bool TreeImportExport::readXmlFile(TiXmlDocument& doc, const WCHAR * xmlFilePath)
	{
	- FILE * pFile = 0;
	- long lSize = 0;
	- char * buffer = 0;
	+ bool success = false;

	- if (_wfopen_s(&pFile, xmlFilePath, L"r") == NULL)
	+ FILE * pFile = 0;
	+ if (_wfopen_s(&pFile, xmlFilePath, L"rb") == 0)
	{
	- fseek(pFile, 0, SEEK_END);
	- lSize = ftell(pFile);
	- fseek(pFile, 0, SEEK_SET);
	-
	- if (lSize > 2)
	- {
	- buffer = new char[lSize + sizeof(char)];
	-
	- ZeroMemory(buffer, lSize + sizeof(char));
	-
	- fread(buffer, sizeof(char), lSize, pFile);
	-
	- if (!feof(pFile) \|\| ferror(pFile))
	- {
	- delete [] buffer;
	- buffer = 0;
	- }
	- }
	-
	+ success = doc.LoadFile(pFile);
	fclose (pFile);
	- return buffer;
	- }
	- else
	- {
	- return 0;
	}
	+
	+ return success;
	}

	-bool TreeImportExport::saveXmlToFile(TiXmlDocument doc, const WCHAR * xmlFilePath)
	+bool TreeImportExport::saveXmlToFile(const TiXmlDocument& doc, const WCHAR * xmlFilePath)
	{
	FILE * pFile = 0;
	-
	- if (_wfopen_s(&pFile, xmlFilePath, L"w") == NULL)
	+ if (_wfopen_s(&pFile, xmlFilePath, L"wb") == 0)
	{
	doc.Print(pFile);
	- fclose (pFile);
	+ fclose(pFile);
	return true;
	}
	else
	{
	return false;
	}
	}

	-void TreeImportExport::addModuleListToRootElement( TiXmlElement * rootElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList )
	+void TreeImportExport::addModuleListToRootElement(TiXmlElement * rootElement, const std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	- std::map<DWORD_PTR, ImportModuleThunk>::iterator mapIt;
	- std::map<DWORD_PTR, ImportThunk>::iterator mapIt2;
	- ImportModuleThunk * importModuleThunk = 0;
	- ImportThunk * importThunk = 0;
	-
	- TiXmlElement * moduleElement;
	- TiXmlElement * importElement;
	-
	- for ( mapIt = moduleList.begin() ; mapIt != moduleList.end(); mapIt++ )
	+ std::map<DWORD_PTR, ImportModuleThunk>::const_iterator it_mod;
	+ for(it_mod = moduleList.begin(); it_mod != moduleList.end(); it_mod++)
	{
	- importModuleThunk = &((*mapIt).second);
	+ const ImportModuleThunk& importModuleThunk = it_mod->second;

	- moduleElement = getModuleXmlElement(importModuleThunk);
	+ TiXmlElement* moduleElement = getModuleXmlElement(&importModuleThunk);

	- for ( mapIt2 = (mapIt).second.thunkList.begin() ; mapIt2 != (mapIt).second.thunkList.end(); mapIt2++ )
	+ std::map<DWORD_PTR, ImportThunk>::const_iterator it_thunk;
	+ for(it_thunk = importModuleThunk.thunkList.begin(); it_thunk != importModuleThunk.thunkList.end(); it_thunk++)
	{
	- importThunk = &((*mapIt2).second);
	+ const ImportThunk& importThunk = it_thunk->second;
	+
	+ TiXmlElement* importElement = getImportXmlElement(&importThunk);

	- importElement = getImportXmlElement(importThunk);
	moduleElement->LinkEndChild(importElement);
	}

	rootElement->LinkEndChild(moduleElement);
	}
	}

	TiXmlElement * TreeImportExport::getModuleXmlElement(const ImportModuleThunk * importModuleThunk)
	{
	TiXmlElement * moduleElement = new TiXmlElement("module");

	StringConversion::ToASCII(importModuleThunk->moduleName, xmlStringBuffer, _countof(xmlStringBuffer));
	moduleElement->SetAttribute("filename", xmlStringBuffer);

	ConvertDwordPtrToString(importModuleThunk->getFirstThunk());
	- moduleElement->SetAttribute("first_thunk_rva",xmlStringBuffer);
	+ moduleElement->SetAttribute("first_thunk_rva", xmlStringBuffer);

	return moduleElement;
	}

	TiXmlElement * TreeImportExport::getImportXmlElement(const ImportThunk * importThunk)
	{
	TiXmlElement * importElement = 0;

	if (importThunk->valid)
	{
	importElement = new TiXmlElement("import_valid");

	if(importThunk->name[0] != '\0')
	{
	- importElement->SetAttribute("name",importThunk->name);
	+ importElement->SetAttribute("name", importThunk->name);
	}

	ConvertWordToString(importThunk->ordinal);
	importElement->SetAttribute("ordinal",xmlStringBuffer);

	ConvertWordToString(importThunk->hint);
	importElement->SetAttribute("hint",xmlStringBuffer);

	ConvertBoolToString(importThunk->suspect);
	importElement->SetAttribute("suspect", xmlStringBuffer);
	}
	else
	{
	importElement = new TiXmlElement("import_invalid");
	}

	ConvertDwordPtrToString(importThunk->rva);
	importElement->SetAttribute("iat_rva", xmlStringBuffer);

	ConvertDwordPtrToString(importThunk->apiAddressVA);
	- importElement->SetAttribute("address_va",xmlStringBuffer);
	+ importElement->SetAttribute("address_va", xmlStringBuffer);

	return importElement;
	}

	void TreeImportExport::ConvertBoolToString(const bool boolValue)
	{
	if (boolValue)
	{
	strcpy_s(xmlStringBuffer, "1");
	}
	else
	{
	strcpy_s(xmlStringBuffer, "0");
	}
	}

	bool TreeImportExport::ConvertStringToBool(const char * strValue)
	{
	if (strValue)
	{
	if (strValue[0] == '1')
	{
	return true;
	}
	}

	return false;
	}

	void TreeImportExport::ConvertDwordPtrToString(const DWORD_PTR dwValue)
	{
	sprintf_s(xmlStringBuffer, PRINTF_DWORD_PTR_FULL_S, dwValue);
	}

	DWORD_PTR TreeImportExport::ConvertStringToDwordPtr(const char * strValue)
	{
	DWORD_PTR result = 0;

	if (strValue)
	{
	#ifdef _WIN64
	result = _strtoi64(strValue, NULL, 16);
	#else
	result = strtoul(strValue, NULL, 16);
	#endif
	}

	return result;
	}

	void TreeImportExport::ConvertWordToString(const WORD dwValue)
	{
	sprintf_s(xmlStringBuffer, "%04X", dwValue);
	}

	WORD TreeImportExport::ConvertStringToWord(const char * strValue)
	{
	WORD result = 0;

	if (strValue)
	{
	result = (WORD)strtoul(strValue, NULL, 16);
	}

	return result;
	}

	-void TreeImportExport::parseAllElementModules( TiXmlElement * targetElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList )
	+void TreeImportExport::parseAllElementModules(TiXmlElement * targetElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	- TiXmlElement * moduleElement = 0;
	ImportModuleThunk importModuleThunk;
	- const char * filename = 0;

	- for(moduleElement = targetElement->FirstChildElement(); moduleElement; moduleElement = moduleElement->NextSiblingElement())
	+ for(TiXmlElement * moduleElement = targetElement->FirstChildElement(); moduleElement; moduleElement = moduleElement->NextSiblingElement())
	{
	- filename = moduleElement->Attribute("filename");
	-
	+ const char * filename = moduleElement->Attribute("filename");
	if (filename)
	{
	StringConversion::ToUTF16(filename, importModuleThunk.moduleName, _countof(importModuleThunk.moduleName));

	importModuleThunk.firstThunk = ConvertStringToDwordPtr(moduleElement->Attribute("first_thunk_rva"));

	importModuleThunk.thunkList.clear();
	-
	parseAllElementImports(moduleElement, &importModuleThunk);

	- moduleList.insert(std::pair<DWORD_PTR,ImportModuleThunk>(importModuleThunk.firstThunk, importModuleThunk));
	-
	+ moduleList[importModuleThunk.firstThunk] = importModuleThunk;
	}
	}
	}

	-void TreeImportExport::parseAllElementImports( TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk )
	+void TreeImportExport::parseAllElementImports(TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk)
	{
	- TiXmlElement * importElement = 0;
	ImportThunk importThunk;
	- const char * temp = 0;

	- for(importElement = moduleElement->FirstChildElement(); importElement; importElement = importElement->NextSiblingElement())
	+ for(TiXmlElement * importElement = moduleElement->FirstChildElement(); importElement; importElement = importElement->NextSiblingElement())
	{
	- temp = importElement->Value();
	+ const char * temp = importElement->Value();

	if (!strcmp(temp, "import_valid"))
	{
	temp = importElement->Attribute("name");
	if (temp)
	{
	strcpy_s(importThunk.name, temp);
	}
	else
	{
	importThunk.name[0] = 0;
	}

	wcscpy_s(importThunk.moduleName, importModuleThunk->moduleName);

	importThunk.suspect = ConvertStringToBool(importElement->Attribute("suspect"));
	importThunk.ordinal = ConvertStringToWord(importElement->Attribute("ordinal"));
	importThunk.hint = ConvertStringToWord(importElement->Attribute("hint"));

	importThunk.valid = true;
	}
	else
	{
	importThunk.valid = false;
	importThunk.suspect = true;
	}

	importThunk.apiAddressVA = ConvertStringToDwordPtr(importElement->Attribute("address_va"));
	importThunk.rva = ConvertStringToDwordPtr(importElement->Attribute("iat_rva"));

	if (importThunk.rva != 0)
	{
	- importModuleThunk->thunkList.insert(std::pair<DWORD_PTR,ImportThunk>(importThunk.rva, importThunk));
	+ importModuleThunk->thunkList[importThunk.rva] = importThunk;
	}

	}
	}
	diff --git a/Scylla/TreeImportExport.h b/Scylla/TreeImportExport.h
	index 443807a..61362d1 100644
	--- a/Scylla/TreeImportExport.h
	+++ b/Scylla/TreeImportExport.h
	@@ -1,38 +1,42 @@
	#pragma once

	#include <windows.h>
	#include "ProcessLister.h"
	#include "Thunks.h"
	#include <tinyxml.h>

	class TreeImportExport
	{
	public:

	- bool exportTreeList(const WCHAR * targetXmlFile, std::map<DWORD_PTR, ImportModuleThunk> & moduleList, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT);
	- bool importTreeList(const WCHAR * targetXmlFile, std::map<DWORD_PTR, ImportModuleThunk> & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT);
	+ TreeImportExport(const WCHAR * targetXmlFile);
	+
	+ bool exportTreeList(const std::map<DWORD_PTR, ImportModuleThunk> & moduleList, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT) ;
	+ bool importTreeList(std::map<DWORD_PTR, ImportModuleThunk> & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT);

	private:

	- char xmlStringBuffer[100];
	+ WCHAR xmlPath[MAX_PATH];
	+
	+ char xmlStringBuffer[MAX_PATH];
	+
	+ void setTargetInformation(TiXmlElement * rootElement, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT);
	+ void addModuleListToRootElement(TiXmlElement * rootElement, const std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
	+
	+ void parseAllElementModules(TiXmlElement * targetElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
	+ void parseAllElementImports(TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk);

	- void addModuleListToRootElement( TiXmlElement * rootElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList );
	TiXmlElement * getModuleXmlElement(const ImportModuleThunk * importModuleThunk);
	TiXmlElement * getImportXmlElement(const ImportThunk * importThunk);

	- bool saveXmlToFile(TiXmlDocument doc, const WCHAR * xmlFilePath);
	- char * readXmlFile(const WCHAR * xmlFilePath);
	-
	- void setTargetInformation(TiXmlElement * rootElement, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT);
	+ bool saveXmlToFile(const TiXmlDocument& doc, const WCHAR * xmlFilePath);
	+ bool readXmlFile(TiXmlDocument& doc, const WCHAR * xmlFilePath);

	void ConvertBoolToString(const bool boolValue);
	void ConvertWordToString(const WORD dwValue);
	void ConvertDwordPtrToString(const DWORD_PTR dwValue);

	DWORD_PTR ConvertStringToDwordPtr(const char * strValue);
	WORD ConvertStringToWord(const char * strValue);
	bool ConvertStringToBool(const char * strValue);
	-
	- void parseAllElementModules( TiXmlElement * targetElement, std::map<DWORD_PTR, ImportModuleThunk> & moduleList );
	- void parseAllElementImports( TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk );
	};

File Metadata

Mime Type: text/x-diff
Expires: Sat, Sep 21, 6:45 AM (1 d, 3 h)
Storage Engine: local-disk
Storage Format: Raw Data
Storage Handle: a4/2c/049bc3c230ece7714d00eded9116

No OneTemporaryActions

View Options

File Metadata

Event Timeline

No OneTemporary
Actions