diff --git a/Scylla/ImportsHandling.cpp b/Scylla/ImportsHandling.cpp index 382dd64..f149aba 100644 --- a/Scylla/ImportsHandling.cpp +++ b/Scylla/ImportsHandling.cpp @@ -1,786 +1,807 @@ #include "ImportsHandling.h" #include "Thunks.h" #include "definitions.h" #include #include #include "multitree.h" // CMultiSelectTreeViewCtrl #include "resource.h" //#define DEBUG_COMMENTS void ImportThunk::invalidate() { ordinal = 0; hint = 0; valid = false; suspect = false; moduleName[0] = 0; name[0] = 0; } bool ImportModuleThunk::isValid() const { std::map::const_iterator iterator = thunkList.begin(); while (iterator != thunkList.end()) { if (iterator->second.valid == false) { return false; } iterator++; } return true; } DWORD_PTR ImportModuleThunk::getFirstThunk() const { if (thunkList.size() > 0) { const std::map::const_iterator iterator = thunkList.begin(); return iterator->first; } else { return 0; } } ImportsHandling::ImportsHandling(CMultiSelectTreeViewCtrl& TreeImports) : TreeImports(TreeImports) { hIconCheck.LoadIcon(IDI_ICON_CHECK, 16, 16); hIconWarning.LoadIcon(IDI_ICON_WARNING, 16, 16); hIconError.LoadIcon(IDI_ICON_ERROR, 16, 16); - CDCHandle dc = TreeImports.GetDC(); + CDCHandle dc = CWindow(::GetDesktopWindow()).GetDC(); int bits = dc.GetDeviceCaps(BITSPIXEL); const UINT FLAGS = bits > 16 ? ILC_COLOR32 : (ILC_COLOR24 | ILC_MASK); TreeIcons.Create(16, 16, FLAGS, 3, 1); TreeIcons.AddIcon(hIconCheck); TreeIcons.AddIcon(hIconWarning); TreeIcons.AddIcon(hIconError); m_thunkCount = m_invalidThunkCount = m_suspectThunkCount = 0; } ImportsHandling::~ImportsHandling() { TreeIcons.Destroy(); } bool ImportsHandling::isModule(CTreeItem item) { return (0 != getModuleThunk(item)); } bool ImportsHandling::isImport(CTreeItem item) { return (0 != getImportThunk(item)); } ImportModuleThunk * ImportsHandling::getModuleThunk(CTreeItem item) { stdext::hash_map::const_iterator it; it = itemData.find(item); if(it != itemData.end()) { const TreeItemData * data = &it->second; if(data->isModule) { return data->module; } } return NULL; } ImportThunk * ImportsHandling::getImportThunk(CTreeItem item) { stdext::hash_map::const_iterator it; TreeItemData * data = getItemData(item); if(data && !data->isModule) { return data->import; } return NULL; } -void ImportsHandling::setItemData(CTreeItem item, const TreeItemData& data) +void ImportsHandling::setItemData(CTreeItem item, const TreeItemData * data) { - itemData[item] = data; + itemData[item] = *data; } ImportsHandling::TreeItemData * ImportsHandling::getItemData(CTreeItem item) { stdext::hash_map::iterator it; it = itemData.find(item); if(it != itemData.end()) { return &it->second; } return NULL; } void ImportsHandling::updateCounts() { std::map::iterator it_module; std::map::iterator it_import; m_thunkCount = m_invalidThunkCount = m_suspectThunkCount = 0; it_module = moduleList.begin(); while (it_module != moduleList.end()) { ImportModuleThunk &moduleThunk = it_module->second; it_import = moduleThunk.thunkList.begin(); while (it_import != moduleThunk.thunkList.end()) { ImportThunk &importThunk = it_import->second; m_thunkCount++; if(!importThunk.valid) m_invalidThunkCount++; else if(importThunk.suspect) m_suspectThunkCount++; it_import++; } it_module++; } } /*bool ImportsHandling::addImport(const WCHAR * moduleName, const CHAR * name, DWORD_PTR va, DWORD_PTR rva, WORD ordinal, bool valid, bool suspect) { ImportThunk import; ImportModuleThunk * module = 0; std::map::iterator iterator1; if (moduleList.size() > 1) { iterator1 = moduleList.begin(); while (iterator1 != moduleList.end()) { if (rva >= iterator1->second.firstThunk) { iterator1++; if (iterator1 == moduleList.end()) { iterator1--; module = &(iterator1->second); break; } else if (rva < iterator1->second.firstThunk) { iterator1--; module = &(iterator1->second); break; } } } } else { iterator1 = moduleList.begin(); module = &(iterator1->second); } if (!module) { Logger::debugLog(TEXT("ImportsHandling::addFunction module not found rva ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(""),rva); return false; } //TODO import.suspect = true; import.valid = false; import.va = va; import.rva = rva; import.ordinal = ordinal; wcscpy_s(import.moduleName, MAX_PATH, moduleName); strcpy_s(import.name, MAX_PATH, name); module->thunkList.insert(std::pair(import.rva, import)); return true; } */ /* bool ImportsHandling::addModule(const WCHAR * moduleName, DWORD_PTR firstThunk) { ImportModuleThunk module; module.firstThunk = firstThunk; wcscpy_s(module.moduleName, MAX_PATH, moduleName); moduleList.insert(std::pair(firstThunk,module)); return true; } */ void ImportsHandling::displayAllImports() { std::map::iterator it_module; std::map::iterator it_import; TreeImports.DeleteAllItems(); itemData.clear(); TreeImports.SetImageList(TreeIcons); it_module = moduleList.begin(); while (it_module != moduleList.end()) { ImportModuleThunk &moduleThunk = it_module->second; + moduleThunk.key = moduleThunk.firstThunk; // This belongs elsewhere... moduleThunk.hTreeItem = addDllToTreeView(TreeImports, &moduleThunk); it_import = moduleThunk.thunkList.begin(); while (it_import != moduleThunk.thunkList.end()) { ImportThunk &importThunk = it_import->second; + importThunk.key = importThunk.rva; // This belongs elsewhere... importThunk.hTreeItem = addApiToTreeView(TreeImports, moduleThunk.hTreeItem, &importThunk); it_import++; } it_module++; } updateCounts(); } void ImportsHandling::clearAllImports() { TreeImports.DeleteAllItems(); itemData.clear(); moduleList.clear(); updateCounts(); } CTreeItem ImportsHandling::addDllToTreeView(CMultiSelectTreeViewCtrl& idTreeView, ImportModuleThunk * moduleThunk) { CTreeItem item = idTreeView.InsertItem(L"", NULL, TVI_ROOT); item.SetData(itemData.size()); TreeItemData data; data.isModule = true; data.module = moduleThunk; - setItemData(item, data); + setItemData(item, &data); updateModuleInTreeView(moduleThunk, item); return item; } CTreeItem ImportsHandling::addApiToTreeView(CMultiSelectTreeViewCtrl& idTreeView, CTreeItem parentDll, ImportThunk * importThunk) { CTreeItem item = idTreeView.InsertItem(L"", parentDll, TVI_LAST); item.SetData(itemData.size()); TreeItemData data; data.isModule = false; data.import = importThunk; - setItemData(item, data); + setItemData(item, &data); updateImportInTreeView(importThunk, item); return item; } void ImportsHandling::selectImports(bool invalid, bool suspect) { std::map::iterator it_module; std::map::iterator it_import; TreeImports.SelectAllItems(FALSE); //remove selection it_module = moduleList.begin(); while (it_module != moduleList.end()) { ImportModuleThunk &moduleThunk = it_module->second; it_import = moduleThunk.thunkList.begin(); while (it_import != moduleThunk.thunkList.end()) { ImportThunk &importThunk = it_import->second; if ((invalid && !importThunk.valid) || (suspect && importThunk.suspect)) { TreeImports.SelectItem(importThunk.hTreeItem, TRUE); importThunk.hTreeItem.EnsureVisible(); } it_import++; } it_module++; } } bool ImportsHandling::invalidateImport(CTreeItem item) { ImportThunk * import = getImportThunk(item); if(import) { CTreeItem parent = item.GetParent(); if(!parent.IsNull()) { const ImportModuleThunk * module = getModuleThunk(parent); if(module) { import->invalidate(); updateImportInTreeView(import, import->hTreeItem); updateModuleInTreeView(module, module->hTreeItem); updateCounts(); return true; } } } return false; } +bool ImportsHandling::invalidateModule(CTreeItem item) +{ + ImportModuleThunk * module = getModuleThunk(item); + if(module) + { + std::map::iterator it_import; + + it_import = module->thunkList.begin(); + while(it_import != module->thunkList.end()) + { + ImportThunk * import = &it_import->second; + import->invalidate(); + updateImportInTreeView(import, import->hTreeItem); + it_import++; + } + + updateModuleInTreeView(module, module->hTreeItem); + + updateCounts(); + return true; + } + return false; +} + bool ImportsHandling::setImport(CTreeItem item, const WCHAR * moduleName, const CHAR * apiName, WORD ordinal, WORD hint, bool valid, bool suspect) { ImportThunk * import = getImportThunk(item); if(import) { CTreeItem parent = item.GetParent(); if(!parent.IsNull()) { const ImportModuleThunk * module = getModuleThunk(parent); if(module) { wcscpy_s(import->moduleName, _countof(import->moduleName), moduleName); strcpy_s(import->name, _countof(import->name), apiName); import->ordinal = ordinal; //import->apiAddressVA = api->va; //?? import->hint = hint; import->valid = valid; import->suspect = suspect; updateImportInTreeView(import, item); updateModuleInTreeView(module, module->hTreeItem); updateCounts(); return true; } } } return false; } void ImportsHandling::updateImportInTreeView(const ImportThunk * importThunk, CTreeItem item) { if (importThunk->valid) { WCHAR tempString[300]; if (importThunk->name[0] != 0x00) { swprintf_s(tempString, _countof(tempString),TEXT("ord: %04X name: %S"),importThunk->ordinal,importThunk->name); } else { swprintf_s(tempString, _countof(tempString),TEXT("ord: %04X"),importThunk->ordinal); } swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(" rva: ")TEXT(PRINTF_DWORD_PTR_HALF)TEXT(" mod: %s %s"),importThunk->rva,importThunk->moduleName,tempString); } else { swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(" rva: ")TEXT(PRINTF_DWORD_PTR_HALF)TEXT(" ptr: ")TEXT(PRINTF_DWORD_PTR_FULL),importThunk->rva,importThunk->apiAddressVA); } item.SetText(stringBuffer); Icon icon = getAppropiateIcon(importThunk); item.SetImage(icon, icon); } void ImportsHandling::updateModuleInTreeView(const ImportModuleThunk * importThunk, CTreeItem item) { swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("%s (%d) FThunk: ")TEXT(PRINTF_DWORD_PTR_HALF),importThunk->moduleName,importThunk->thunkList.size(), importThunk->firstThunk); item.SetText(stringBuffer); Icon icon = getAppropiateIcon(importThunk->isValid()); item.SetImage(icon, icon); } ImportsHandling::Icon ImportsHandling::getAppropiateIcon(const ImportThunk * importThunk) { if(importThunk->valid) { if(importThunk->suspect) { return iconWarning; } else { return iconCheck; } } else { return iconError; } } ImportsHandling::Icon ImportsHandling::getAppropiateIcon(bool valid) { if(valid) { return iconCheck; } else { return iconError; } } bool ImportsHandling::cutImport(CTreeItem item) { ImportThunk * import = getImportThunk(item); if(import) { CTreeItem parent = item.GetParent(); if(!parent.IsNull()) { ImportModuleThunk * module = getModuleThunk(parent); if(module) { itemData.erase(item); import->hTreeItem.Delete(); - module->thunkList.erase(import->rva); + module->thunkList.erase(import->key); import = 0; if (module->thunkList.empty()) { itemData.erase(parent); module->hTreeItem.Delete(); - moduleList.erase(module->firstThunk); + moduleList.erase(module->key); module = 0; } else { if (module->isValid() && module->moduleName[0] == L'?') { //update module name wcscpy_s(module->moduleName,_countof(module->moduleName),(*module->thunkList.begin()).second.moduleName); } module->firstThunk = (*module->thunkList.begin()).second.rva; - updateModuleInTreeView(module, module->hTreeItem); } updateCounts(); return true; } } } return false; } bool ImportsHandling::cutModule(CTreeItem item) { ImportModuleThunk * module = getModuleThunk(item); if(module) { CTreeItem child = item.GetChild(); while(!child.IsNull()) { itemData.erase(child); child = child.GetNextSibling(); } itemData.erase(item); module->hTreeItem.Delete(); - moduleList.erase(module->firstThunk); + moduleList.erase(module->key); module = 0; updateCounts(); return true; } return false; } DWORD_PTR ImportsHandling::getApiAddressByNode(CTreeItem item) { const ImportThunk * import = getImportThunk(item); if(import) { return import->apiAddressVA; } return 0; } void ImportsHandling::scanAndFixModuleList() { - std::map::iterator iterator1; - std::map::iterator iterator2; - ImportModuleThunk * moduleThunk; - ImportThunk * importThunk; - - - iterator1 = moduleList.begin(); + std::map::iterator it_module; + std::map::iterator it_import; - while (iterator1 != moduleList.end()) + it_module = moduleList.begin(); + while (it_module != moduleList.end()) { - moduleThunk = &(iterator1->second); - - iterator2 = moduleThunk->thunkList.begin(); + ImportModuleThunk &moduleThunk = it_module->second; - while (iterator2 != moduleThunk->thunkList.end()) + it_import = moduleThunk.thunkList.begin(); + while (it_import != moduleThunk.thunkList.end()) { - importThunk = &(iterator2->second); + ImportThunk &importThunk = it_import->second; - if (importThunk->moduleName[0] == 0 || importThunk->moduleName[0] == L'?') + if (importThunk.moduleName[0] == 0 || importThunk.moduleName[0] == L'?') { - addNotFoundApiToModuleList(importThunk); + addNotFoundApiToModuleList(&importThunk); } else { - if (isNewModule(importThunk->moduleName)) + if (isNewModule(importThunk.moduleName)) { - addModuleToModuleList(importThunk->moduleName, importThunk->rva); + addModuleToModuleList(importThunk.moduleName, importThunk.rva); } - addFunctionToModuleList(importThunk); + addFunctionToModuleList(&importThunk); } - iterator2++; + it_import++; } - moduleThunk->thunkList.clear(); + moduleThunk.thunkList.clear(); - iterator1++; + it_module++; } moduleList = moduleListNew; moduleListNew.clear(); } bool ImportsHandling::findNewModules(std::map & thunkList) { throw std::exception("The method or operation is not implemented."); } bool ImportsHandling::addModuleToModuleList(const WCHAR * moduleName, DWORD_PTR firstThunk) { ImportModuleThunk module; module.firstThunk = firstThunk; wcscpy_s(module.moduleName, _countof(module.moduleName), moduleName); - moduleListNew.insert(std::pair(firstThunk,module)); - + module.key = module.firstThunk; + moduleListNew[module.key] = module; return true; } bool ImportsHandling::isNewModule(const WCHAR * moduleName) { std::map::iterator it_module; it_module = moduleListNew.begin(); while (it_module != moduleListNew.end()) { if (!_wcsicmp(it_module->second.moduleName, moduleName)) { return false; } it_module++; } return true; } void ImportsHandling::addUnknownModuleToModuleList(DWORD_PTR firstThunk) { ImportModuleThunk module; module.firstThunk = firstThunk; wcscpy_s(module.moduleName, _countof(module.moduleName), L"?"); - moduleListNew.insert(std::pair(firstThunk,module)); + module.key = module.firstThunk; + moduleListNew[module.key] = module; } bool ImportsHandling::addNotFoundApiToModuleList(const ImportThunk * apiNotFound) { ImportThunk import; ImportModuleThunk * module = 0; std::map::iterator it_module; DWORD_PTR rva = apiNotFound->rva; if (moduleListNew.size() > 0) { it_module = moduleListNew.begin(); while (it_module != moduleListNew.end()) { if (rva >= it_module->second.firstThunk) { it_module++; if (it_module == moduleListNew.end()) { it_module--; //new unknown module if (it_module->second.moduleName[0] == L'?') { module = &(it_module->second); } else { addUnknownModuleToModuleList(apiNotFound->rva); module = &(moduleListNew.find(rva)->second); } break; } else if (rva < it_module->second.firstThunk) { it_module--; module = &(it_module->second); break; } } else { #ifdef DEBUG_COMMENTS Logger::debugLog("Error iterator1 != (*moduleThunkList).end()\r\n"); #endif break; } } } else { //new unknown module addUnknownModuleToModuleList(apiNotFound->rva); module = &(moduleListNew.find(rva)->second); } if (!module) { #ifdef DEBUG_COMMENTS Logger::debugLog(TEXT("ImportsHandling::addFunction module not found rva ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("\r\n"),rva); #endif return false; } import.suspect = true; import.valid = false; import.va = apiNotFound->va; import.rva = apiNotFound->rva; import.apiAddressVA = apiNotFound->apiAddressVA; import.ordinal = 0; wcscpy_s(import.moduleName, _countof(import.moduleName), L"?"); strcpy_s(import.name, _countof(import.name), "?"); - module->thunkList.insert(std::pair(import.rva, import)); - + import.key = import.rva; + module->thunkList[import.key] = import; return true; } bool ImportsHandling::addFunctionToModuleList(const ImportThunk * apiFound) { ImportThunk import; ImportModuleThunk * module = 0; std::map::iterator it_module; if (moduleListNew.size() > 1) { it_module = moduleListNew.begin(); while (it_module != moduleListNew.end()) { if (apiFound->rva >= it_module->second.firstThunk) { it_module++; if (it_module == moduleListNew.end()) { it_module--; module = &(it_module->second); break; } else if (apiFound->rva < it_module->second.firstThunk) { it_module--; module = &(it_module->second); break; } } else { #ifdef DEBUG_COMMENTS Logger::debugLog(TEXT("Error iterator1 != moduleListNew.end()\r\n")); #endif break; } } } else { it_module = moduleListNew.begin(); module = &(it_module->second); } if (!module) { #ifdef DEBUG_COMMENTS Logger::debugLog(TEXT("ImportsHandling::addFunction module not found rva ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("\r\n"),apiFound->rva); #endif return false; } import.suspect = apiFound->suspect; import.valid = apiFound->valid; import.va = apiFound->va; import.rva = apiFound->rva; import.apiAddressVA = apiFound->apiAddressVA; import.ordinal = apiFound->ordinal; import.hint = apiFound->hint; wcscpy_s(import.moduleName, _countof(import.moduleName), apiFound->moduleName); strcpy_s(import.name, _countof(import.name), apiFound->name); - module->thunkList.insert(std::pair(import.rva, import)); - + import.key = import.rva; + module->thunkList[import.key] = import; return true; } void ImportsHandling::expandAllTreeNodes() { changeExpandStateOfTreeNodes(TVE_EXPAND); } void ImportsHandling::collapseAllTreeNodes() { changeExpandStateOfTreeNodes(TVE_COLLAPSE); } void ImportsHandling::changeExpandStateOfTreeNodes(UINT flag) { std::map::iterator it_module; it_module = moduleList.begin(); while (it_module != moduleList.end()) { ImportModuleThunk &moduleThunk = it_module->second; moduleThunk.hTreeItem.Expand(flag); it_module++; } } diff --git a/Scylla/ImportsHandling.h b/Scylla/ImportsHandling.h index f02fa26..ef4444d 100644 --- a/Scylla/ImportsHandling.h +++ b/Scylla/ImportsHandling.h @@ -1,111 +1,112 @@ #pragma once #include #include // WTL #include #include #include // CTreeItem class CMultiSelectTreeViewCtrl; class ImportThunk; class ImportModuleThunk; class ImportsHandling { public: std::map moduleList; std::map moduleListNew; ImportsHandling(CMultiSelectTreeViewCtrl& TreeImports); ~ImportsHandling(); unsigned int thunkCount() const { return m_thunkCount; } unsigned int invalidThunkCount() const { return m_invalidThunkCount; } unsigned int suspectThunkCount() const { return m_suspectThunkCount; } bool isModule(CTreeItem item); bool isImport(CTreeItem item); ImportModuleThunk * getModuleThunk(CTreeItem item); ImportThunk * getImportThunk(CTreeItem item); void displayAllImports(); void clearAllImports(); void selectImports(bool invalid, bool suspect); bool invalidateImport(CTreeItem item); + bool invalidateModule(CTreeItem item); bool setImport(CTreeItem item, const WCHAR * moduleName, const CHAR * apiName, WORD ordinal = 0, WORD hint = 0, bool valid = true, bool suspect = false); bool cutImport(CTreeItem item); bool cutModule(CTreeItem item); //bool addImport(const WCHAR * moduleName, const CHAR * name, DWORD_PTR va, DWORD_PTR rva, WORD ordinal = 0, bool valid = true, bool suspect = false); //bool addModule(const WCHAR * moduleName, DWORD_PTR firstThunk); DWORD_PTR getApiAddressByNode(CTreeItem selectedTreeNode); void scanAndFixModuleList(); void expandAllTreeNodes(); void collapseAllTreeNodes(); private: DWORD numberOfFunctions; unsigned int m_thunkCount; unsigned int m_invalidThunkCount; unsigned int m_suspectThunkCount; struct TreeItemData { bool isModule; union { ImportModuleThunk * module; ImportThunk * import; }; }; stdext::hash_map itemData; - void setItemData(CTreeItem item, const TreeItemData& data); + void setItemData(CTreeItem item, const TreeItemData * data); TreeItemData * getItemData(CTreeItem item); WCHAR stringBuffer[600]; CMultiSelectTreeViewCtrl& TreeImports; CImageList TreeIcons; CIcon hIconCheck; CIcon hIconWarning; CIcon hIconError; // They have to be added to the image list in that order! enum Icon { iconCheck = 0, iconWarning, iconError }; void updateCounts(); CTreeItem addDllToTreeView(CMultiSelectTreeViewCtrl& idTreeView, ImportModuleThunk * moduleThunk); CTreeItem addApiToTreeView(CMultiSelectTreeViewCtrl& idTreeView, CTreeItem parentDll, ImportThunk * importThunk); void updateImportInTreeView(const ImportThunk * importThunk, CTreeItem item); void updateModuleInTreeView(const ImportModuleThunk * importThunk, CTreeItem item); //bool isItemSelected(CTreeItem hItem); //void unselectItem(CTreeItem htItem); //bool selectItem(CTreeItem hItem, bool select = true); bool findNewModules(std::map & thunkList); Icon getAppropiateIcon(const ImportThunk * importThunk); Icon getAppropiateIcon(bool valid); bool addModuleToModuleList(const WCHAR * moduleName, DWORD_PTR firstThunk); void addUnknownModuleToModuleList(DWORD_PTR firstThunk); bool addNotFoundApiToModuleList(const ImportThunk * apiNotFound); bool addFunctionToModuleList(const ImportThunk * apiFound); bool isNewModule(const WCHAR * moduleName); void changeExpandStateOfTreeNodes(UINT flag); }; diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp index 8fa7c0c..41fe979 100644 --- a/Scylla/MainGui.cpp +++ b/Scylla/MainGui.cpp @@ -1,1243 +1,1266 @@ #include "MainGui.h" -#include // 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" +#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); 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.GetFocusItem(); - if(!selected.IsNull() && !selected.GetParent().IsNull()) + if(!selected.IsNull() && importsHandling.isImport(selected)) { pickApiActionHandler(selected); } } return 1; case VK_DELETE: 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) { invalidateSelectedImportsActionHandler(); } void MainGui::OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { deleteSelectedImportsActionHandler(); } void MainGui::OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl) { - // TODO + saveTreeActionHandler(); } void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl) { - // TODO + 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::fillStatusBar() { // 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, _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(); processAccessHelp.targetImageBase = processAccessHelp.selectedModule->modBaseAddr; 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) { 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); } } 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& 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); + EditOEPAddress.SetValue(process.entryPoint + process.imageBase); selectedProcess = &process; enableDialogControls(TRUE); fillStatusBar(); } void MainGui::fillProcessListComboBox(CComboBox& hCombo) { hCombo.ResetContent(); std::vector& 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; + size_t bufsize = 0; for(int i = 0; i < ListLog.GetCount(); i++) { - int size = ListLog.GetTextLen(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, size * sizeof(WCHAR), buffer); + 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); + } + fillStatusBar(); +} + +void MainGui::invalidateSelectedImportsActionHandler() +{ + CTreeItem selected = TreeImports.GetFirstSelectedItem(); + while(!selected.IsNull()) + { + if(importsHandling.isImport(selected)) + { + importsHandling.invalidateImport(selected); + } + selected = TreeImports.GetNextSelectedItem(selected); + } + fillStatusBar(); +} + +void MainGui::loadTreeActionHandler() +{ + if(!selectedProcess) + return; + + WCHAR selectedFilePath[MAX_PATH]; + TreeImportExport treeIO; + DWORD_PTR addrOEP = 0; + DWORD_PTR addrIAT = 0; + DWORD sizeIAT = 0; + + if(showFileDialog(selectedFilePath, false, NULL, filterXml)) + { + if(!treeIO.importTreeList(selectedFilePath, importsHandling.moduleList, &addrOEP, &addrIAT, &sizeIAT)) + { + Logger::printfDialog(TEXT("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(); + fillStatusBar(); + + Logger::printfDialog(TEXT("Loaded tree file %s"), selectedFilePath); + Logger::printfDialog(TEXT("-> OEP: ")TEXT(PRINTF_DWORD_PTR_FULL), addrOEP); + Logger::printfDialog(TEXT("-> IAT: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: ")TEXT(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; + + if(showFileDialog(selectedFilePath, true, NULL, filterXml, L"xml")) + { + addrOEP = EditOEPAddress.GetValue(); + addrIAT = EditIATAddress.GetValue(); + sizeIAT = EditIATSize.GetValue(); + + if(!treeIO.exportTreeList(selectedFilePath, importsHandling.moduleList, selectedProcess, addrOEP, addrIAT, sizeIAT)) + { + Logger::printfDialog(TEXT("Saving tree file failed %s"), selectedFilePath); + MessageBox(L"Saving tree file failed.", L"Failure", MB_ICONERROR); + } + else + { + Logger::printfDialog(TEXT("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.GetWindowText(stringBuffer, _countof(stringBuffer)) > 1) + if(EditOEPAddress.GetWindowTextLength() > 0) { - searchAddress = stringToDwordPtr(stringBuffer); + searchAddress = EditOEPAddress.GetValue(); 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); + EditIATAddress.SetValue(addressIAT); + EditIATSize.SetValue(sizeIAT); swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("IAT found:\r\n\r\nStart: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("\r\nSize: 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) + if (EditIATAddress.GetWindowTextLength() > 0) { - addressIAT = stringToDwordPtr(stringBuffer); + addressIAT = EditIATAddress.GetValue(); } - if (EditIATSize.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0) + if (EditIATSize.GetWindowTextLength() > 0) { - sizeIAT = wcstoul(stringBuffer, NULL, 16); + sizeIAT = EditIATSize.GetValue(); } if (addressIAT && sizeIAT) { apiReader.readAndParseIAT(addressIAT, sizeIAT,importsHandling.moduleList); importsHandling.displayAllImports(); } fillStatusBar(); } -DWORD_PTR MainGui::stringToDwordPtr(const WCHAR * hexString) +void MainGui::SetupImportsMenuItems(CTreeItem item) { - DWORD_PTR address = 0; - -#ifdef _WIN64 - address = _wcstoui64(hexString, NULL, 16); -#else - address = wcstoul(hexString, NULL, 16); -#endif - - if (address == 0) + bool isItem, isImport = false; + isItem = !item.IsNull(); + if(isItem) { -#ifdef DEBUG_COMMENTS - Logger::debugLog(L"stringToDwordPtr :: address == 0, %s",hexString); -#endif - return 0; + isImport = importsHandling.isImport(item); } - 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; + UINT importOnly = isImport ? MF_ENABLED : MF_GRAYED; - hSub.EnableMenuItem(ID__INVALIDATEFUNCTION, thunkOnly); - hSub.EnableMenuItem(ID__DISASSEMBLE, thunkOnly); - hSub.EnableMenuItem(ID__CUTTHUNK, thunkOnly); + 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 CPoint client = pt; TreeImports.ScreenToClient(&client); UINT flags; over = TreeImports.HitTest(client, &flags); if(over && !(flags & TVHT_ONITEM)) { over = NULL; } } if (hMenuImports) { - // Prepare hmenuImports - SetupImportsMenuItems(!over.IsNull(), importsHandling.isImport(over)); + 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)(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.invalidateImport(over); + 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; } } } 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 &scyllaPluginList = PluginLoader::getScyllaPluginList(); std::vector &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(); + peDump.entryPoint = EditOEPAddress.GetValue(); wcscpy_s(peDump.fullpath, _countof(peDump.fullpath), processAccessHelp.selectedModule->fullPath); } else { peDump.imageBase = ProcessAccessHelp::targetImageBase; peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage; //get it from gui - peDump.entryPoint = getOEPFromGui(); + peDump.entryPoint = EditOEPAddress.GetValue(); wcscpy_s(peDump.fullpath, _countof(peDump.fullpath), 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,_countof(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, _countof(newFilePath), L"_SCY"); if(extension) { wcscat_s(newFilePath, _countof(newFilePath), 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) { 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_SAVETREE).EnableWindow(FALSE); - GetDlgItem(IDC_BTN_LOADTREE).EnableWindow(FALSE); GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE); - - 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() { importsHandling.clearAllImports(); fillStatusBar(); } void MainGui::pluginActionHandler( int menuItem ) { if(!selectedProcess) return; DllInjectionPlugin dllInjectionPlugin; std::vector &scyllaPluginList = PluginLoader::getScyllaPluginList(); std::vector &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 49a40bc..4d54d00 100644 --- a/Scylla/MainGui.h +++ b/Scylla/MainGui.h @@ -1,284 +1,283 @@ #pragma once #include #include "resource.h" // WTL #include // base ATL classes #include // base WTL classes #include // ATL GUI classes #include // WTL window frame helpers #include // WTL utility classes #include // WTL enhanced msg map macros #include // WTL controls #include // WTL dialog data exchange #include "multitree.h" +#include "hexedit.h" //#define _CRTDBG_MAP_ALLOC //#include //#include #include #include "Logger.h" #include "ProcessLister.h" #include "IATSearch.h" #include "PickDllGui.h" #include "ImportsHandling.h" class MainGui : public CDialogImpl, public CWinDataExchange, public CDialogResize, public CMessageFilter { public: enum { IDD = IDD_DLG_MAIN }; + // DDX_CONTROL : subclass + // DDX_CONTROL_HANDLE : attach BEGIN_DDX_MAP(MainGui) - DDX_CONTROL(IDC_TREE_IMPORTS, TreeImportsSubclass) // subclass + DDX_CONTROL(IDC_TREE_IMPORTS, TreeImportsSubclass) // needed for message reflection DDX_CONTROL(IDC_TREE_IMPORTS, TreeImports) - DDX_CONTROL_HANDLE(IDC_CBO_PROCESSLIST, ComboProcessList) // attach + DDX_CONTROL_HANDLE(IDC_CBO_PROCESSLIST, ComboProcessList) 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) + DDX_CONTROL(IDC_EDIT_OEPADDRESS, EditOEPAddress) + DDX_CONTROL(IDC_EDIT_IATADDRESS, EditIATAddress) + DDX_CONTROL(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) 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[]; + static const WCHAR filterXml[]; // Controls CMultiSelectTreeViewCtrl TreeImports; CComboBox ComboProcessList; - CEdit EditOEPAddress; - CEdit EditIATAddress; - CEdit EditIATSize; + CHexEdit EditOEPAddress; + CHexEdit EditIATAddress; + CHexEdit 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 loadTreeActionHandler(); + void saveTreeActionHandler(); 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 SetupImportsMenuItems(CTreeItem item); 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/MainGui.rc b/Scylla/MainGui.rc index 93fc449..9b23364 100644 Binary files a/Scylla/MainGui.rc and b/Scylla/MainGui.rc differ diff --git a/Scylla/Thunks.h b/Scylla/Thunks.h index 93cf06b..9cb935b 100644 --- a/Scylla/Thunks.h +++ b/Scylla/Thunks.h @@ -1,42 +1,44 @@ #pragma once #include #include #include // WTL #include #include #include //CTreeItem class ImportThunk { public: WCHAR moduleName[MAX_PATH]; char name[MAX_PATH]; DWORD_PTR va; DWORD_PTR rva; WORD ordinal; DWORD_PTR apiAddressVA; WORD hint; bool valid; bool suspect; CTreeItem hTreeItem; + DWORD_PTR key; void invalidate(); }; class ImportModuleThunk { public: WCHAR moduleName[MAX_PATH]; std::map thunkList; DWORD_PTR firstThunk; CTreeItem hTreeItem; + DWORD_PTR key; DWORD_PTR getFirstThunk() const; bool isValid() const; }; diff --git a/Scylla/hexedit.h b/Scylla/hexedit.h new file mode 100644 index 0000000..0e68680 --- /dev/null +++ b/Scylla/hexedit.h @@ -0,0 +1,235 @@ +#pragma once + +#ifndef __cplusplus + #error WTL requires C++ compilation (use a .cpp suffix) +#endif + +#ifndef __ATLMISC_H__ + #error hexedit.h requires atlmisc.h to be included first +#endif + +#ifndef __ATLCTRLS_H__ + #error hexedit.h requires atlmisc.h to be included first requires atlctrls.h to be included first +#endif + +/* +#ifdef _WIN64 + address = _wcstoui64(hexString, NULL, 16); +#else + address = wcstoul(hexString, NULL, 16); +#endif +*/ + +template< class T, typename NUM_T, class TBase = CEdit, class TWinTraits = CControlWinTraits > +class ATL_NO_VTABLE CHexEditImpl : public CWindowImpl< T, TBase, TWinTraits > +{ +public: + + static const short int BASE = 16; + static const size_t DIGITS = sizeof(NUM_T) * 2; // 2 digits / byte + static const size_t STRSIZE = DIGITS + 1; + + static const TCHAR Digits[]; + static const TCHAR OldProcProp[]; + + // 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; + } + + NUM_T GetValue() const + { + ATLASSERT(::IsWindow(m_hWnd)); + + TCHAR String[STRSIZE] = { 0 }; + GetWindowText(String, _countof(String)); + return _StringToNum(String); + } + + void SetValue(NUM_T Num, bool Fill = true) + { + ATLASSERT(::IsWindow(m_hWnd)); + + TCHAR String[STRSIZE] = { 0 }; + _NumToString(Num, String, Fill); + SetWindowText(String); + } + + // Implementation + + void _Init() + { + ATLASSERT(::IsWindow(m_hWnd)); + + LimitText(DIGITS); + } + + bool _IsValidChar(TCHAR Char) const + { + return ((NUM_T)-1 != _CharToNum(Char)); + } + + bool _IsValidString(const TCHAR String[STRSIZE]) const + { + for(int i = 0; String[i]; i++) + { + if(!_IsValidChar(String[i])) + return false; + } + return true; + } + + NUM_T _CharToNum(TCHAR Char) const + { + Char = _totupper(Char); + + for(int i = 0; Digits[i]; i++) + { + if(Char == Digits[i]) + return i; + } + + return -1; + } + + NUM_T _StringToNum(const TCHAR String[STRSIZE]) const + { + NUM_T CharNum; + NUM_T Num = 0; + + for(int i = 0; String[i]; i++) + { + CharNum = _CharToNum(String[i]); + if(CharNum == (NUM_T)-1) + break; + + Num *= BASE; + Num += CharNum; + } + + return Num; + } + + TCHAR _NumToChar(NUM_T Num) const + { + return Digits[Num % BASE]; + } + + void _NumToString(NUM_T Num, TCHAR String[STRSIZE], bool Fill) const + { + NUM_T Nums[DIGITS]; + int i, j; + + for(i = DIGITS-1; i >= 0; i--) + { + Nums[i] = Num % BASE; + Num /= BASE; + } + for(i = j = 0; i < DIGITS; i++) + { + // Only copy num if : non-null OR Fill OR non-null encountered before OR last num + if(Nums[i] || Fill || j || i == DIGITS-1) + { + String[j++] = _NumToChar(Nums[i]); + } + } + String[j] = '\0'; + } + + bool _GetClipboardText(TCHAR String[STRSIZE]) + { + #ifdef UNICODE + const UINT Format = CF_UNICODETEXT; + #else + const UINT Format = CF_TEXT; + #endif + + bool RetVal = false; + + if(IsClipboardFormatAvailable(Format) && OpenClipboard()) + { + HANDLE HMem = GetClipboardData(Format); + if(HMem) + { + _tcsncpy_s(String, STRSIZE, (TCHAR *)GlobalLock(HMem), STRSIZE-1); + String[STRSIZE-1] = '\0'; + GlobalUnlock(HMem); + RetVal = true; + } + CloseClipboard(); + } + + return RetVal; + } + + // Message map and handlers + + BEGIN_MSG_MAP_EX(CHexEditImpl) + MESSAGE_HANDLER_EX(WM_CREATE, OnCreate) + MSG_WM_SETTEXT(OnSetText) + MSG_WM_PASTE(OnPaste) + MSG_WM_CHAR(OnChar) + END_MSG_MAP() + + LRESULT OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam) + { + LRESULT lRes = DefWindowProc(); + _Init(); + return lRes; + } + + int OnSetText(LPCTSTR lpstrText) + { + bool PassThrough = (_tcslen(lpstrText) <= DIGITS) && _IsValidString(lpstrText); + if(!PassThrough) + { + MessageBeep(-1); + return FALSE; + } + + SetMsgHandled(FALSE); + return TRUE; + } + + void OnPaste() + { + TCHAR String[STRSIZE]; + bool PassThrough = !_GetClipboardText(String) || _IsValidString(String); + if(!PassThrough) + { + MessageBeep(-1); + return; + } + + SetMsgHandled(FALSE); + return; + } + + void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags) + { + // ignore all printable chars (incl. space) which are not valid digits + bool PassThrough = !_istprint((TCHAR)nChar) || _IsValidChar((TCHAR)nChar); + if(!PassThrough) + { + MessageBeep(-1); + return; + } + + SetMsgHandled(FALSE); + return; + } +}; + +template< class T, typename NUM_T, class TBase, class TWinTraits > const TCHAR CHexEditImpl< T, NUM_T, TBase, TWinTraits >::Digits[] = _T("0123456789ABCDEF"); + +template class CHexEdit : public CHexEditImpl, NUM_T> +{ +public: + DECLARE_WND_CLASS(_T("WTL_HexEdit")) +}; diff --git a/Scylla/main.cpp b/Scylla/main.cpp index d7d4073..18fe331 100644 --- a/Scylla/main.cpp +++ b/Scylla/main.cpp @@ -1,38 +1,40 @@ +//#include // Visual Leak Detector + #include // base ATL classes #include // base WTL classes CAppModule _Module; #include "MainGui.h" MainGui* pMainGui = NULL; // for Logger int APIENTRY _tWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPTSTR lpCmdLine, int nCmdShow) { CoInitialize(NULL); AtlInitCommonControls(ICC_LISTVIEW_CLASSES | ICC_TREEVIEW_CLASSES); HRESULT hRes = _Module.Init(NULL, hInstance); ATLASSERT(SUCCEEDED(hRes)); int nRet = 0; // BLOCK: Run application { MainGui dlgMain; pMainGui = &dlgMain; // o_O CMessageLoop loop; _Module.AddMessageLoop(&loop); dlgMain.Create(GetDesktopWindow()); dlgMain.ShowWindow(SW_SHOW); loop.Run(); } _Module.Term(); CoUninitialize(); return nRet; } diff --git a/Scylla/multitree.h b/Scylla/multitree.h index 7213139..4b793a0 100644 --- a/Scylla/multitree.h +++ b/Scylla/multitree.h @@ -1,600 +1,608 @@ #pragma once ///////////////////////////////////////////////////////////////////////////// // MultiSelectTree - A tree control with multi-select capabilities // // Written by Bjarke Viksoe (bjarke@viksoe.dk) // 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==(const CTreeItem& ti1, const CTreeItem& ti2) { 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 m_aData; + CSimpleMap 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; } BOOL SelectItem(HTREEITEM hItem, BOOL bSelect) { ATLASSERT(::IsWindow(m_hWnd)); _SelectItem(hItem, bSelect == TRUE); if( bSelect ) TBase::SelectItem(hItem); return TRUE; } 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)); + int iIndex = m_aData.FindKey(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(_MakeItem(hItem)); + int iIndex = m_aData.FindKey(hItem); if( iIndex >= 0 ) { nRes &= ~TVIS_SELECTED; if( m_aData.GetValueAt(iIndex) ) nRes |= TVIS_SELECTED; } } return nRes; } CTreeItem GetFirstSelectedItem() const { + HTREEITEM item = NULL; if( m_aData.GetSize() > 0 ) { for( int i = 0; i < m_aData.GetSize(); i++ ) { if( m_aData.GetValueAt(i) ) - return m_aData.GetKeyAt(i); + { + item = m_aData.GetKeyAt(i); + break; + } } } - return _MakeItem(NULL); + return _MakeItem(item); } CTreeItem GetNextSelectedItem(HTREEITEM hItem) const { - int iIndex = m_aData.FindKey(_MakeItem(hItem)); + HTREEITEM item = NULL; + int iIndex = m_aData.FindKey(hItem); if( iIndex >= 0 ) { for( int i = iIndex + 1; i < m_aData.GetSize(); i++ ) { if( m_aData.GetValueAt(i) ) - return m_aData.GetKeyAt(i); + { + item = m_aData.GetKeyAt(i); + break; + } } } - return _MakeItem(NULL); + return _MakeItem(item); } 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); + HTREEITEM hItem = 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.hItem = hItem; + nmtv.itemNew.lParam = GetItemData(hItem); 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); + m_aData.SetAtIndex(iIndex, hItem, bSelect); // Repaint item CRect rcItem; - if( GetItemRect(cItem, &rcItem, FALSE) ) + if( GetItemRect(hItem, &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(_MakeItem(hItem)), bSelect, action); + _SelectItem(m_aData.FindKey(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) { HTREEITEM hItem = GetFirstVisibleItem(); while( hItem != NULL ) { - int i = m_aData.FindKey(_MakeItem(hItem)); + int i = m_aData.FindKey(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); + _SelectItem(hItem, IsItemSelected(hItem) == TRUE/*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 if 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(FALSE); //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(_MakeItem(hItem)); + int iIndex = m_aData.FindKey(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; // 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(_MakeItem(hItem), false); + m_aData.Add(hItem, false); return (LRESULT) hItem; } LRESULT OnDeleteItem(NMHDR* pnmh) { const NMTREEVIEW* lpNMTV = (NMTREEVIEW*) pnmh; - m_aData.Remove(_MakeItem(lpNMTV->itemNew.hItem)); + m_aData.Remove(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(_MakeItem(hItem)); + int iIndex = m_aData.FindKey(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 > { public: DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName()) }; diff --git a/Scylla/resource.h b/Scylla/resource.h index 6819568..a61d5eb 100644 Binary files a/Scylla/resource.h and b/Scylla/resource.h differ