| diff --git a/Scylla/ImportsHandling.cpp b/Scylla/ImportsHandling.cpp | |
| index fde1745..382dd64 100644 | |
| --- a/Scylla/ImportsHandling.cpp | |
| +++ b/Scylla/ImportsHandling.cpp | |
| @@ -1,779 +1,786 @@ | |
| #include "ImportsHandling.h" | |
| #include "Thunks.h" | |
| #include "definitions.h" | |
| #include <atlmisc.h> | |
| #include <atlcrack.h> | |
| #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<DWORD_PTR, ImportThunk>::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<DWORD_PTR, ImportThunk>::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); | |
| - TreeIcons.Create(16, 16, ILC_COLOR32, 3, 1); | |
| + CDCHandle dc = TreeImports.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<HTREEITEM, TreeItemData>::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<HTREEITEM, TreeItemData>::const_iterator it; | |
| TreeItemData * data = getItemData(item); | |
| if(data && !data->isModule) | |
| { | |
| return data->import; | |
| } | |
| return NULL; | |
| } | |
| void ImportsHandling::setItemData(CTreeItem item, const TreeItemData& data) | |
| { | |
| itemData[item] = data; | |
| } | |
| ImportsHandling::TreeItemData * ImportsHandling::getItemData(CTreeItem item) | |
| { | |
| stdext::hash_map<HTREEITEM, TreeItemData>::iterator it; | |
| it = itemData.find(item); | |
| if(it != itemData.end()) | |
| { | |
| return &it->second; | |
| } | |
| return NULL; | |
| } | |
| void ImportsHandling::updateCounts() | |
| { | |
| std::map<DWORD_PTR, ImportModuleThunk>::iterator it_module; | |
| std::map<DWORD_PTR, ImportThunk>::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<DWORD_PTR, ImportModuleThunk>::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<DWORD_PTR,ImportThunk>(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<DWORD_PTR,ImportModuleThunk>(firstThunk,module)); | |
| return true; | |
| } | |
| */ | |
| void ImportsHandling::displayAllImports() | |
| { | |
| std::map<DWORD_PTR, ImportModuleThunk>::iterator it_module; | |
| std::map<DWORD_PTR, ImportThunk>::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.hTreeItem = addDllToTreeView(TreeImports, &moduleThunk); | |
| it_import = moduleThunk.thunkList.begin(); | |
| while (it_import != moduleThunk.thunkList.end()) | |
| { | |
| ImportThunk &importThunk = it_import->second; | |
| 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); | |
| 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); | |
| updateImportInTreeView(importThunk, item); | |
| return item; | |
| } | |
| void ImportsHandling::selectImports(bool invalid, bool suspect) | |
| { | |
| std::map<DWORD_PTR, ImportModuleThunk>::iterator it_module; | |
| std::map<DWORD_PTR, ImportThunk>::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::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); | |
| import = 0; | |
| if (module->thunkList.empty()) | |
| { | |
| itemData.erase(parent); | |
| module->hTreeItem.Delete(); | |
| moduleList.erase(module->firstThunk); | |
| 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); | |
| 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<DWORD_PTR, ImportModuleThunk>::iterator iterator1; | |
| std::map<DWORD_PTR, ImportThunk>::iterator iterator2; | |
| ImportModuleThunk * moduleThunk; | |
| ImportThunk * importThunk; | |
| iterator1 = moduleList.begin(); | |
| while (iterator1 != moduleList.end()) | |
| { | |
| moduleThunk = &(iterator1->second); | |
| iterator2 = moduleThunk->thunkList.begin(); | |
| while (iterator2 != moduleThunk->thunkList.end()) | |
| { | |
| importThunk = &(iterator2->second); | |
| if (importThunk->moduleName[0] == 0 || importThunk->moduleName[0] == L'?') | |
| { | |
| addNotFoundApiToModuleList(importThunk); | |
| } | |
| else | |
| { | |
| if (isNewModule(importThunk->moduleName)) | |
| { | |
| addModuleToModuleList(importThunk->moduleName, importThunk->rva); | |
| } | |
| addFunctionToModuleList(importThunk); | |
| } | |
| iterator2++; | |
| } | |
| moduleThunk->thunkList.clear(); | |
| iterator1++; | |
| } | |
| moduleList = moduleListNew; | |
| moduleListNew.clear(); | |
| } | |
| bool ImportsHandling::findNewModules(std::map<DWORD_PTR, ImportThunk> & 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<DWORD_PTR,ImportModuleThunk>(firstThunk,module)); | |
| return true; | |
| } | |
| bool ImportsHandling::isNewModule(const WCHAR * moduleName) | |
| { | |
| std::map<DWORD_PTR, ImportModuleThunk>::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<DWORD_PTR,ImportModuleThunk>(firstThunk,module)); | |
| } | |
| bool ImportsHandling::addNotFoundApiToModuleList(const ImportThunk * apiNotFound) | |
| { | |
| ImportThunk import; | |
| ImportModuleThunk * module = 0; | |
| std::map<DWORD_PTR, ImportModuleThunk>::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<DWORD_PTR,ImportThunk>(import.rva, import)); | |
| return true; | |
| } | |
| bool ImportsHandling::addFunctionToModuleList(const ImportThunk * apiFound) | |
| { | |
| ImportThunk import; | |
| ImportModuleThunk * module = 0; | |
| std::map<DWORD_PTR, ImportModuleThunk>::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<DWORD_PTR,ImportThunk>(import.rva, import)); | |
| return true; | |
| } | |
| void ImportsHandling::expandAllTreeNodes() | |
| { | |
| changeExpandStateOfTreeNodes(TVE_EXPAND); | |
| } | |
| void ImportsHandling::collapseAllTreeNodes() | |
| { | |
| changeExpandStateOfTreeNodes(TVE_COLLAPSE); | |
| } | |
| void ImportsHandling::changeExpandStateOfTreeNodes(UINT flag) | |
| { | |
| std::map<DWORD_PTR, ImportModuleThunk>::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/MainGui.cpp b/Scylla/MainGui.cpp | |
| index e6f01cc..4c94462 100644 | |
| --- a/Scylla/MainGui.cpp | |
| +++ b/Scylla/MainGui.cpp | |
| @@ -1,1241 +1,1243 @@ | |
| #include "MainGui.h" | |
| #include <atldlgs.h> // WTL common dialogs | |
| #include "definitions.h" | |
| #include "PluginLoader.h" | |
| #include "ConfigurationHolder.h" | |
| #include "PeDump.h" | |
| #include "PeRebuild.h" | |
| #include "DllInjectionPlugin.h" | |
| #include "DisassemblerGui.h" | |
| #include "PickApiGui.h" | |
| #include "NativeWinApi.h" | |
| #include "ImportRebuild.h" | |
| #include "SystemInformation.h" | |
| #include "AboutGui.h" | |
| #include "OptionsGui.h" | |
| extern CAppModule _Module; // o_O | |
| const WCHAR MainGui::filterExe[] = L"Executable (*.exe)\0*.exe\0All files\0*.*\0"; | |
| const WCHAR MainGui::filterDll[] = L"Dynamic Link Library (*.dll)\0*.dll\0All files\0*.*\0"; | |
| const WCHAR MainGui::filterExeDll[] = L"Executable (*.exe)\0*.exe\0Dynamic Link Library (*.dll)\0*.dll\0All files\0*.*\0"; | |
| const WCHAR MainGui::filterTxt[] = L"Text file (*.txt)\0*.txt\0All files\0*.*\0"; | |
| MainGui::MainGui() : selectedProcess(0), importsHandling(TreeImports), TreeImportsSubclass(this, IDC_TREE_IMPORTS) | |
| { | |
| Logger::getDebugLogFilePath(); | |
| ConfigurationHolder::loadConfiguration(); | |
| PluginLoader::findAllPlugins(); | |
| NativeWinApi::initialize(); | |
| SystemInformation::getSystemInformation(); | |
| if(ConfigurationHolder::getConfigObject(DEBUG_PRIVILEGE)->isTrue()) | |
| { | |
| processLister.setDebugPrivileges(); | |
| } | |
| processAccessHelp.getProcessModules(GetCurrentProcessId(), processAccessHelp.ownModuleList); | |
| hIcon.LoadIcon(IDI_ICON_SCYLLA); | |
| hMenuImports.LoadMenu(IDR_MENU_IMPORTS); | |
| hMenuLog.LoadMenu(IDR_MENU_LOG); | |
| accelerators.LoadAccelerators(IDR_ACCELERATOR_MAIN); | |
| hIconCheck.LoadIcon(IDI_ICON_CHECK, 16, 16); | |
| hIconWarning.LoadIcon(IDI_ICON_WARNING, 16, 16); | |
| hIconError.LoadIcon(IDI_ICON_ERROR, 16, 16); | |
| if(hMenuImports) | |
| { | |
| appendPluginListToMenu(hMenuImports.GetSubMenu(0)); | |
| } | |
| } | |
| BOOL MainGui::PreTranslateMessage(MSG* pMsg) | |
| { | |
| if(accelerators.TranslateAccelerator(m_hWnd, pMsg)) | |
| { | |
| return TRUE; | |
| } | |
| else if(IsDialogMessage(pMsg)) | |
| { | |
| return TRUE; | |
| } | |
| return FALSE; | |
| } | |
| BOOL MainGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) | |
| { | |
| if (SystemInformation::currenOS == UNKNOWN_OS) | |
| { | |
| if(IDCANCEL == MessageBox(L"Operating System is not supported\r\nContinue anyway?", L"Scylla", MB_ICONWARNING | MB_OKCANCEL)) | |
| { | |
| SendMessage(WM_CLOSE); | |
| return FALSE; | |
| } | |
| } | |
| CMessageLoop* pLoop = _Module.GetMessageLoop(); | |
| pLoop->AddMessageFilter(this); | |
| setupStatusBar(); | |
| fillStatusBar(); | |
| DoDataExchange(); // attach controls | |
| DlgResize_Init(true, true); | |
| EditOEPAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); | |
| EditIATAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); | |
| EditIATSize.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); | |
| appendPluginListToMenu(CMenuHandle(GetMenu()).GetSubMenu(MenuImportsOffsetTrace)); | |
| enableDialogControls(FALSE); | |
| setIconAndDialogCaption(); | |
| return TRUE; | |
| } | |
| void MainGui::OnDestroy() | |
| { | |
| PostQuitMessage(0); | |
| } | |
| void MainGui::OnSize(UINT nType, CSize size) | |
| { | |
| StatusBar.SendMessage(WM_SIZE); | |
| SetMsgHandled(FALSE); | |
| } | |
| void MainGui::OnLButtonDown(UINT nFlags, CPoint point) | |
| { | |
| SetMsgHandled(FALSE); | |
| } | |
| void MainGui::OnContextMenu(CWindow wnd, CPoint point) | |
| { | |
| // point = -1, -1 for keyboard invoked shortcut! | |
| switch(wnd.GetDlgCtrlID()) | |
| { | |
| case IDC_TREE_IMPORTS: | |
| DisplayContextMenuImports(wnd, point); | |
| return; | |
| case IDC_LIST_LOG: | |
| DisplayContextMenuLog(wnd, point); | |
| return; | |
| //default: // wnd == m_hWnd? | |
| // DisplayContextMenu(wnd, point); | |
| // return; | |
| } | |
| SetMsgHandled(FALSE); | |
| } | |
| void MainGui::OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl) | |
| { | |
| // Make sure it's a menu | |
| if(uNotifyCode == 0 && !wndCtl.IsWindow()) | |
| { | |
| if ((nID >= PLUGIN_MENU_BASE_ID) && (nID <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID))) | |
| { | |
| pluginActionHandler(nID); | |
| return; | |
| } | |
| } | |
| SetMsgHandled(FALSE); | |
| } | |
| LRESULT MainGui::OnTreeImportsClick(const NMHDR* pnmh) | |
| { | |
| SetMsgHandled(FALSE); | |
| return 0; | |
| } | |
| LRESULT MainGui::OnTreeImportsDoubleClick(const NMHDR* pnmh) | |
| { | |
| if(TreeImports.GetCount() < 1) | |
| return 0; | |
| // Get item under cursor | |
| CPoint client = GetMessagePos(); | |
| TreeImports.ScreenToClient(&client); | |
| UINT flags; | |
| CTreeItem over = TreeImports.HitTest(client, &flags); | |
| CTreeItem parent; | |
| if(over) | |
| { | |
| if(!(flags & TVHT_ONITEM)) | |
| { | |
| over = NULL; | |
| } | |
| else | |
| { | |
| parent = over.GetParent(); | |
| } | |
| } | |
| if(!over.IsNull() && !parent.IsNull()) | |
| { | |
| pickApiActionHandler(over); | |
| } | |
| return 0; | |
| } | |
| LRESULT MainGui::OnTreeImportsRightClick(const NMHDR* pnmh) | |
| { | |
| SetMsgHandled(FALSE); | |
| return 0; | |
| } | |
| LRESULT MainGui::OnTreeImportsRightDoubleClick(const NMHDR* pnmh) | |
| { | |
| SetMsgHandled(FALSE); | |
| return 0; | |
| } | |
| LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh) | |
| { | |
| const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh; | |
| switch(tkd->wVKey) | |
| { | |
| case VK_RETURN: | |
| { | |
| - CTreeItem selected = TreeImports.GetFirstSelectedItem(); | |
| + CTreeItem selected = TreeImports.GetFocusItem(); | |
| if(!selected.IsNull() && !selected.GetParent().IsNull()) | |
| { | |
| 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 | |
| } | |
| void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl) | |
| { | |
| // TODO | |
| } | |
| void MainGui::OnAutotrace(UINT uNotifyCode, int nID, CWindow wndCtl) | |
| { | |
| // TODO | |
| } | |
| void MainGui::OnExit(UINT uNotifyCode, int nID, CWindow wndCtl) | |
| { | |
| DestroyWindow(); | |
| } | |
| void MainGui::OnAbout(UINT uNotifyCode, int nID, CWindow wndCtl) | |
| { | |
| showAboutDialog(); | |
| } | |
| void MainGui::setupStatusBar() | |
| { | |
| StatusBar.Create(m_hWnd, NULL, L"", WS_CHILD | WS_VISIBLE | WS_CLIPCHILDREN | WS_CLIPSIBLINGS | SBARS_TOOLTIPS, NULL, 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<Process>& processList = processLister.getProcessList(); | |
| Process &process = processList.at(index); | |
| selectedProcess = 0; | |
| clearImportsActionHandler(); | |
| Logger::printfDialog(TEXT("Analyzing %s"),process.fullPath); | |
| if (processAccessHelp.hProcess != 0) | |
| { | |
| processAccessHelp.closeProcessHandle(); | |
| apiReader.clearAll(); | |
| } | |
| if (!processAccessHelp.openProcessHandle(process.PID)) | |
| { | |
| enableDialogControls(FALSE); | |
| Logger::printfDialog(TEXT("Error: Cannot open process handle.")); | |
| fillStatusBar(); | |
| return; | |
| } | |
| processAccessHelp.getProcessModules(process.PID, processAccessHelp.moduleList); | |
| apiReader.readApisFromModuleList(); | |
| Logger::printfDialog(TEXT("Loading modules done.")); | |
| //TODO improve | |
| processAccessHelp.selectedModule = 0; | |
| processAccessHelp.targetSizeOfImage = process.imageSize; | |
| processAccessHelp.targetImageBase = process.imageBase; | |
| ProcessAccessHelp::getSizeOfImageCurrentProcess(); | |
| process.imageSize = (DWORD)processAccessHelp.targetSizeOfImage; | |
| Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),process.imageBase, process.imageSize); | |
| process.entryPoint = ProcessAccessHelp::getEntryPointFromFile(process.fullPath); | |
| swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),process.entryPoint + process.imageBase); | |
| EditOEPAddress.SetWindowText(stringBuffer); | |
| selectedProcess = &process; | |
| enableDialogControls(TRUE); | |
| fillStatusBar(); | |
| } | |
| void MainGui::fillProcessListComboBox(CComboBox& hCombo) | |
| { | |
| hCombo.ResetContent(); | |
| std::vector<Process>& processList = processLister.getProcessListSnapshot(); | |
| for (size_t i = 0; i < processList.size(); i++) | |
| { | |
| swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("0x%04X - %s - %s"),processList[i].PID,processList[i].filename,processList[i].fullPath); | |
| hCombo.AddString(stringBuffer); | |
| } | |
| } | |
| void MainGui::addTextToOutputLog(const WCHAR * text) | |
| { | |
| if (m_hWnd) | |
| { | |
| ListLog.SetCurSel(ListLog.AddString(text)); | |
| } | |
| } | |
| void MainGui::clearOutputLog() | |
| { | |
| if (m_hWnd) | |
| { | |
| ListLog.ResetContent(); | |
| } | |
| } | |
| bool MainGui::saveLogToFile(const WCHAR * file) | |
| { | |
| const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian | |
| const WCHAR newLine[] = L"\r\n"; | |
| bool success = true; | |
| HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0); | |
| if(hFile != INVALID_HANDLE_VALUE) | |
| { | |
| ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM); | |
| WCHAR * buffer = 0; | |
| int bufsize = 0; | |
| for(int i = 0; i < ListLog.GetCount(); i++) | |
| { | |
| int size = ListLog.GetTextLen(i); | |
| size += _countof(newLine)-1; | |
| if(size+1 > bufsize) | |
| { | |
| bufsize = size+1; | |
| delete[] buffer; | |
| try | |
| { | |
| buffer = new WCHAR[bufsize]; | |
| } | |
| catch(std::bad_alloc&) | |
| { | |
| buffer = 0; | |
| success = false; | |
| break; | |
| } | |
| } | |
| ListLog.GetText(i, buffer); | |
| wcscat_s(buffer, bufsize, newLine); | |
| ProcessAccessHelp::writeMemoryToFileEnd(hFile, size * sizeof(WCHAR), buffer); | |
| } | |
| delete[] buffer; | |
| CloseHandle(hFile); | |
| } | |
| return success; | |
| } | |
| void MainGui::showInvalidImportsActionHandler() | |
| { | |
| importsHandling.selectImports(true, false); | |
| + GotoDlgCtrl(TreeImports); | |
| } | |
| void MainGui::showSuspectImportsActionHandler() | |
| { | |
| importsHandling.selectImports(false, true); | |
| + GotoDlgCtrl(TreeImports); | |
| } | |
| void MainGui::iatAutosearchActionHandler() | |
| { | |
| DWORD_PTR searchAddress = 0; | |
| DWORD_PTR addressIAT = 0; | |
| DWORD sizeIAT = 0; | |
| IATSearch iatSearch; | |
| if(!selectedProcess) | |
| return; | |
| if(EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 1) | |
| { | |
| searchAddress = stringToDwordPtr(stringBuffer); | |
| if (searchAddress) | |
| { | |
| if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT)) | |
| { | |
| Logger::printfDialog(TEXT("IAT found at VA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" RVA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d)"),addressIAT, addressIAT - processAccessHelp.targetImageBase,sizeIAT,sizeIAT); | |
| swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),addressIAT); | |
| EditIATAddress.SetWindowText(stringBuffer); | |
| swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("%08X"),sizeIAT); | |
| EditIATSize.SetWindowText(stringBuffer); | |
| - swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("IAT found! Start Address ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d) "),addressIAT,sizeIAT,sizeIAT); | |
| + 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) | |
| { | |
| addressIAT = stringToDwordPtr(stringBuffer); | |
| } | |
| if (EditIATSize.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0) | |
| { | |
| sizeIAT = wcstoul(stringBuffer, NULL, 16); | |
| } | |
| if (addressIAT && sizeIAT) | |
| { | |
| apiReader.readAndParseIAT(addressIAT, sizeIAT,importsHandling.moduleList); | |
| importsHandling.displayAllImports(); | |
| } | |
| fillStatusBar(); | |
| } | |
| DWORD_PTR MainGui::stringToDwordPtr(const WCHAR * hexString) | |
| { | |
| DWORD_PTR address = 0; | |
| #ifdef _WIN64 | |
| address = _wcstoui64(hexString, NULL, 16); | |
| #else | |
| address = wcstoul(hexString, NULL, 16); | |
| #endif | |
| if (address == 0) | |
| { | |
| #ifdef DEBUG_COMMENTS | |
| Logger::debugLog(L"stringToDwordPtr :: address == 0, %s",hexString); | |
| #endif | |
| return 0; | |
| } | |
| else | |
| { | |
| return address; | |
| } | |
| } | |
| void MainGui::SetupImportsMenuItems(bool isItem, bool isThunk) | |
| { | |
| // assert(!(!isItem && isThunk)); | |
| CMenuHandle hSub = hMenuImports.GetSubMenu(0); | |
| UINT itemOnly = isItem ? MF_ENABLED : MF_GRAYED; | |
| UINT thunkOnly = isThunk ? MF_ENABLED : MF_GRAYED; | |
| hSub.EnableMenuItem(ID__INVALIDATEFUNCTION, thunkOnly); | |
| hSub.EnableMenuItem(ID__DISASSEMBLE, thunkOnly); | |
| hSub.EnableMenuItem(ID__CUTTHUNK, thunkOnly); | |
| hSub.EnableMenuItem(ID__DELETETREENODE, itemOnly); | |
| } | |
| void MainGui::DisplayContextMenuImports(CWindow hwnd, CPoint pt) | |
| { | |
| if(TreeImports.GetCount() < 1) | |
| return; | |
| CTreeItem over, parent; | |
| if(pt.x == -1 && pt.y == -1) // invoked by keyboard | |
| { | |
| CRect pos; | |
| - over = TreeImports.GetFirstSelectedItem(); | |
| + 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)); | |
| 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); | |
| 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(); // ? | |
| + fillStatusBar(); | |
| } | |
| void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt) | |
| { | |
| if (hMenuLog) | |
| { | |
| if(pt.x == -1 && pt.y == -1) // invoked by keyboard | |
| { | |
| CRect pos; | |
| ListLog.GetWindowRect(&pos); | |
| pt = pos.TopLeft(); | |
| } | |
| CMenuHandle hSub = hMenuLog.GetSubMenu(0); | |
| BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD, pt.x, pt.y, hwnd); | |
| if (menuItem) | |
| { | |
| switch (menuItem) | |
| { | |
| case ID__SAVE: | |
| WCHAR selectedFilePath[MAX_PATH]; | |
| if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt")) | |
| { | |
| saveLogToFile(selectedFilePath); | |
| } | |
| break; | |
| case ID__CLEAR: | |
| clearOutputLog(); | |
| break; | |
| } | |
| } | |
| } | |
| } | |
| void MainGui::appendPluginListToMenu(CMenuHandle hMenu) | |
| { | |
| std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList(); | |
| std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList(); | |
| if (scyllaPluginList.size() > 0) | |
| { | |
| CMenuHandle newMenu; | |
| newMenu.CreatePopupMenu(); | |
| for (size_t i = 0; i < scyllaPluginList.size(); i++) | |
| { | |
| newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName); | |
| } | |
| hMenu.AppendMenu(MF_MENUBARBREAK); | |
| hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins"); | |
| } | |
| if (imprecPluginList.size() > 0) | |
| { | |
| CMenuHandle newMenu; | |
| newMenu.CreatePopupMenu(); | |
| for (size_t i = 0; i < imprecPluginList.size(); i++) | |
| { | |
| newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName); | |
| } | |
| hMenu.AppendMenu(MF_MENUBARBREAK); | |
| hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins"); | |
| } | |
| } | |
| void MainGui::dumpActionHandler() | |
| { | |
| if(!selectedProcess) | |
| return; | |
| WCHAR selectedFilePath[MAX_PATH]; | |
| const WCHAR * fileFilter; | |
| const WCHAR * defExtension; | |
| PeDump peDump; | |
| if (processAccessHelp.selectedModule) | |
| { | |
| fileFilter = filterDll; | |
| defExtension = L"dll"; | |
| } | |
| else | |
| { | |
| fileFilter = filterExe; | |
| defExtension = L"exe"; | |
| } | |
| if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension)) | |
| { | |
| if (processAccessHelp.selectedModule) | |
| { | |
| //dump DLL | |
| peDump.imageBase = processAccessHelp.selectedModule->modBaseAddr; | |
| peDump.sizeOfImage = processAccessHelp.selectedModule->modBaseSize; | |
| //get it from gui | |
| peDump.entryPoint = getOEPFromGui(); | |
| wcscpy_s(peDump.fullpath, MAX_PATH, processAccessHelp.selectedModule->fullPath); | |
| } | |
| else | |
| { | |
| peDump.imageBase = ProcessAccessHelp::targetImageBase; | |
| peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage; | |
| //get it from gui | |
| peDump.entryPoint = getOEPFromGui(); | |
| wcscpy_s(peDump.fullpath, MAX_PATH, selectedProcess->fullPath); | |
| } | |
| peDump.useHeaderFromDisk = ConfigurationHolder::getConfigObject(USE_PE_HEADER_FROM_DISK)->isTrue(); | |
| if (peDump.dumpCompleteProcessToDisk(selectedFilePath)) | |
| { | |
| Logger::printfDialog(TEXT("Dump success %s"),selectedFilePath); | |
| //MessageBox(L"Image dumped successfully.", L"Success"); | |
| } | |
| else | |
| { | |
| Logger::printfDialog(TEXT("Error: Cannot dump image.")); | |
| MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR); | |
| } | |
| } | |
| } | |
| DWORD_PTR MainGui::getOEPFromGui() | |
| { | |
| if (EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0) | |
| { | |
| return stringToDwordPtr(stringBuffer); | |
| } | |
| else | |
| { | |
| return 0; | |
| } | |
| } | |
| void MainGui::peRebuildActionHandler() | |
| { | |
| DWORD newSize = 0; | |
| WCHAR selectedFilePath[MAX_PATH]; | |
| PeRebuild peRebuild; | |
| if(showFileDialog(selectedFilePath, false, NULL, filterExeDll)) | |
| { | |
| if (ConfigurationHolder::getConfigObject(CREATE_BACKUP)->isTrue()) | |
| { | |
| if (!ProcessAccessHelp::createBackupFile(selectedFilePath)) | |
| { | |
| Logger::printfDialog(TEXT("Creating backup file failed %s"), selectedFilePath); | |
| } | |
| } | |
| LONGLONG fileSize = ProcessAccessHelp::getFileSize(selectedFilePath); | |
| LPVOID mapped = peRebuild.createFileMappingViewFull(selectedFilePath); | |
| newSize = peRebuild.realignPE(mapped, (DWORD)fileSize); | |
| peRebuild.closeAllMappingHandles(); | |
| if (newSize < 10) | |
| { | |
| Logger::printfDialog(TEXT("Rebuild failed %s"), selectedFilePath); | |
| MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR); | |
| } | |
| else | |
| { | |
| peRebuild.truncateFile(selectedFilePath, newSize); | |
| Logger::printfDialog(TEXT("Rebuild success %s"), selectedFilePath); | |
| Logger::printfDialog(TEXT("-> Old file size 0x%08X new file size 0x%08X (%d %%)"), (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) ); | |
| //MessageBox(L"Image rebuilded successfully.", L"Success", MB_ICONINFORMATION); | |
| } | |
| } | |
| } | |
| void MainGui::dumpFixActionHandler() | |
| { | |
| if(!selectedProcess) | |
| return; | |
| if (TreeImports.GetCount() < 2) | |
| { | |
| Logger::printfDialog(TEXT("Nothing to rebuild")); | |
| return; | |
| } | |
| WCHAR newFilePath[MAX_PATH]; | |
| WCHAR selectedFilePath[MAX_PATH]; | |
| const WCHAR * fileFilter; | |
| ImportRebuild importRebuild; | |
| if (processAccessHelp.selectedModule) | |
| { | |
| fileFilter = filterDll; | |
| } | |
| else | |
| { | |
| fileFilter = filterExe; | |
| } | |
| if (showFileDialog(selectedFilePath, false, NULL, fileFilter)) | |
| { | |
| wcscpy_s(newFilePath,MAX_PATH,selectedFilePath); | |
| const WCHAR * extension = 0; | |
| WCHAR* dot = wcsrchr(newFilePath, L'.'); | |
| if (dot) | |
| { | |
| *dot = L'\0'; | |
| extension = selectedFilePath + (dot - newFilePath); //wcsrchr(selectedFilePath, L'.'); | |
| } | |
| wcscat_s(newFilePath, MAX_PATH, L"_SCY"); | |
| if(extension) | |
| { | |
| wcscat_s(newFilePath, MAX_PATH, extension); | |
| } | |
| if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList)) | |
| { | |
| //MessageBox(L"Imports rebuilding successful", L"Success", MB_ICONINFORMATION); | |
| Logger::printfDialog(TEXT("Import Rebuild success %s"), newFilePath); | |
| } | |
| else | |
| { | |
| Logger::printfDialog(TEXT("Import Rebuild failed, target %s"), selectedFilePath); | |
| MessageBox(L"Imports rebuilding failed", L"Failure", MB_ICONERROR); | |
| } | |
| } | |
| } | |
| void MainGui::enableDialogControls(BOOL value) | |
| { | |
| 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_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<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList(); | |
| std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList(); | |
| menuItem -= PLUGIN_MENU_BASE_ID; | |
| dllInjectionPlugin.hProcess = ProcessAccessHelp::hProcess; | |
| dllInjectionPlugin.apiReader = &apiReader; | |
| if (menuItem < (int)scyllaPluginList.size()) | |
| { | |
| //scylla plugin | |
| dllInjectionPlugin.injectPlugin(scyllaPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); | |
| } | |
| else | |
| { | |
| #ifndef _WIN64 | |
| menuItem -= (int)scyllaPluginList.size(); | |
| //imprec plugin | |
| dllInjectionPlugin.injectImprecPlugin(imprecPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); | |
| #endif | |
| } | |
| importsHandling.scanAndFixModuleList(); | |
| importsHandling.displayAllImports(); | |
| fillStatusBar(); | |
| } | |
| diff --git a/Scylla/multitree.h b/Scylla/multitree.h | |
| index 59b8dc1..7213139 100644 | |
| --- a/Scylla/multitree.h | |
| +++ b/Scylla/multitree.h | |
| @@ -1,600 +1,600 @@ | |
| #pragma once | |
| ///////////////////////////////////////////////////////////////////////////// | |
| // MultiSelectTree - A tree control with multi-select capabilities | |
| // | |
| // Written by Bjarke Viksoe ([email protected]) | |
| // Copyright (c) 2005 Bjarke Viksoe. | |
| // | |
| // Add the following macro to the parent's message map: | |
| // REFLECT_NOTIFICATIONS() | |
| // | |
| // This code may be used in compiled form in any way you desire. This | |
| // source file may be redistributed by any means PROVIDING it is | |
| // not sold for profit without the authors written consent, and | |
| // providing that this notice and the authors name is included. | |
| // | |
| // This file is provided "as is" with no expressed or implied warranty. | |
| // The author accepts no liability if it causes any damage to you or your | |
| // computer whatsoever. It's free, so don't hassle me about it. | |
| // | |
| // Beware of bugs. | |
| // | |
| #ifndef __cplusplus | |
| #error WTL requires C++ compilation (use a .cpp suffix) | |
| #endif | |
| #ifndef __ATLMISC_H__ | |
| #error multitree.h requires atlmisc.h to be included first | |
| #endif | |
| #ifndef __ATLCRACK_H__ | |
| #error multitree.h requires atlcrack.h to be included first | |
| #endif | |
| #ifndef __ATLCTRLS_H__ | |
| #error multitree.h requires atlctrls.h to be included first | |
| #endif | |
| // Extended MultiSelectTree styles | |
| static const DWORD MTVS_EX_NOMARQUEE = 0x00000001; | |
| // New control notifications | |
| static const UINT TVN_ITEMSELECTING = 0x0001; | |
| static const UINT TVN_ITEMSELECTED = 0x0002; | |
| static bool operator==(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<CTreeItem, bool> m_aData; | |
| CMultiSelectTreeViewImpl() : m_dwExStyle(0), m_bMarquee(false) | |
| { | |
| } | |
| CTreeItem _MakeItem(HTREEITEM hItem) const | |
| { | |
| return CTreeItem(hItem, (CTreeViewCtrlEx*)this); | |
| } | |
| // Operations | |
| BOOL SubclassWindow(HWND hWnd) | |
| { | |
| ATLASSERT(m_hWnd == NULL); | |
| ATLASSERT(::IsWindow(hWnd)); | |
| BOOL bRet = CWindowImpl< T, TBase, TWinTraits >::SubclassWindow(hWnd); | |
| if( bRet ) | |
| _Init(); | |
| return bRet; | |
| } | |
| DWORD SetExtendedTreeStyle(DWORD dwStyle) | |
| { | |
| ATLASSERT(!m_ctrlTree.IsWindow()); // Before control is created, please! | |
| DWORD dwOldStyle = m_dwTreeStyle; | |
| m_dwTreeStyle = dwStyle; | |
| return dwOldStyle; | |
| } | |
| 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)); | |
| 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)); | |
| if( iIndex >= 0 ) | |
| { | |
| nRes &= ~TVIS_SELECTED; | |
| if( m_aData.GetValueAt(iIndex) ) | |
| nRes |= TVIS_SELECTED; | |
| } | |
| } | |
| return nRes; | |
| } | |
| CTreeItem GetFirstSelectedItem() const | |
| { | |
| if( m_aData.GetSize() > 0 ) | |
| { | |
| for( int i = 0; i < m_aData.GetSize(); i++ ) | |
| { | |
| if( m_aData.GetValueAt(i) ) | |
| return m_aData.GetKeyAt(i); | |
| } | |
| } | |
| return _MakeItem(NULL); | |
| } | |
| CTreeItem GetNextSelectedItem(HTREEITEM hItem) const | |
| { | |
| int iIndex = m_aData.FindKey(_MakeItem(hItem)); | |
| if( iIndex >= 0 ) | |
| { | |
| for( int i = iIndex + 1; i < m_aData.GetSize(); i++ ) | |
| { | |
| if( m_aData.GetValueAt(i) ) | |
| return m_aData.GetKeyAt(i); | |
| } | |
| } | |
| return _MakeItem(NULL); | |
| } | |
| int GetSelectedCount() const | |
| { | |
| int nCount = 0; | |
| for( int i = 0; i < m_aData.GetSize(); i++ ) | |
| { | |
| if( m_aData.GetValueAt(i) ) | |
| nCount++; | |
| } | |
| return nCount; | |
| } | |
| // Implementation | |
| void _Init() | |
| { | |
| ATLASSERT(::IsWindow(m_hWnd)); | |
| ModifyStyle(TVS_SHOWSELALWAYS, 0); | |
| } | |
| void _SelectItem(int iIndex, bool bSelect, int action = TVC_UNKNOWN) | |
| { | |
| if( iIndex < 0 ) | |
| return; | |
| bool bSelected = m_aData.GetValueAt(iIndex); | |
| // Don't change if state is already updated (avoids flicker) | |
| if( bSelected == bSelect ) | |
| return; | |
| CTreeItem cItem = m_aData.GetKeyAt(iIndex); | |
| CWindow parent = GetParent(); | |
| // Send notifications | |
| NMTREEVIEW nmtv = { 0 }; | |
| nmtv.hdr.code = TVN_ITEMSELECTING; | |
| nmtv.hdr.hwndFrom = m_hWnd; | |
| nmtv.hdr.idFrom = GetDlgCtrlID(); | |
| nmtv.action = action; | |
| nmtv.itemNew.hItem = cItem; | |
| nmtv.itemNew.lParam = GetItemData(cItem); | |
| nmtv.itemNew.state = bSelect ? TVIS_SELECTED : 0; | |
| nmtv.itemNew.stateMask = TVIS_SELECTED; | |
| if( parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv) != 0 ) | |
| return; | |
| // Change state | |
| m_aData.SetAtIndex(iIndex, cItem, bSelect); | |
| // Repaint item | |
| CRect rcItem; | |
| if( GetItemRect(cItem, &rcItem, FALSE) ) | |
| InvalidateRect(&rcItem, TRUE); | |
| // More notifications | |
| nmtv.hdr.code = TVN_ITEMSELECTED; | |
| parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv); | |
| } | |
| void _SelectItem(HTREEITEM hItem, bool bSelect, int action = TVC_UNKNOWN) | |
| { | |
| _SelectItem(m_aData.FindKey(_MakeItem(hItem)), bSelect, action); | |
| } | |
| void _SelectTree(HTREEITEM hItem, HTREEITEM hGoal, int action) | |
| { | |
| if( !_SelectTreeSub(hItem, hGoal, action) ) | |
| return; | |
| hItem = GetParentItem(hItem); | |
| while( (hItem = GetNextSiblingItem(hItem)) != NULL ) | |
| { | |
| if( !_SelectTreeSub(hItem, hGoal, action) ) | |
| return; | |
| } | |
| } | |
| bool _SelectTreeSub(HTREEITEM hItem, HTREEITEM hGoal, int action) | |
| { | |
| while( hItem != NULL ) | |
| { | |
| _SelectItem(hItem, true, action); | |
| if( hItem == hGoal ) | |
| return false; | |
| if( (TBase::GetItemState(hItem, TVIS_EXPANDED) & TVIS_EXPANDED) != 0 ) | |
| { | |
| if( !_SelectTreeSub(GetChildItem(hItem), hGoal, action) ) | |
| return false; | |
| } | |
| hItem = GetNextSiblingItem(hItem); | |
| } | |
| return true; | |
| } | |
| /* | |
| TODO: keep list of 'temporary' selected items | |
| if you select an item with the mouse and it gets | |
| out of the selection rectangle it stays selected. | |
| it should unselect it but keep old selected items | |
| (when holding CTRL) | |
| */ | |
| void _SelectBox(CRect rc) | |
| { | |
| HTREEITEM hItem = GetFirstVisibleItem(); | |
| while( hItem != NULL ) | |
| { | |
| int i = m_aData.FindKey(_MakeItem(hItem)); | |
| if(i >= 0 && !m_aData.GetValueAt(i)) // ignore already selected | |
| { | |
| CRect rcItem, rcTemp; | |
| GetItemRect(hItem, &rcItem, TRUE); | |
| _SelectItem(hItem, rcTemp.IntersectRect(&rcItem, &rc) == TRUE, TVC_BYMOUSE); | |
| } | |
| hItem = GetNextVisibleItem(hItem); | |
| } | |
| } | |
| void _DrawDragRect(CPoint pt) | |
| { | |
| CClientDC dc(m_hWnd); | |
| CSize szFrame(1, 1); | |
| CRect rect(m_ptDragStart, pt); | |
| rect.NormalizeRect(); | |
| //CRect rectOld(m_ptDragStart, m_ptDragOld); | |
| //rectOld.NormalizeRect(); | |
| dc.DrawDragRect(&rect, szFrame, NULL/*&rectOld*/, szFrame); | |
| } | |
| // Message map and handlers | |
| BEGIN_MSG_MAP_EX(CMultiSelectTreeViewImpl) | |
| MESSAGE_HANDLER_EX(WM_CREATE, OnCreate) | |
| MSG_WM_DESTROY(OnDestroy) | |
| MSG_WM_KEYDOWN(OnKeyDown) | |
| MSG_WM_KEYUP(OnKeyUp) | |
| MSG_WM_CHAR(OnChar) | |
| MSG_WM_SETFOCUS(OnSetFocus) | |
| MSG_WM_LBUTTONDOWN(OnLButtonDown) | |
| MSG_WM_LBUTTONUP(OnLButtonUp) | |
| MSG_WM_MOUSEMOVE(OnMouseMove) | |
| MSG_WM_CAPTURECHANGED(OnCaptureChanged) | |
| MESSAGE_HANDLER_EX(TVM_INSERTITEM, OnInsertItem) | |
| REFLECTED_NOTIFY_CODE_HANDLER_EX(TVN_DELETEITEM, OnDeleteItem) | |
| CHAIN_MSG_MAP_ALT( CCustomDraw< T >, 1 ) | |
| END_MSG_MAP() | |
| LRESULT OnCreate(UINT uMsg, WPARAM wParam, LPARAM lParam) | |
| { | |
| LRESULT lRes = DefWindowProc(); | |
| _Init(); | |
| return lRes; | |
| } | |
| void OnDestroy() | |
| { | |
| m_aData.RemoveAll(); | |
| SetMsgHandled(FALSE); | |
| } | |
| void OnKeyDown(UINT nChar, UINT nRepCnt, UINT nFlags) | |
| { | |
| if( nChar == VK_SHIFT ) | |
| m_hExtSelStart = GetFocusItem(); | |
| if( ::GetKeyState(VK_SHIFT) < 0 && m_hExtSelStart == GetFocusItem() ) | |
| { | |
| switch( nChar ) | |
| { | |
| case VK_UP: | |
| case VK_DOWN: | |
| case VK_HOME: | |
| case VK_END: | |
| case VK_NEXT: | |
| case VK_PRIOR: | |
| for( int i = 0; i < m_aData.GetSize(); i++ ) | |
| { | |
| _SelectItem(i, false, TVC_BYKEYBOARD); | |
| } | |
| } | |
| } | |
| SetMsgHandled(FALSE); | |
| } | |
| void OnKeyUp(UINT nChar, UINT nRepCnt, UINT nFlags) | |
| { | |
| if( ::GetKeyState(VK_SHIFT) < 0 ) | |
| { | |
| switch( nChar ) | |
| { | |
| case VK_UP: | |
| case VK_DOWN: | |
| case VK_HOME: | |
| case VK_END: | |
| case VK_NEXT: | |
| case VK_PRIOR: | |
| HTREEITEM hItem = GetFocusItem(); | |
| // Is current or first-shift-item the upper item? | |
| CRect rcItem1, rcItem2; | |
| GetItemRect(m_hExtSelStart, &rcItem1, TRUE); | |
| GetItemRect(hItem, &rcItem2, TRUE); | |
| // Select from current item to item where SHIFT was pressed | |
| if( rcItem1.top > rcItem2.top ) | |
| _SelectTree(hItem, m_hExtSelStart, TVC_BYKEYBOARD); | |
| else | |
| _SelectTree(m_hExtSelStart, hItem, TVC_BYKEYBOARD); | |
| _SelectItem(hItem, true, TVC_BYKEYBOARD); | |
| } | |
| } | |
| SetMsgHandled(FALSE); | |
| } | |
| void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags) | |
| { | |
| if( nChar == VK_SPACE ) | |
| { | |
| HTREEITEM hItem = GetFocusItem(); | |
| _SelectItem(hItem, (GetItemState(hItem, TVIS_SELECTED) & TVIS_SELECTED) == 0, TVC_BYKEYBOARD); | |
| return; | |
| } | |
| SetMsgHandled(FALSE); | |
| } | |
| void OnSetFocus(CWindow wndOld) | |
| { | |
| DefWindowProc(); | |
| // FIX: We really need the focus-rectangle in this control since it | |
| // improves the navigation a lot. So let's ask Windows to display it. | |
| SendMessage(WM_UPDATEUISTATE, MAKEWPARAM(UIS_CLEAR, UISF_HIDEFOCUS)); | |
| } | |
| void OnLButtonDown(UINT nFlags, CPoint point) | |
| { | |
| SetMsgHandled(FALSE); | |
| // Hit-test and figure out where we're clicking... | |
| TVHITTESTINFO hti = { 0 }; | |
| hti.pt = point; | |
| HTREEITEM hItem = HitTest(&hti); | |
| if( (hItem == NULL || (hti.flags & TVHT_ONITEMRIGHT) != 0) ) | |
| { | |
| if( (m_dwExStyle & MTVS_EX_NOMARQUEE) == 0 && ::DragDetect(m_hWnd, point) ) | |
| { | |
| // Great we're dragging a rubber-band | |
| - // Clear selection of CTRL is not down | |
| + // 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(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)); | |
| 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); | |
| return (LRESULT) hItem; | |
| } | |
| LRESULT OnDeleteItem(NMHDR* pnmh) | |
| { | |
| const NMTREEVIEW* lpNMTV = (NMTREEVIEW*) pnmh; | |
| m_aData.Remove(_MakeItem(lpNMTV->itemNew.hItem)); | |
| return 0; | |
| } | |
| // Custom Draw | |
| DWORD OnPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* /*lpNMCustomDraw*/) | |
| { | |
| return CDRF_NOTIFYITEMDRAW; // We need per-item notifications | |
| } | |
| DWORD OnItemPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* lpNMCustomDraw) | |
| { | |
| NMTVCUSTOMDRAW* lpTVCD = (NMTVCUSTOMDRAW*) lpNMCustomDraw; | |
| HTREEITEM hItem = (HTREEITEM) lpTVCD->nmcd.dwItemSpec; | |
| int iIndex = m_aData.FindKey(_MakeItem(hItem)); | |
| if( iIndex >= 0 ) | |
| { | |
| bool bSelected = m_aData.GetValueAt(iIndex); | |
| // Trick TreeView into displaying correct selection colors | |
| if( bSelected ) | |
| { | |
| lpTVCD->clrText = ::GetSysColor(COLOR_HIGHLIGHTTEXT); | |
| lpTVCD->clrTextBk = ::GetSysColor(COLOR_HIGHLIGHT); | |
| } | |
| else | |
| { | |
| // Special case of tree-item actually have selection, but our | |
| // state says it is currently not selected (CTRL+click on same item twice). | |
| if( (lpTVCD->nmcd.uItemState & CDIS_SELECTED) != 0 ) | |
| { | |
| COLORREF clrText = GetTextColor(); | |
| if( clrText == CLR_NONE ) | |
| clrText = ::GetSysColor(COLOR_WINDOWTEXT); | |
| COLORREF clrBack = GetBkColor(); | |
| if( clrBack == CLR_NONE ) | |
| clrBack = ::GetSysColor(COLOR_WINDOW); | |
| //CDCHandle dc = lpTVCD->nmcd.hdc; | |
| //dc.SetTextColor(clrText); | |
| //dc.SetBkColor(clrBack); | |
| lpTVCD->clrText = clrText; | |
| lpTVCD->clrTextBk = clrBack; | |
| } | |
| } | |
| return CDRF_NEWFONT; | |
| } | |
| return CDRF_DODEFAULT; | |
| } | |
| }; | |
| class CMultiSelectTreeViewCtrl : public CMultiSelectTreeViewImpl<CMultiSelectTreeViewCtrl, CTreeViewCtrlEx, CWinTraitsOR<TVS_SHOWSELALWAYS> > | |
| { | |
| public: | |
| DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName()) | |
| }; |
File Metadata
File Metadata
- Mime Type
- text/x-diff
- Expires
- Tue, Jan 7, 12:03 AM (2 h, 7 m)
- Storage Engine
- local-disk
- Storage Format
- Raw Data
- Storage Handle
- d3/31/8a035b0c38a3c2d0ef0674096b96