Page MenuHomedesp's stash

No OneTemporary

diff --git a/Scylla/ImportsHandling.cpp b/Scylla/ImportsHandling.cpp
index bbe7ab4..7b12ac0 100644
--- a/Scylla/ImportsHandling.cpp
+++ b/Scylla/ImportsHandling.cpp
@@ -1,784 +1,788 @@
#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
-ImportsHandling::ImportsHandling(CTreeViewCtrlEx& TreeImports) : TreeImports(TreeImports)
+ImportsHandling::ImportsHandling(CMultiSelectTreeViewCtrl& TreeImports) : TreeImports(TreeImports)
{
hIconCheck.LoadIcon(IDI_ICON_CHECK);
hIconWarning.LoadIcon(IDI_ICON_WARNING);
hIconError.LoadIcon(IDI_ICON_ERROR);
TreeIcons.Create(16, 16, ILC_COLOR32, 3, 1);
TreeIcons.AddIcon(hIconCheck);
TreeIcons.AddIcon(hIconWarning);
TreeIcons.AddIcon(hIconError);
}
ImportsHandling::~ImportsHandling()
{
TreeIcons.Destroy();
}
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;
}
}
/*bool ImportsHandling::addModule(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;
}*/
/*bool ImportsHandling::addFunction(WCHAR * moduleName, char * name, DWORD_PTR va, DWORD_PTR rva, DWORD_PTR 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;
}*/
void ImportsHandling::displayAllImports()
{
std::map<DWORD_PTR, ImportModuleThunk>::iterator it_module;
std::map<DWORD_PTR, ImportThunk>::iterator it_thunk;
ImportModuleThunk * moduleThunk;
ImportThunk * importThunk;
CTreeItem module;
CTreeItem apiFunction;
TreeImports.DeleteAllItems();
TreeImports.SetImageList(TreeIcons);
it_module = moduleList.begin();
while (it_module != moduleList.end())
{
moduleThunk = &(it_module->second);
module = addDllToTreeView(TreeImports,moduleThunk);
moduleThunk->hTreeItem = module;
it_thunk = moduleThunk->thunkList.begin();
while (it_thunk != moduleThunk->thunkList.end())
{
importThunk = &(it_thunk->second);
apiFunction = addApiToTreeView(TreeImports,module,importThunk);
importThunk->hTreeItem = apiFunction;
it_thunk++;
}
it_module++;
}
}
-CTreeItem ImportsHandling::addDllToTreeView(CTreeViewCtrlEx& idTreeView, const ImportModuleThunk * importThunk)
+CTreeItem ImportsHandling::addDllToTreeView(CMultiSelectTreeViewCtrl& idTreeView, const ImportModuleThunk * importThunk)
{
CTreeItem item = idTreeView.InsertItem(L"", NULL, TVI_ROOT);
updateModuleInTreeView(importThunk, item);
return item;
}
-CTreeItem ImportsHandling::addApiToTreeView(CTreeViewCtrlEx& idTreeView, CTreeItem parentDll, const ImportThunk * importThunk)
+CTreeItem ImportsHandling::addApiToTreeView(CMultiSelectTreeViewCtrl& idTreeView, CTreeItem parentDll, const ImportThunk * importThunk)
{
CTreeItem item = idTreeView.InsertItem(L"", parentDll, TVI_LAST);
updateImportInTreeView(importThunk, item);
return item;
}
void ImportsHandling::showImports(bool invalid, bool suspect)
{
std::map<DWORD_PTR, ImportModuleThunk>::iterator iterator1;
std::map<DWORD_PTR, ImportThunk>::iterator iterator2;
ImportModuleThunk * moduleThunk;
ImportThunk * importThunk;
- TreeImports.SetFocus();
- TreeImports.SelectItem(NULL); //remove selection
+ TreeImports.SetFocus(); // should be GotoDlgCtrl...
+ TreeImports.SelectAllItems(FALSE); //remove selection
iterator1 = moduleList.begin();
while (iterator1 != moduleList.end())
{
moduleThunk = &(iterator1->second);
iterator2 = moduleThunk->thunkList.begin();
while (iterator2 != moduleThunk->thunkList.end())
{
importThunk = &(iterator2->second);
if (invalid && !importThunk->valid)
{
selectItem(importThunk->hTreeItem);
setFocus(TreeImports, importThunk->hTreeItem);
}
else if (suspect && importThunk->suspect)
{
selectItem(importThunk->hTreeItem);
setFocus(TreeImports, importThunk->hTreeItem);
}
else
{
unselectItem(importThunk->hTreeItem);
}
iterator2++;
}
iterator1++;
}
}
bool ImportsHandling::isItemSelected(CTreeItem hItem)
{
const UINT state = TVIS_SELECTED;
return ((hItem.GetState(state) & state) == state);
}
void ImportsHandling::unselectItem(CTreeItem htItem)
{
selectItem(htItem, false);
}
bool ImportsHandling::selectItem(CTreeItem hItem, bool select)
{
const UINT state = TVIS_SELECTED;
return FALSE != hItem.SetState((select ? state : 0), state);
}
-void ImportsHandling::setFocus(CTreeViewCtrlEx& hwndTV, CTreeItem htItem)
+void ImportsHandling::setFocus(CMultiSelectTreeViewCtrl& hwndTV, CTreeItem htItem)
{
// the current focus
- CTreeItem htFocus = hwndTV.GetSelectedItem();
+ CTreeItem htFocus = hwndTV.GetFirstSelectedItem();
if ( htItem )
{
// set the focus
if ( htItem != htFocus )
{
// remember the selection state of the item
bool wasSelected = isItemSelected(htItem);
if ( htFocus && isItemSelected(htFocus) )
{
// prevent the tree from unselecting the old focus which it
// would do by default (TreeView_SelectItem unselects the
// focused item)
- hwndTV.SelectItem(NULL);
+ hwndTV.SelectAllItems(FALSE);
selectItem(htFocus);
}
hwndTV.SelectItem(htItem);
if ( !wasSelected )
{
// need to clear the selection which TreeView_SelectItem() gave
// us
unselectItem(htItem);
}
//else: was selected, still selected - ok
}
//else: nothing to do, focus already there
}
else
{
if ( htFocus )
{
bool wasFocusSelected = isItemSelected(htFocus);
// just clear the focus
hwndTV.SelectItem(NULL);
if ( wasFocusSelected )
{
// restore the selection state
selectItem(htFocus);
}
}
//else: nothing to do, no focus already
}
}
bool ImportsHandling::invalidateFunction(CTreeItem selectedTreeNode)
{
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->hTreeItem == selectedTreeNode)
{
importThunk->ordinal = 0;
importThunk->hint = 0;
importThunk->valid = false;
importThunk->suspect = false;
importThunk->moduleName[0] = 0;
importThunk->name[0] = 0;
updateImportInTreeView(importThunk, importThunk->hTreeItem);
updateModuleInTreeView(moduleThunk, moduleThunk->hTreeItem);
return true;
}
iterator2++;
}
iterator1++;
}
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::cutThunk(CTreeItem selectedTreeNode)
{
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->hTreeItem == selectedTreeNode)
{
importThunk->hTreeItem.Delete();
moduleThunk->thunkList.erase(iterator2);
if (moduleThunk->thunkList.empty())
{
moduleThunk->hTreeItem.Delete();
moduleList.erase(iterator1);
}
else
{
updateModuleInTreeView(moduleThunk, moduleThunk->hTreeItem);
}
return true;
}
iterator2++;
}
iterator1++;
}
return false;
}
bool ImportsHandling::deleteTreeNode(CTreeItem selectedTreeNode)
{
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);
if (moduleThunk->hTreeItem == selectedTreeNode)
{
moduleThunk->hTreeItem.Delete();
moduleThunk->thunkList.clear();
moduleList.erase(iterator1);
return true;
}
else
{
iterator2 = moduleThunk->thunkList.begin();
while (iterator2 != moduleThunk->thunkList.end())
{
importThunk = &(iterator2->second);
if (importThunk->hTreeItem == selectedTreeNode)
{
moduleThunk->hTreeItem.Delete();
moduleThunk->thunkList.clear();
moduleList.erase(iterator1);
return true;
}
iterator2++;
}
}
iterator1++;
}
return false;
}
DWORD_PTR ImportsHandling::getApiAddressByNode(CTreeItem selectedTreeNode)
{
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->hTreeItem == selectedTreeNode)
{
return importThunk->apiAddressVA;
}
iterator2++;
}
iterator1++;
}
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, MAX_PATH, moduleName);
moduleListNew.insert(std::pair<DWORD_PTR,ImportModuleThunk>(firstThunk,module));
return true;
}
bool ImportsHandling::isNewModule(const WCHAR * moduleName)
{
std::map<DWORD_PTR, ImportModuleThunk>::iterator iterator1;
iterator1 = moduleListNew.begin();
while (iterator1 != moduleListNew.end())
{
if (!_wcsicmp(iterator1->second.moduleName, moduleName))
{
return false;
}
iterator1++;
}
return true;
}
void ImportsHandling::addUnknownModuleToModuleList(DWORD_PTR firstThunk)
{
ImportModuleThunk module;
module.firstThunk = firstThunk;
wcscpy_s(module.moduleName, MAX_PATH, 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 iterator1;
DWORD_PTR rva = apiNotFound->rva;
if (moduleListNew.size() > 0)
{
iterator1 = moduleListNew.begin();
while (iterator1 != moduleListNew.end())
{
if (rva >= iterator1->second.firstThunk)
{
iterator1++;
if (iterator1 == moduleListNew.end())
{
iterator1--;
//new unknown module
if (iterator1->second.moduleName[0] == L'?')
{
module = &(iterator1->second);
}
else
{
addUnknownModuleToModuleList(apiNotFound->rva);
module = &(moduleListNew.find(rva)->second);
}
break;
}
else if (rva < iterator1->second.firstThunk)
{
iterator1--;
module = &(iterator1->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, MAX_PATH, L"?");
strcpy_s(import.name, MAX_PATH, "?");
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 iterator1;
if (moduleListNew.size() > 1)
{
iterator1 = moduleListNew.begin();
while (iterator1 != moduleListNew.end())
{
if (apiFound->rva >= iterator1->second.firstThunk)
{
iterator1++;
if (iterator1 == moduleListNew.end())
{
iterator1--;
module = &(iterator1->second);
break;
}
else if (apiFound->rva < iterator1->second.firstThunk)
{
iterator1--;
module = &(iterator1->second);
break;
}
}
else
{
#ifdef DEBUG_COMMENTS
Logger::debugLog(TEXT("Error iterator1 != moduleListNew.end()\r\n"));
#endif
break;
}
}
}
else
{
iterator1 = moduleListNew.begin();
module = &(iterator1->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, MAX_PATH, apiFound->moduleName);
strcpy_s(import.name, MAX_PATH, 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 iterator1;
ImportModuleThunk * moduleThunk;
iterator1 = moduleList.begin();
while (iterator1 != moduleList.end())
{
moduleThunk = &(iterator1->second);
moduleThunk->hTreeItem.Expand(flag);
iterator1++;
}
}
diff --git a/Scylla/ImportsHandling.h b/Scylla/ImportsHandling.h
index 3c4f8ab..e8fc5e6 100644
--- a/Scylla/ImportsHandling.h
+++ b/Scylla/ImportsHandling.h
@@ -1,75 +1,77 @@
#pragma once
#include <map>
// WTL
#include <atlbase.h>
#include <atlapp.h>
-#include <atlctrls.h> // CTreeViewCtrlEx, CTreeItem
+#include <atlctrls.h>
class ImportThunk;
class ImportModuleThunk;
+class CMultiSelectTreeViewCtrl;
+
class ImportsHandling
{
public:
std::map<DWORD_PTR, ImportModuleThunk> moduleList;
std::map<DWORD_PTR, ImportModuleThunk> moduleListNew;
//bool addFunction(WCHAR * moduleName, char * name, DWORD_PTR va, DWORD_PTR rva, DWORD_PTR ordinal, bool valid, bool suspect);
//bool addModule(WCHAR * moduleName, DWORD_PTR firstThunk);
- ImportsHandling(CTreeViewCtrlEx& TreeImports);
+ ImportsHandling(CMultiSelectTreeViewCtrl& TreeImports);
~ImportsHandling();
void displayAllImports();
void showImports(bool invalid, bool suspect);
bool invalidateFunction(CTreeItem selectedTreeNode);
bool cutThunk(CTreeItem selectedTreeNode);
bool deleteTreeNode(CTreeItem selectedTreeNode);
void updateImportInTreeView(const ImportThunk * importThunk, CTreeItem item);
void updateModuleInTreeView(const ImportModuleThunk * importThunk, CTreeItem item);
DWORD_PTR getApiAddressByNode(CTreeItem selectedTreeNode);
void scanAndFixModuleList();
void expandAllTreeNodes();
void collapseAllTreeNodes();
private:
DWORD numberOfFunctions;
WCHAR stringBuffer[600];
- CTreeViewCtrlEx& TreeImports;
+ 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
};
- CTreeItem addDllToTreeView(CTreeViewCtrlEx& idTreeView, const ImportModuleThunk * importThunk);
- CTreeItem addApiToTreeView(CTreeViewCtrlEx& idTreeView, CTreeItem parentDll, const ImportThunk * importThunk);
+ CTreeItem addDllToTreeView(CMultiSelectTreeViewCtrl& idTreeView, const ImportModuleThunk * importThunk);
+ CTreeItem addApiToTreeView(CMultiSelectTreeViewCtrl& idTreeView, CTreeItem parentDll, const ImportThunk * importThunk);
bool isItemSelected(CTreeItem hItem);
void unselectItem(CTreeItem htItem);
bool selectItem(CTreeItem hItem, bool select = true);
- void setFocus(CTreeViewCtrlEx& hwndTV, CTreeItem htItem);
+ void setFocus(CMultiSelectTreeViewCtrl& hwndTV, CTreeItem htItem);
bool findNewModules(std::map<DWORD_PTR, ImportThunk> & 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 c95b0a0..1906406 100644
--- a/Scylla/MainGui.cpp
+++ b/Scylla/MainGui.cpp
@@ -1,1224 +1,1225 @@
#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"
#include "WindowDeferrer.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);
if(hMenuImports)
{
appendPluginListToMenu(hMenuImports.GetSubMenu(0));
}
accelerators.LoadAccelerators(IDR_ACCELERATOR_MAIN);
}
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;
}
}
DoDataExchange(); // attach controls
CMessageLoop* pLoop = _Module.GetMessageLoop();
pLoop->AddMessageFilter(this);
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();
GetWindowRect(&minDlgSize);
return TRUE;
}
void MainGui::OnDestroy()
{
PostQuitMessage(0);
}
void MainGui::OnGetMinMaxInfo(MINMAXINFO* lpMMI)
{
lpMMI->ptMinTrackSize = CPoint(minDlgSize.Size());
}
void MainGui::OnSizing(UINT fwSide, const RECT* pRect)
{
// Get size difference
CRect rectOld;
GetWindowRect(&rectOld);
CRect rectNew = *pRect;
sizeOffset = rectNew.Size() - rectOld.Size();
}
void MainGui::OnSize(UINT nType, CSize size)
{
const WindowDeferrer::Deferrable controls[] =
{
{IDC_GROUP_ATTACH, false, false, true, false},
{IDC_CBO_PROCESSLIST, false, false, true, false},
{IDC_BTN_PICKDLL, true, false, false, false},
{IDC_GROUP_IMPORTS, false, false, true, true},
{IDC_TREE_IMPORTS, false, false, true, true},
{IDC_BTN_INVALIDIMPORTS, false, true, false, false},
{IDC_BTN_SUSPECTIMPORTS, false, true, false, false},
{IDC_BTN_SAVETREE, true, true, false, false},
{IDC_BTN_LOADTREE, true, true, false, false},
{IDC_BTN_CLEARIMPORTS, true, true, false, false},
{IDC_GROUP_IATINFO, false, true, false, false},
{IDC_STATIC_OEPADDRESS, false, true, false, false},
{IDC_STATIC_IATADDRESS, false, true, false, false},
{IDC_STATIC_IATSIZE, false, true, false, false},
{IDC_EDIT_OEPADDRESS, false, true, false, false},
{IDC_EDIT_IATADDRESS, false, true, false, false},
{IDC_EDIT_IATSIZE, false, true, false, false},
{IDC_BTN_IATAUTOSEARCH, false, true, false, false},
{IDC_BTN_GETIMPORTS, false, true, false, false},
{IDC_GROUP_ACTIONS, false, true, false, false},
{IDC_BTN_AUTOTRACE, false, true, false, false},
{IDC_GROUP_DUMP, false, true, false, false},
{IDC_BTN_DUMP, false, true, false, false},
{IDC_BTN_PEREBUILD, false, true, false, false},
{IDC_BTN_FIXDUMP, false, true, false, false},
{IDC_GROUP_LOG, false, true, true, false},
{IDC_LIST_LOG, false, true, true, false}
};
if(nType == SIZE_RESTORED)
{
WindowDeferrer deferrer(m_hWnd, controls, _countof(controls));
deferrer.defer(sizeOffset.cx, sizeOffset.cy);
sizeOffset.SetSize(0, 0);
}
}
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::OnTreeImportsOnKey(const NMHDR* pnmh)
{
const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh;
switch(tkd->wVKey)
{
case VK_RETURN:
{
- CTreeItem selected = TreeImports.GetSelectedItem();
+ CTreeItem selected = TreeImports.GetFirstSelectedItem();
if(!selected.IsNull() && !selected.GetParent().IsNull())
{
pickApiActionHandler(selected);
}
}
return 1;
case VK_DELETE:
{
- CTreeItem selected = TreeImports.GetSelectedItem();
- if(!selected.IsNull())
+ CTreeItem selected = TreeImports.GetFirstSelectedItem();
+ while(!selected.IsNull())
{
if(selected.GetParent().IsNull())
{
importsHandling.deleteTreeNode(selected);
}
else
{
importsHandling.cutThunk(selected);
}
+ selected = TreeImports.GetNextSelectedItem(selected);
}
}
return 1;
}
SetMsgHandled(FALSE);
return 0;
}
UINT MainGui::OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg)
{
//TreeImportsSubclass.ProcessWindowMessage();
//UINT original = 0;
if(lpMsg)
{
switch(lpMsg->wParam)
{
case VK_RETURN:
return DLGC_WANTMESSAGE;
}
}
SetMsgHandled(FALSE);
return 0;
}
void MainGui::OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags)
{
switch(nChar)
{
case VK_RETURN:
break;
default:
SetMsgHandled(FALSE);
break;
}
}
void MainGui::OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl)
{
fillProcessListComboBox(ComboProcessList);
}
void MainGui::OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
{
processSelectedActionHandler(ComboProcessList.GetCurSel());
}
void MainGui::OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl)
{
pickDllActionHandler();
}
void MainGui::OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl)
{
optionsActionHandler();
}
void MainGui::OnDump(UINT uNotifyCode, int nID, CWindow wndCtl)
{
dumpActionHandler();
}
void MainGui::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl)
{
dumpFixActionHandler();
}
void MainGui::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl)
{
peRebuildActionHandler();
}
void MainGui::OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl)
{
dllInjectActionHandler();
}
void MainGui::OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl)
{
iatAutosearchActionHandler();
}
void MainGui::OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl)
{
getImportsActionHandler();
}
void MainGui::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl)
{
showInvalidImportsActionHandler();
}
void MainGui::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl)
{
showSuspectImportsActionHandler();
}
void MainGui::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl)
{
clearImportsActionHandler();
}
void MainGui::OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
{
// TODO
}
void MainGui::OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
{
// TODO
}
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();
}
bool MainGui::showFileDialog(WCHAR * selectedFile, bool save, const WCHAR * defFileName, const WCHAR * filter, const WCHAR * defExtension, const WCHAR * directory)
{
OPENFILENAME ofn = {0};
// WTL doesn't support new explorer styles on Vista and up
// This is because it uses a custom hook, we could remove it or derive
// from CFileDialog but this solution is easier and allows more control anyway (e.g. initial dir)
if(defFileName)
{
wcscpy_s(selectedFile, MAX_PATH, defFileName);
}
else
{
selectedFile[0] = _T('\0');
}
ofn.lStructSize = sizeof(ofn);
ofn.hwndOwner = m_hWnd;
ofn.lpstrFilter = filter;
ofn.lpstrDefExt = defExtension; // only first 3 chars are used, no dots!
ofn.lpstrFile = selectedFile;
ofn.lpstrInitialDir = directory;
ofn.nMaxFile = MAX_PATH;
ofn.Flags = OFN_PATHMUSTEXIST | OFN_HIDEREADONLY;
/*
*OFN_EXPLORER is automatically used, it only has to be specified
*if using a custom hook
*OFN_LONGNAMES is automatically used by explorer-style dialogs
*/
if(save)
ofn.Flags |= OFN_OVERWRITEPROMPT;
else
ofn.Flags |= OFN_FILEMUSTEXIST;
if(save)
return 0 != GetSaveFileName(&ofn);
else
return 0 != GetOpenFileName(&ofn);
}
void MainGui::setIconAndDialogCaption()
{
SetIcon(hIcon, TRUE);
SetIcon(hIcon, FALSE);
SetWindowText(TEXT(APPNAME)TEXT(" ")TEXT(ARCHITECTURE)TEXT(" ")TEXT(APPVERSION));
}
void MainGui::pickDllActionHandler()
{
if(!selectedProcess)
return;
PickDllGui dlgPickDll(processAccessHelp.moduleList);
if(dlgPickDll.DoModal())
{
//get selected module
processAccessHelp.selectedModule = dlgPickDll.getSelectedModule();
Logger::printfDialog(TEXT("->>> Module %s selected."), processAccessHelp.selectedModule->getFilename());
Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),processAccessHelp.selectedModule->modBaseAddr,processAccessHelp.selectedModule->modBaseSize);
}
else
{
processAccessHelp.selectedModule = 0;
}
}
void MainGui::pickApiActionHandler(CTreeItem item)
{
CTreeItem parent = item.GetParent();
if(parent.IsNull())
return;
// TODO: new node when user picked an API from another DLL?
PickApiGui dlgPickApi(processAccessHelp.moduleList);
if(dlgPickApi.DoModal())
{
const ApiInfo* api = dlgPickApi.getSelectedApi();
if(api && api->module)
{
std::map<DWORD_PTR, ImportModuleThunk>::iterator iterator1;
std::map<DWORD_PTR, ImportThunk>::iterator iterator2;
iterator1 = importsHandling.moduleList.begin();
while(iterator1 != importsHandling.moduleList.end())
{
if(iterator1->second.hTreeItem == parent)
{
iterator2 = iterator1->second.thunkList.begin();
while(iterator2 != iterator1->second.thunkList.end())
{
if(iterator2->second.hTreeItem == item)
{
ImportThunk &imp = iterator2->second;
wcscpy_s(imp.moduleName, MAX_PATH, api->module->getFilename());
strcpy_s(imp.name, MAX_PATH, api->name);
imp.ordinal = api->ordinal;
//imp.apiAddressVA = api->va; //??
imp.hint = api->hint;
imp.valid = true;
imp.suspect = api->isForwarded;
importsHandling.updateImportInTreeView(&imp, item);
break;
}
iterator2++;
}
break;
}
iterator1++;
}
}
}
}
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."));
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);
}
void MainGui::fillProcessListComboBox(CComboBox& hCombo)
{
hCombo.ResetContent();
std::vector<Process>& processList = processLister.getProcessListSnapshot();
for (size_t i = 0; i < processList.size(); i++)
{
swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("0x%04X - %s - %s"),processList[i].PID,processList[i].filename,processList[i].fullPath);
hCombo.AddString(stringBuffer);
}
}
void MainGui::addTextToOutputLog(const WCHAR * text)
{
if (m_hWnd)
{
ListLog.SetCurSel(ListLog.AddString(text));
}
}
void MainGui::clearOutputLog()
{
if (m_hWnd)
{
ListLog.ResetContent();
}
}
bool MainGui::saveLogToFile(const WCHAR * file)
{
const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian
const WCHAR newLine[] = L"\r\n";
bool success = true;
HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
if(hFile != INVALID_HANDLE_VALUE)
{
ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM);
WCHAR * buffer = 0;
int bufsize = 0;
for(int i = 0; i < ListLog.GetCount(); i++)
{
int size = ListLog.GetTextLen(i);
size += _countof(newLine)-1;
if(size+1 > bufsize)
{
bufsize = size+1;
delete[] buffer;
try
{
buffer = new WCHAR[bufsize];
}
catch(std::bad_alloc&)
{
buffer = 0;
success = false;
break;
}
}
ListLog.GetText(i, buffer);
wcscat_s(buffer, bufsize, newLine);
ProcessAccessHelp::writeMemoryToFileEnd(hFile, size * sizeof(WCHAR), buffer);
}
delete[] buffer;
CloseHandle(hFile);
}
return success;
}
void MainGui::showInvalidImportsActionHandler()
{
importsHandling.showImports(true, false);
}
void MainGui::showSuspectImportsActionHandler()
{
importsHandling.showImports(false, true);
}
void MainGui::iatAutosearchActionHandler()
{
DWORD_PTR searchAddress = 0;
DWORD_PTR addressIAT = 0;
DWORD sizeIAT = 0;
IATSearch iatSearch;
if(!selectedProcess)
return;
if(EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 1)
{
searchAddress = stringToDwordPtr(stringBuffer);
if (searchAddress)
{
if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT))
{
Logger::printfDialog(TEXT("IAT found at VA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" RVA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d)"),addressIAT, addressIAT - processAccessHelp.targetImageBase,sizeIAT,sizeIAT);
swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),addressIAT);
EditIATAddress.SetWindowText(stringBuffer);
swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("%08X"),sizeIAT);
EditIATSize.SetWindowText(stringBuffer);
swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("IAT found! Start Address ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d) "),addressIAT,sizeIAT,sizeIAT);
MessageBox(stringBuffer, L"IAT found", MB_ICONINFORMATION);
}
else
{
Logger::printfDialog(TEXT("IAT not found at OEP ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("!"),searchAddress);
}
}
}
}
void MainGui::getImportsActionHandler()
{
DWORD_PTR addressIAT = 0;
DWORD sizeIAT = 0;
if(!selectedProcess)
return;
if (EditIATAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
{
addressIAT = stringToDwordPtr(stringBuffer);
}
if (EditIATSize.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
{
sizeIAT = wcstoul(stringBuffer, NULL, 16);
}
if (addressIAT && sizeIAT)
{
apiReader.readAndParseIAT(addressIAT, sizeIAT,importsHandling.moduleList);
importsHandling.displayAllImports();
}
}
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.GetSelectedItem();
+ over = TreeImports.GetFirstSelectedItem();
if(over)
{
over.EnsureVisible();
over.GetRect(&pos, TRUE);
TreeImports.ClientToScreen(&pos);
}
else
{
TreeImports.GetWindowRect(&pos);
}
pt = pos.TopLeft();
}
else
{
// Get item under cursor
CPoint client = pt;
TreeImports.ScreenToClient(&client);
UINT flags;
over = TreeImports.HitTest(client, &flags);
if(over && !(flags & TVHT_ONITEM))
{
over = NULL;
}
}
if(over)
{
parent = over.GetParent();
}
if (hMenuImports)
{
// Prepare hmenuImports
SetupImportsMenuItems(!over.IsNull(), !parent.IsNull());
CMenuHandle hSub = hMenuImports.GetSubMenu(0);
BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD, pt.x, pt.y, hwnd);
if (menuItem)
{
if ((menuItem >= PLUGIN_MENU_BASE_ID) && (menuItem <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID)))
{
//wsprintf(stringBuffer, L"%d %s\n",menuItem,pluginList[menuItem - PLUGIN_MENU_BASE_ID].pluginName);
//MessageBox(stringBuffer, L"plugin selection");
pluginActionHandler(menuItem);
return;
}
switch (menuItem)
{
case ID__INVALIDATEFUNCTION:
importsHandling.invalidateFunction(over);
break;
case ID__DISASSEMBLE:
startDisassemblerGui(over);
break;
case ID__EXPANDALLNODES:
importsHandling.expandAllTreeNodes();
break;
case ID__COLLAPSEALLNODES:
importsHandling.collapseAllTreeNodes();
break;
case ID__CUTTHUNK:
importsHandling.cutThunk(over);
break;
case ID__DELETETREENODE:
importsHandling.deleteTreeNode(parent ? parent : over);
break;
}
}
}
}
void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt)
{
if (hMenuLog)
{
if(pt.x == -1 && pt.y == -1) // invoked by keyboard
{
CRect pos;
ListLog.GetWindowRect(&pos);
pt = pos.TopLeft();
}
CMenuHandle hSub = hMenuLog.GetSubMenu(0);
BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD, pt.x, pt.y, hwnd);
if (menuItem)
{
switch (menuItem)
{
case ID__SAVE:
WCHAR selectedFilePath[MAX_PATH];
if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt"))
{
saveLogToFile(selectedFilePath);
}
break;
case ID__CLEAR:
clearOutputLog();
break;
}
}
}
}
void MainGui::appendPluginListToMenu(CMenuHandle hMenu)
{
std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList();
std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList();
if (scyllaPluginList.size() > 0)
{
CMenuHandle newMenu;
newMenu.CreatePopupMenu();
for (size_t i = 0; i < scyllaPluginList.size(); i++)
{
newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName);
}
hMenu.AppendMenu(MF_MENUBARBREAK);
hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins");
}
if (imprecPluginList.size() > 0)
{
CMenuHandle newMenu;
newMenu.CreatePopupMenu();
for (size_t i = 0; i < imprecPluginList.size(); i++)
{
newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName);
}
hMenu.AppendMenu(MF_MENUBARBREAK);
hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins");
}
}
void MainGui::dumpActionHandler()
{
if(!selectedProcess)
return;
WCHAR selectedFilePath[MAX_PATH];
const WCHAR * fileFilter;
const WCHAR * defExtension;
PeDump peDump;
if (processAccessHelp.selectedModule)
{
fileFilter = filterDll;
defExtension = L"dll";
}
else
{
fileFilter = filterExe;
defExtension = L"exe";
}
if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension))
{
if (processAccessHelp.selectedModule)
{
//dump DLL
peDump.imageBase = processAccessHelp.selectedModule->modBaseAddr;
peDump.sizeOfImage = processAccessHelp.selectedModule->modBaseSize;
//get it from gui
peDump.entryPoint = getOEPFromGui();
wcscpy_s(peDump.fullpath, MAX_PATH, processAccessHelp.selectedModule->fullPath);
}
else
{
peDump.imageBase = ProcessAccessHelp::targetImageBase;
peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage;
//get it from gui
peDump.entryPoint = getOEPFromGui();
wcscpy_s(peDump.fullpath, MAX_PATH, selectedProcess->fullPath);
}
peDump.useHeaderFromDisk = ConfigurationHolder::getConfigObject(USE_PE_HEADER_FROM_DISK)->isTrue();
if (peDump.dumpCompleteProcessToDisk(selectedFilePath))
{
Logger::printfDialog(TEXT("Dump success %s"),selectedFilePath);
//MessageBox(L"Image dumped successfully.", L"Success");
}
else
{
Logger::printfDialog(TEXT("Error: Cannot dump image."));
MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR);
}
}
}
DWORD_PTR MainGui::getOEPFromGui()
{
if (EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
{
return stringToDwordPtr(stringBuffer);
}
else
{
return 0;
}
}
void MainGui::peRebuildActionHandler()
{
DWORD newSize = 0;
WCHAR selectedFilePath[MAX_PATH];
PeRebuild peRebuild;
if(showFileDialog(selectedFilePath, false, NULL, filterExeDll))
{
if (ConfigurationHolder::getConfigObject(CREATE_BACKUP)->isTrue())
{
if (!ProcessAccessHelp::createBackupFile(selectedFilePath))
{
Logger::printfDialog(TEXT("Creating backup file failed %s"), selectedFilePath);
}
}
LONGLONG fileSize = ProcessAccessHelp::getFileSize(selectedFilePath);
LPVOID mapped = peRebuild.createFileMappingViewFull(selectedFilePath);
newSize = peRebuild.realignPE(mapped, (DWORD)fileSize);
peRebuild.closeAllMappingHandles();
if (newSize < 10)
{
Logger::printfDialog(TEXT("Rebuild failed %s"), selectedFilePath);
MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR);
}
else
{
peRebuild.truncateFile(selectedFilePath, newSize);
Logger::printfDialog(TEXT("Rebuild success %s"), selectedFilePath);
Logger::printfDialog(TEXT("-> Old file size 0x%08X new file size 0x%08X (%d %%)"), (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) );
//MessageBox(L"Image rebuilded successfully.", L"Success", MB_ICONINFORMATION);
}
}
}
void MainGui::dumpFixActionHandler()
{
if(!selectedProcess)
return;
if (TreeImports.GetCount() < 2)
{
Logger::printfDialog(TEXT("Nothing to rebuild"));
return;
}
WCHAR newFilePath[MAX_PATH];
WCHAR selectedFilePath[MAX_PATH];
const WCHAR * fileFilter;
ImportRebuild importRebuild;
if (processAccessHelp.selectedModule)
{
fileFilter = filterDll;
}
else
{
fileFilter = filterExe;
}
if (showFileDialog(selectedFilePath, false, NULL, fileFilter))
{
wcscpy_s(newFilePath,MAX_PATH,selectedFilePath);
const WCHAR * extension = 0;
WCHAR* dot = wcsrchr(newFilePath, L'.');
if (dot)
{
*dot = L'\0';
extension = selectedFilePath + (dot - newFilePath); //wcsrchr(selectedFilePath, L'.');
}
wcscat_s(newFilePath, MAX_PATH, L"_SCY");
if(extension)
{
wcscat_s(newFilePath, MAX_PATH, extension);
}
if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList))
{
//MessageBox(L"Imports rebuilding successful", L"Success", MB_ICONINFORMATION);
Logger::printfDialog(TEXT("Import Rebuild success %s"), newFilePath);
}
else
{
Logger::printfDialog(TEXT("Import Rebuild failed, target %s"), selectedFilePath);
MessageBox(L"Imports rebuilding failed", L"Failure", MB_ICONERROR);
}
}
}
void MainGui::enableDialogControls(BOOL value)
{
GetDlgItem(IDC_BTN_PICKDLL).EnableWindow(value);
GetDlgItem(IDC_BTN_DUMP).EnableWindow(value);
GetDlgItem(IDC_BTN_FIXDUMP).EnableWindow(value);
GetDlgItem(IDC_BTN_IATAUTOSEARCH).EnableWindow(value);
GetDlgItem(IDC_BTN_GETIMPORTS).EnableWindow(value);
GetDlgItem(IDC_BTN_SUSPECTIMPORTS).EnableWindow(value);
GetDlgItem(IDC_BTN_INVALIDIMPORTS).EnableWindow(value);
GetDlgItem(IDC_BTN_CLEARIMPORTS).EnableWindow(value);
CMenuHandle menu = GetMenu();
menu.EnableMenuItem(ID_FILE_DUMP, value ? MF_ENABLED : MF_GRAYED);
menu.EnableMenuItem(ID_FILE_FIXDUMP, value ? MF_ENABLED : MF_GRAYED);
menu.EnableMenuItem(ID_MISC_DLLINJECTION, value ? MF_ENABLED : MF_GRAYED);
menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetScylla, MF_BYPOSITION | (value ? MF_ENABLED : MF_GRAYED));
menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetImpRec, MF_BYPOSITION | (value ? MF_ENABLED : MF_GRAYED));
//not yet implemented
GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE);
GetDlgItem(IDC_BTN_SAVETREE).EnableWindow(FALSE);
GetDlgItem(IDC_BTN_LOADTREE).EnableWindow(FALSE);
menu.EnableMenuItem(ID_IMPORTS_INVALIDATESELECTED, MF_GRAYED);
menu.EnableMenuItem(ID_IMPORTS_CUTSELECTED, MF_GRAYED);
menu.EnableMenuItem(ID_IMPORTS_SAVETREE, MF_GRAYED);
menu.EnableMenuItem(ID_IMPORTS_SAVETREE, MF_GRAYED);
menu.EnableMenuItem(ID_IMPORTS_LOADTREE, MF_GRAYED);
menu.EnableMenuItem(ID_TRACE_AUTOTRACE, MF_GRAYED);
}
void MainGui::showAboutDialog()
{
AboutGui dlgAbout;
dlgAbout.DoModal();
}
void MainGui::dllInjectActionHandler()
{
if(!selectedProcess)
return;
WCHAR selectedFilePath[MAX_PATH];
HMODULE hMod = 0;
DllInjection dllInjection;
if (showFileDialog(selectedFilePath, false, NULL, filterDll))
{
hMod = dllInjection.dllInjection(ProcessAccessHelp::hProcess, selectedFilePath);
if (hMod && ConfigurationHolder::getConfigObject(DLL_INJECTION_AUTO_UNLOAD)->isTrue())
{
if (!dllInjection.unloadDllInProcess(ProcessAccessHelp::hProcess, hMod))
{
Logger::printfDialog(TEXT("DLL unloading failed, target %s"), selectedFilePath);
}
}
if (hMod)
{
Logger::printfDialog(TEXT("DLL Injection was successful, target %s"), selectedFilePath);
}
else
{
Logger::printfDialog(TEXT("DLL Injection failed, target %s"), selectedFilePath);
}
}
}
void MainGui::optionsActionHandler()
{
OptionsGui dlgOptions;
dlgOptions.DoModal();
}
void MainGui::clearImportsActionHandler()
{
TreeImports.DeleteAllItems();
importsHandling.moduleList.clear();
}
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();
}
diff --git a/Scylla/MainGui.h b/Scylla/MainGui.h
index 15e7f18..7da3192 100644
--- a/Scylla/MainGui.h
+++ b/Scylla/MainGui.h
@@ -1,246 +1,242 @@
#pragma once
#include <windows.h>
#include "resource.h"
// WTL
#include <atlbase.h> // base ATL classes
#include <atlapp.h> // base WTL classes
#include <atlwin.h> // ATL GUI classes
#include <atlmisc.h> // WTL utility classes like CString
#include <atlcrack.h> // WTL enhanced msg map macros
#include <atlctrls.h> // WTL controls
#include <atlddx.h> // WTL dialog data exchange
+#include "multitree.h"
//#define _CRTDBG_MAP_ALLOC
//#include <cstdlib>
//#include <crtdbg.h>
#include <cstdio>
#include "Logger.h"
#include "ProcessLister.h"
#include "IATSearch.h"
#include "PickDllGui.h"
#include "ImportsHandling.h"
-/*
-TODO: Multiselect treeview:
- - GetSelectedCount(), GetFirstSelected(), GetNextSelected(), GetAllSelected() -> vector, SetSelection(), AddSelection(), ClearSelection()
- - pass vector of items to cutThunk for 'Cut selected thunks' menu item
- - otherwise use GetFirstSelected
-*/
-
class MainGui : public CDialogImpl<MainGui>, public CWinDataExchange<MainGui>, public CMessageFilter
{
public:
enum { IDD = IDD_DLG_MAIN };
BEGIN_DDX_MAP(MainGui)
DDX_CONTROL(IDC_TREE_IMPORTS, TreeImportsSubclass) // subclass
- DDX_CONTROL_HANDLE(IDC_TREE_IMPORTS, TreeImports) // attach
- DDX_CONTROL_HANDLE(IDC_CBO_PROCESSLIST, ComboProcessList)
+ DDX_CONTROL(IDC_TREE_IMPORTS, TreeImports)
+ DDX_CONTROL_HANDLE(IDC_CBO_PROCESSLIST, ComboProcessList) // attach
DDX_CONTROL_HANDLE(IDC_LIST_LOG, ListLog)
DDX_CONTROL_HANDLE(IDC_EDIT_OEPADDRESS, EditOEPAddress)
DDX_CONTROL_HANDLE(IDC_EDIT_IATADDRESS, EditIATAddress)
DDX_CONTROL_HANDLE(IDC_EDIT_IATSIZE, EditIATSize)
END_DDX_MAP()
BEGIN_MSG_MAP_EX(MainGui)
MSG_WM_INITDIALOG(OnInitDialog)
MSG_WM_DESTROY(OnDestroy)
MSG_WM_GETMINMAXINFO(OnGetMinMaxInfo)
MSG_WM_SIZING(OnSizing)
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, OnTreeImportsOnKey)
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_NOTIFICATIONS()
+
ALT_MSG_MAP(IDC_TREE_IMPORTS) // message map for subclassed treeview
MSG_WM_GETDLGCODE(OnTreeImportsSubclassGetDlgCode)
MSG_WM_CHAR(OnTreeImportsSubclassChar)
END_MSG_MAP()
MainGui();
void addTextToOutputLog(const WCHAR * text);
protected:
// Variables
ProcessLister processLister;
WCHAR stringBuffer[600];
ImportsHandling importsHandling;
ProcessAccessHelp processAccessHelp;
ApiReader apiReader;
Process * selectedProcess;
// File selection stuff
static const WCHAR filterExe[];
static const WCHAR filterDll[];
static const WCHAR filterExeDll[];
static const WCHAR filterTxt[];
// Controls
- CTreeViewCtrlEx TreeImports;
+ CMultiSelectTreeViewCtrl TreeImports;
CComboBox ComboProcessList;
CEdit EditOEPAddress;
CEdit EditIATAddress;
CEdit EditIATSize;
CListBox ListLog;
CContainedWindow TreeImportsSubclass;
CRect minDlgSize;
CSize sizeOffset;
CAccelerator accelerators;
// Handles
CIcon hIcon;
CMenu hMenuImports;
CMenu hMenuLog;
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 OnGetMinMaxInfo(MINMAXINFO* lpMMI);
void OnSizing(UINT fwSide, const RECT* pRect);
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 OnTreeImportsOnKey(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 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 iatAutosearchActionHandler();
void getImportsActionHandler();
void dumpActionHandler();
void peRebuildActionHandler();
void startDisassemblerGui(CTreeItem selectedTreeNode);
void dumpFixActionHandler();
void showAboutDialog();
void dllInjectActionHandler();
void optionsActionHandler();
void clearImportsActionHandler();
void pluginActionHandler(int menuItem);
// Popup menu functions
void SetupImportsMenuItems(bool isItem, bool isThunk);
void appendPluginListToMenu(CMenuHandle hMenuTrackPopup);
void DisplayContextMenuImports(CWindow, CPoint);
void DisplayContextMenuLog(CWindow, CPoint);
// Log
void clearOutputLog();
bool saveLogToFile(const WCHAR * file);
// Misc
DWORD_PTR getOEPFromGui();
static DWORD_PTR stringToDwordPtr(const WCHAR * hexString);
};
diff --git a/Scylla/Scylla.vcxproj b/Scylla/Scylla.vcxproj
index f04e963..eb6a702 100644
--- a/Scylla/Scylla.vcxproj
+++ b/Scylla/Scylla.vcxproj
@@ -1,227 +1,228 @@
<?xml version="1.0" encoding="utf-8"?>
<Project DefaultTargets="Build" ToolsVersion="4.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
<ItemGroup Label="ProjectConfigurations">
<ProjectConfiguration Include="Debug|Win32">
<Configuration>Debug</Configuration>
<Platform>Win32</Platform>
</ProjectConfiguration>
<ProjectConfiguration Include="Debug|x64">
<Configuration>Debug</Configuration>
<Platform>x64</Platform>
</ProjectConfiguration>
<ProjectConfiguration Include="Release|Win32">
<Configuration>Release</Configuration>
<Platform>Win32</Platform>
</ProjectConfiguration>
<ProjectConfiguration Include="Release|x64">
<Configuration>Release</Configuration>
<Platform>x64</Platform>
</ProjectConfiguration>
</ItemGroup>
<PropertyGroup Label="Globals">
<ProjectGuid>{710434C9-FC4B-4F1D-B318-E10ADC78499F}</ProjectGuid>
<Keyword>Win32Proj</Keyword>
<RootNamespace>Scylla</RootNamespace>
</PropertyGroup>
<Import Project="$(VCTargetsPath)\Microsoft.Cpp.Default.props" />
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'" Label="Configuration">
<ConfigurationType>Application</ConfigurationType>
<UseDebugLibraries>true</UseDebugLibraries>
<CharacterSet>Unicode</CharacterSet>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'" Label="Configuration">
<ConfigurationType>Application</ConfigurationType>
<UseDebugLibraries>true</UseDebugLibraries>
<CharacterSet>Unicode</CharacterSet>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'" Label="Configuration">
<ConfigurationType>Application</ConfigurationType>
<UseDebugLibraries>false</UseDebugLibraries>
<WholeProgramOptimization>true</WholeProgramOptimization>
<CharacterSet>Unicode</CharacterSet>
<PlatformToolset>v100</PlatformToolset>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'" Label="Configuration">
<ConfigurationType>Application</ConfigurationType>
<UseDebugLibraries>false</UseDebugLibraries>
<WholeProgramOptimization>true</WholeProgramOptimization>
<CharacterSet>Unicode</CharacterSet>
</PropertyGroup>
<Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
<ImportGroup Label="ExtensionSettings">
</ImportGroup>
<ImportGroup Label="PropertySheets" Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
<Import Project="$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props" Condition="exists('$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props')" Label="LocalAppDataPlatform" />
</ImportGroup>
<ImportGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'" Label="PropertySheets">
<Import Project="$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props" Condition="exists('$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props')" Label="LocalAppDataPlatform" />
</ImportGroup>
<ImportGroup Label="PropertySheets" Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
<Import Project="$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props" Condition="exists('$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props')" Label="LocalAppDataPlatform" />
</ImportGroup>
<ImportGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'" Label="PropertySheets">
<Import Project="$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props" Condition="exists('$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props')" Label="LocalAppDataPlatform" />
</ImportGroup>
<PropertyGroup Label="UserMacros" />
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
<LinkIncremental>true</LinkIncremental>
<OutDir>$(SolutionDir)$(Platform)\$(Configuration)\</OutDir>
<IntDir>$(Platform)\$(Configuration)\</IntDir>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
<LinkIncremental>true</LinkIncremental>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
<LinkIncremental>false</LinkIncremental>
<OutDir>$(SolutionDir)$(Platform)\$(Configuration)\</OutDir>
<IntDir>$(Platform)\$(Configuration)\</IntDir>
</PropertyGroup>
<PropertyGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
<LinkIncremental>false</LinkIncremental>
</PropertyGroup>
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
<ClCompile>
<PrecompiledHeader>
</PrecompiledHeader>
<WarningLevel>Level3</WarningLevel>
<Optimization>Disabled</Optimization>
<PreprocessorDefinitions>WIN32;_DEBUG;_WINDOWS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<AdditionalIncludeDirectories>$(SolutionDir)diStorm\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
</ClCompile>
<Link>
<SubSystem>Windows</SubSystem>
<GenerateDebugInformation>true</GenerateDebugInformation>
<AdditionalDependencies>$(SolutionDir)$(Platform)\$(Configuration)\diStorm.lib;%(AdditionalDependencies)</AdditionalDependencies>
<AdditionalManifestDependencies>type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*' publicKeyToken='6595b64144ccf1df' </AdditionalManifestDependencies>
</Link>
</ItemDefinitionGroup>
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
<ClCompile>
<PrecompiledHeader>
</PrecompiledHeader>
<WarningLevel>Level3</WarningLevel>
<Optimization>Disabled</Optimization>
<PreprocessorDefinitions>WIN32;_DEBUG;_WINDOWS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<AdditionalIncludeDirectories>$(SolutionDir)diStorm\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
</ClCompile>
<Link>
<SubSystem>Windows</SubSystem>
<GenerateDebugInformation>true</GenerateDebugInformation>
<AdditionalDependencies>$(SolutionDir)$(Platform)\$(Configuration)\diStorm.lib;%(AdditionalDependencies)</AdditionalDependencies>
<AdditionalManifestDependencies>type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*' publicKeyToken='6595b64144ccf1df' </AdditionalManifestDependencies>
</Link>
</ItemDefinitionGroup>
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
<ClCompile>
<WarningLevel>Level3</WarningLevel>
<PrecompiledHeader>
</PrecompiledHeader>
<Optimization>MaxSpeed</Optimization>
<FunctionLevelLinking>true</FunctionLevelLinking>
<IntrinsicFunctions>true</IntrinsicFunctions>
<PreprocessorDefinitions>WIN32;NDEBUG;_WINDOWS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<RuntimeLibrary>MultiThreaded</RuntimeLibrary>
<AdditionalIncludeDirectories>$(SolutionDir)diStorm\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
<MinimalRebuild>true</MinimalRebuild>
</ClCompile>
<Link>
<SubSystem>Windows</SubSystem>
<GenerateDebugInformation>false</GenerateDebugInformation>
<EnableCOMDATFolding>true</EnableCOMDATFolding>
<OptimizeReferences>true</OptimizeReferences>
<AdditionalDependencies>$(SolutionDir)$(Platform)\$(Configuration)\diStorm.lib;%(AdditionalDependencies)</AdditionalDependencies>
<AdditionalManifestDependencies>type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*' publicKeyToken='6595b64144ccf1df' </AdditionalManifestDependencies>
</Link>
</ItemDefinitionGroup>
<ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
<ClCompile>
<WarningLevel>Level3</WarningLevel>
<PrecompiledHeader>
</PrecompiledHeader>
<Optimization>MaxSpeed</Optimization>
<FunctionLevelLinking>true</FunctionLevelLinking>
<IntrinsicFunctions>true</IntrinsicFunctions>
<PreprocessorDefinitions>WIN32;NDEBUG;_WINDOWS;%(PreprocessorDefinitions)</PreprocessorDefinitions>
<RuntimeLibrary>MultiThreaded</RuntimeLibrary>
<AdditionalIncludeDirectories>$(SolutionDir)diStorm\include;%(AdditionalIncludeDirectories)</AdditionalIncludeDirectories>
<MinimalRebuild>true</MinimalRebuild>
</ClCompile>
<Link>
<SubSystem>Windows</SubSystem>
<GenerateDebugInformation>false</GenerateDebugInformation>
<EnableCOMDATFolding>true</EnableCOMDATFolding>
<OptimizeReferences>true</OptimizeReferences>
<AdditionalDependencies>$(SolutionDir)$(Platform)\$(Configuration)\diStorm.lib;%(AdditionalDependencies)</AdditionalDependencies>
<AdditionalManifestDependencies>type='win32' name='Microsoft.Windows.Common-Controls' version='6.0.0.0' processorArchitecture='*' publicKeyToken='6595b64144ccf1df' </AdditionalManifestDependencies>
</Link>
</ItemDefinitionGroup>
<ItemGroup>
<ClCompile Include="AboutGui.cpp" />
<ClCompile Include="ApiReader.cpp" />
<ClCompile Include="ConfigurationHolder.cpp" />
<ClCompile Include="DisassemblerGui.cpp" />
<ClCompile Include="DllInjection.cpp" />
<ClCompile Include="DllInjectionPlugin.cpp" />
<ClCompile Include="IATSearch.cpp" />
<ClCompile Include="ImportRebuild.cpp" />
<ClCompile Include="ImportsHandling.cpp" />
<ClCompile Include="Logger.cpp" />
<ClCompile Include="main.cpp" />
<ClCompile Include="MainGui.cpp" />
<ClCompile Include="NativeWinApi.cpp" />
<ClCompile Include="OptionsGui.cpp" />
<ClCompile Include="PeDump.cpp" />
<ClCompile Include="PeRebuild.cpp" />
<ClCompile Include="PickApiGui.cpp" />
<ClCompile Include="PickDllGui.cpp" />
<ClCompile Include="PluginLoader.cpp" />
<ClCompile Include="ProcessAccessHelp.cpp" />
<ClCompile Include="ProcessLister.cpp" />
<ClCompile Include="SystemInformation.cpp" />
<ClCompile Include="WindowDeferrer.cpp" />
</ItemGroup>
<ItemGroup>
<ClInclude Include="AboutGui.h" />
<ClInclude Include="ApiReader.h" />
<ClInclude Include="ConfigurationHolder.h" />
<ClInclude Include="definitions.h" />
<ClInclude Include="DisassemblerGui.h" />
<ClInclude Include="DllInjection.h" />
<ClInclude Include="DllInjectionPlugin.h" />
<ClInclude Include="IATSearch.h" />
<ClInclude Include="ImportRebuild.h" />
<ClInclude Include="ImportsHandling.h" />
<ClInclude Include="Logger.h" />
<ClInclude Include="MainGui.h" />
+ <ClInclude Include="multitree.h" />
<ClInclude Include="NativeWinApi.h" />
<ClInclude Include="OptionsGui.h" />
<ClInclude Include="PeDump.h" />
<ClInclude Include="PeRebuild.h" />
<ClInclude Include="PickApiGui.h" />
<ClInclude Include="PickDllGui.h" />
<ClInclude Include="PluginLoader.h" />
<ClInclude Include="ProcessAccessHelp.h" />
<ClInclude Include="ProcessLister.h" />
<ClInclude Include="resource.h" />
<ClInclude Include="SystemInformation.h" />
<ClInclude Include="Thunks.h" />
<ClInclude Include="WindowDeferrer.h" />
</ItemGroup>
<ItemGroup>
<ResourceCompile Include="MainGui.rc" />
</ItemGroup>
<ItemGroup>
<None Include="check.ico" />
<None Include="error.ico" />
<None Include="scylla.ico" />
<None Include="warning.ico" />
</ItemGroup>
<Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
<ImportGroup Label="ExtensionTargets">
</ImportGroup>
<ProjectExtensions>
<VisualStudio>
<UserProperties RESOURCE_FILE="MainGui.rc" />
</VisualStudio>
</ProjectExtensions>
</Project>
\ No newline at end of file
diff --git a/Scylla/Scylla.vcxproj.filters b/Scylla/Scylla.vcxproj.filters
index 5a2eb3e..2ab74cf 100644
--- a/Scylla/Scylla.vcxproj.filters
+++ b/Scylla/Scylla.vcxproj.filters
@@ -1,190 +1,193 @@
<?xml version="1.0" encoding="utf-8"?>
<Project ToolsVersion="4.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
<ItemGroup>
<Filter Include="Quelldateien">
<UniqueIdentifier>{4FC737F1-C7A5-4376-A066-2A32D752A2FF}</UniqueIdentifier>
<Extensions>cpp;c;cc;cxx;def;odl;idl;hpj;bat;asm;asmx</Extensions>
</Filter>
<Filter Include="Headerdateien">
<UniqueIdentifier>{93995380-89BD-4b04-88EB-625FBE52EBFB}</UniqueIdentifier>
<Extensions>h;hpp;hxx;hm;inl;inc;xsd</Extensions>
</Filter>
<Filter Include="Ressourcendateien">
<UniqueIdentifier>{67DA6AB6-F800-4c08-8B7A-83BB121AAD01}</UniqueIdentifier>
<Extensions>rc;ico;cur;bmp;dlg;rc2;rct;bin;rgs;gif;jpg;jpeg;jpe;resx;tiff;tif;png;wav;mfcribbon-ms</Extensions>
</Filter>
<Filter Include="Headerdateien\GUI">
<UniqueIdentifier>{e037d0d5-35ad-4034-83db-746a56a4fee7}</UniqueIdentifier>
</Filter>
<Filter Include="Quelldateien\GUI">
<UniqueIdentifier>{6f76186f-b79c-41e2-8939-05d9de028aad}</UniqueIdentifier>
</Filter>
</ItemGroup>
<ItemGroup>
<ClCompile Include="main.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="MainGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="Logger.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="ProcessLister.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="PickDllGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="ImportsHandling.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="ApiReader.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="IATSearch.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="ProcessAccessHelp.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="ImportRebuild.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="PeDump.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="PeRebuild.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="DllInjection.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="DllInjectionPlugin.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="PluginLoader.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="ConfigurationHolder.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="DisassemblerGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="NativeWinApi.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="SystemInformation.cpp">
<Filter>Quelldateien</Filter>
</ClCompile>
<ClCompile Include="AboutGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="OptionsGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="WindowDeferrer.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
<ClCompile Include="PickApiGui.cpp">
<Filter>Quelldateien\GUI</Filter>
</ClCompile>
</ItemGroup>
<ItemGroup>
<ClInclude Include="resource.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="MainGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="Logger.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="ProcessLister.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="PickDllGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="ImportsHandling.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="ApiReader.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="IATSearch.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="ProcessAccessHelp.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="Thunks.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="PeRebuild.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="PeDump.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="ImportRebuild.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="ConfigurationHolder.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="DllInjectionPlugin.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="DllInjection.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="PluginLoader.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="definitions.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="DisassemblerGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="NativeWinApi.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="SystemInformation.h">
<Filter>Headerdateien</Filter>
</ClInclude>
<ClInclude Include="AboutGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="OptionsGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="WindowDeferrer.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
<ClInclude Include="PickApiGui.h">
<Filter>Headerdateien\GUI</Filter>
</ClInclude>
+ <ClInclude Include="multitree.h">
+ <Filter>Headerdateien\GUI</Filter>
+ </ClInclude>
</ItemGroup>
<ItemGroup>
<ResourceCompile Include="MainGui.rc">
<Filter>Ressourcendateien</Filter>
</ResourceCompile>
</ItemGroup>
<ItemGroup>
<None Include="scylla.ico">
<Filter>Ressourcendateien</Filter>
</None>
<None Include="check.ico">
<Filter>Ressourcendateien</Filter>
</None>
<None Include="error.ico">
<Filter>Ressourcendateien</Filter>
</None>
<None Include="warning.ico">
<Filter>Ressourcendateien</Filter>
</None>
</ItemGroup>
</Project>
\ No newline at end of file
diff --git a/Scylla/multitree.h b/Scylla/multitree.h
new file mode 100644
index 0000000..2cd5df5
--- /dev/null
+++ b/Scylla/multitree.h
@@ -0,0 +1,580 @@
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////
+// MultiSelectTree - A tree control with multi-select capabilities
+//
+// Written by Bjarke Viksoe ([email protected])
+// Copyright (c) 2005 Bjarke Viksoe.
+//
+// Add the following macro to the parent's message map:
+// REFLECT_NOTIFICATIONS()
+//
+// This code may be used in compiled form in any way you desire. This
+// source file may be redistributed by any means PROVIDING it is
+// not sold for profit without the authors written consent, and
+// providing that this notice and the authors name is included.
+//
+// This file is provided "as is" with no expressed or implied warranty.
+// The author accepts no liability if it causes any damage to you or your
+// computer whatsoever. It's free, so don't hassle me about it.
+//
+// Beware of bugs.
+//
+
+#ifndef __cplusplus
+ #error WTL requires C++ compilation (use a .cpp suffix)
+#endif
+
+#ifndef __ATLMISC_H__
+ #error multitree.h requires atlmisc.h to be included first
+#endif
+
+#ifndef __ATLCRACK_H__
+ #error multitree.h requires atlcrack.h to be included first
+#endif
+
+#ifndef __ATLCTRLS_H__
+ #error multitree.h requires atlctrls.h to be included first
+#endif
+
+// Extended MultiSelectTree styles
+static const DWORD MTVS_EX_NOMARQUEE = 0x00000001;
+
+// New control notifications
+static const UINT TVN_ITEMSELECTING = 0x0001;
+static const UINT TVN_ITEMSELECTED = 0x0002;
+
+static bool operator==(CTreeItem ti1, CTreeItem ti2)
+{
+ return ti1.m_hTreeItem == ti2.m_hTreeItem && ti1.m_pTreeView == ti2.m_pTreeView;
+}
+
+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)
+ {
+ }
+
+ // Operations
+
+ BOOL SubclassWindow(HWND hWnd)
+ {
+ ATLASSERT(m_hWnd == NULL);
+ ATLASSERT(::IsWindow(hWnd));
+ BOOL bRet = CWindowImpl< T, TBase, TWinTraits >::SubclassWindow(hWnd);
+ if( bRet )
+ _Init();
+ return bRet;
+ }
+
+ DWORD SetExtendedTreeStyle(DWORD dwStyle)
+ {
+ ATLASSERT(!m_ctrlTree.IsWindow()); // Before control is created, please!
+ DWORD dwOldStyle = m_dwTreeStyle;
+ m_dwTreeStyle = dwStyle;
+ return dwOldStyle;
+ }
+
+ void SelectItem(HTREEITEM hItem, BOOL bSelect)
+ {
+ ATLASSERT(::IsWindow(m_hWnd));
+ _SelectItem(hItem, bSelect == TRUE);
+ if( bSelect )
+ TBase::SelectItem(hItem);
+ }
+
+ void SelectAllItems(BOOL bSelect)
+ {
+ ATLASSERT(::IsWindow(m_hWnd));
+
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ _SelectItem(i, bSelect);
+ }
+ }
+
+ CTreeItem GetSelectedItem() const
+ {
+ ATLASSERT(false); // Not usable!
+ return GetFirstSelectedItem();
+ }
+
+ CTreeItem GetFocusItem() const
+ {
+ return TBase::GetSelectedItem();
+ }
+
+ UINT GetItemState(HTREEITEM hItem, UINT nStateMask) const
+ {
+ UINT nRes = TBase::GetItemState(hItem, nStateMask);
+ if( (nStateMask & TVIS_SELECTED) != 0 )
+ {
+ int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
+ if( iIndex >= 0 )
+ {
+ nRes &= ~TVIS_SELECTED;
+ if( m_aData.GetValueAt(iIndex) )
+ nRes |= TVIS_SELECTED;
+ }
+ }
+ return nRes;
+ }
+
+ CTreeItem GetFirstSelectedItem() const
+ {
+ if( m_aData.GetSize() == 0 )
+ return NULL;
+
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ if( m_aData.GetValueAt(i) )
+ return m_aData.GetKeyAt(i);
+ }
+ return NULL;
+ }
+
+ CTreeItem GetNextSelectedItem(HTREEITEM hItem) const
+ {
+ int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
+ if( iIndex < 0 )
+ return NULL;
+
+ for( int i = iIndex + 1; i < m_aData.GetSize(); i++ )
+ {
+ if( m_aData.GetValueAt(i) )
+ return m_aData.GetKeyAt(i);
+ }
+ return NULL;
+ }
+
+ int GetSelectedCount() const
+ {
+ int nCount = 0;
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ if( m_aData.GetValueAt(i) )
+ nCount++;
+ }
+ return nCount;
+ }
+
+ // Implementation
+
+ void _Init()
+ {
+ ATLASSERT(::IsWindow(m_hWnd));
+
+ ModifyStyle(TVS_SHOWSELALWAYS, 0);
+ }
+
+ void _SelectItem(int iIndex, bool bSelect, int action = TVC_UNKNOWN)
+ {
+ if( iIndex < 0 )
+ return;
+ bool bSelected = m_aData.GetValueAt(iIndex);
+ // Don't change if state is already updated (avoids flicker)
+ if( bSelected == bSelect )
+ return;
+ CTreeItem cItem = m_aData.GetKeyAt(iIndex);
+ CWindow parent = GetParent();
+ // Send notifications
+ NMTREEVIEW nmtv = { 0 };
+ nmtv.hdr.code = TVN_ITEMSELECTING;
+ nmtv.hdr.hwndFrom = m_hWnd;
+ nmtv.hdr.idFrom = GetDlgCtrlID();
+ nmtv.action = action;
+ nmtv.itemNew.hItem = cItem;
+ nmtv.itemNew.lParam = GetItemData(cItem);
+ nmtv.itemNew.state = bSelect ? TVIS_SELECTED : 0;
+ nmtv.itemNew.stateMask = TVIS_SELECTED;
+ if( parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv) != 0 )
+ return;
+ // Change state
+ m_aData.SetAtIndex(iIndex, cItem, bSelect);
+ // Repaint item
+ CRect rcItem;
+ if( GetItemRect(cItem, &rcItem, FALSE) )
+ InvalidateRect(&rcItem, TRUE);
+ // More notifications
+ nmtv.hdr.code = TVN_ITEMSELECTED;
+ parent.SendMessage(WM_NOTIFY, nmtv.hdr.idFrom, (LPARAM) &nmtv);
+ }
+
+ void _SelectItem(HTREEITEM hItem, bool bSelect, int action = TVC_UNKNOWN)
+ {
+ _SelectItem(m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this)), bSelect, action);
+ }
+
+ 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;
+ }
+
+ void _SelectBox(CRect rc)
+ {
+ CTreeItem hItem = GetFirstVisibleItem();
+ while( hItem != NULL )
+ {
+ CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this);
+ int i = m_aData.FindKey(cItem);
+ 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();
+ CBrush brush = CDCHandle::GetHalftoneBrush();
+ if( !brush.IsNull() )
+ {
+ CBrushHandle hOldBrush = dc.SelectBrush(brush);
+ dc.PatBlt(rect.left, rect.top, rect.Width(), szFrame.cy, PATINVERT);
+ dc.PatBlt(rect.left, rect.bottom - szFrame.cy, rect.Width(), szFrame.cy, PATINVERT);
+ dc.PatBlt(rect.left, rect.top + szFrame.cy, szFrame.cx, rect.Height() - (szFrame.cy * 2), PATINVERT);
+ dc.PatBlt(rect.right - szFrame.cx, rect.top + szFrame.cy, szFrame.cx, rect.Height() - (szFrame.cy * 2), PATINVERT);
+ dc.SelectBrush(hOldBrush);
+ }
+ }
+
+ // Message map and handlers
+
+ BEGIN_MSG_MAP_EX(CMultiSelectTreeViewImpl)
+ MSG_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()
+
+ int OnCreate(CREATESTRUCT* lpCreateStruct)
+ {
+ 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( ::GetAsyncKeyState(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( ::GetAsyncKeyState(VK_SHIFT) < 0 )
+ {
+ switch( nChar )
+ {
+ case VK_UP:
+ case VK_DOWN:
+ case VK_HOME:
+ case VK_END:
+ case VK_NEXT:
+ case VK_PRIOR:
+ HTREEITEM hItem = GetFocusItem();
+ // Is current or first-shift-item the upper item?
+ CRect rcItem1, rcItem2;
+ GetItemRect(m_hExtSelStart, &rcItem1, TRUE);
+ GetItemRect(hItem, &rcItem2, TRUE);
+ // Select from current item to item where SHIFT was pressed
+ if( rcItem1.top > rcItem2.top )
+ _SelectTree(hItem, m_hExtSelStart, TVC_BYKEYBOARD);
+ else
+ _SelectTree(m_hExtSelStart, hItem, TVC_BYKEYBOARD);
+ _SelectItem(hItem, true, TVC_BYKEYBOARD);
+ }
+ }
+ SetMsgHandled(FALSE);
+ }
+
+ void OnChar(UINT nChar, UINT nRepCnt, UINT nFlags)
+ {
+ if( nChar == VK_SPACE )
+ {
+ HTREEITEM hItem = GetFocusItem();
+ _SelectItem(hItem, (GetItemState(hItem, TVIS_SELECTED) & TVIS_SELECTED) == 0, TVC_BYKEYBOARD);
+ return;
+ }
+ SetMsgHandled(FALSE);
+ }
+
+ void OnSetFocus(CWindow wndOld)
+ {
+ DefWindowProc();
+ // FIX: We really need the focus-rectangle in this control since it
+ // improves the navigation a lot. So let's ask Windows to display it.
+ SendMessage(WM_UPDATEUISTATE, MAKEWPARAM(UIS_CLEAR, UISF_HIDEFOCUS));
+ }
+
+ void OnLButtonDown(UINT nFlags, CPoint point)
+ {
+ SetMsgHandled(FALSE);
+
+ // Hit-test and figure out where we're clicking...
+ TVHITTESTINFO hti = { 0 };
+ hti.pt = point;
+ HTREEITEM hItem = HitTest(&hti);
+ if( (hItem == NULL || (hti.flags & TVHT_ONITEMRIGHT) != 0) )
+ {
+ if( (m_dwExStyle & MTVS_EX_NOMARQUEE) == 0 && ::DragDetect(m_hWnd, point) )
+ {
+ // Great we're dragging a rubber-band
+ // Clear selection of CTRL is not down
+ if( ::GetAsyncKeyState(VK_CONTROL) >= 0 )
+ {
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ _SelectItem(i, false, TVC_BYMOUSE);
+ }
+ UpdateWindow();
+ }
+ // Now start drawing the rubber-band...
+ SetCapture();
+ m_ptDragStart = m_ptDragOld = point;
+ _DrawDragRect(point);
+ m_bMarquee = true;
+ SetMsgHandled(TRUE);
+ return;
+ }
+ }
+
+ if( hItem == NULL )
+ return;
+
+ if( (hti.flags & TVHT_ONITEMBUTTON) != 0 )
+ return;
+
+ // Great, let's do an advanced selection
+ if( (hti.flags & TVHT_ONITEMRIGHT) != 0 )
+ {
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ _SelectItem(i, false, TVC_BYMOUSE);
+ }
+ return;
+ }
+ int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
+ 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( ::GetAsyncKeyState(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( ::GetAsyncKeyState(VK_CONTROL) < 0 )
+ {
+ // Just toggle item
+ _SelectItem(iIndex, !bSelected, TVC_BYMOUSE);
+ }
+ else
+ {
+ // Remove current selection and replace it with clicked item
+ for( int i = 0; i < m_aData.GetSize(); i++ )
+ {
+ _SelectItem(i, i == iIndex, TVC_BYMOUSE);
+ }
+ }
+ }
+
+ void OnLButtonUp(UINT nFlags, CPoint point)
+ {
+ if( m_bMarquee )
+ ReleaseCapture();
+ SetMsgHandled(FALSE);
+ }
+
+ void OnMouseMove(UINT nFlags, CPoint point)
+ {
+ if( m_bMarquee )
+ {
+ CRect rc(m_ptDragStart, point);
+ _DrawDragRect(m_ptDragOld);
+ rc.NormalizeRect();
+ _SelectBox(rc);
+ UpdateWindow();
+ _DrawDragRect(point);
+ m_ptDragOld = point;
+ }
+ SetMsgHandled(FALSE);
+ }
+
+ void OnCaptureChanged(CWindow wnd)
+ {
+ if( m_bMarquee )
+ {
+ _DrawDragRect(m_ptDragOld);
+ m_bMarquee = false;
+ }
+ SetMsgHandled(FALSE);
+ }
+
+ LRESULT OnInsertItem(UINT uMsg, WPARAM wParam, LPARAM lParam)
+ {
+ HTREEITEM hItem = (HTREEITEM) DefWindowProc(uMsg, wParam, lParam);
+ if( hItem == NULL )
+ return (LRESULT) hItem;
+ CTreeItem cItem(hItem, (CTreeViewCtrlEx*)this);
+ // We manage a bit of extra information for each item. We'll store
+ // this in an ATL::CSimpleMap. Not a particular speedy structure for lookups.
+ // Don't keep too many items around in the tree!
+ m_aData.Add(cItem, false);
+ return (LRESULT) hItem;
+ }
+
+ LRESULT OnDeleteItem(NMHDR* pnmh)
+ {
+ const NMTREEVIEW* lpNMTV = (NMTREEVIEW*) pnmh;
+ m_aData.Remove(CTreeItem(lpNMTV->itemNew.hItem, (CTreeViewCtrlEx*)this));
+ return 0;
+ }
+
+ // Custom Draw
+
+ DWORD OnPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* /*lpNMCustomDraw*/)
+ {
+ return CDRF_NOTIFYITEMDRAW; // We need per-item notifications
+ }
+
+ DWORD OnItemPrePaint(int /*idCtrl*/, NMCUSTOMDRAW* lpNMCustomDraw)
+ {
+ NMTVCUSTOMDRAW* lpTVCD = (NMTVCUSTOMDRAW*) lpNMCustomDraw;
+ HTREEITEM hItem = (HTREEITEM) lpTVCD->nmcd.dwItemSpec;
+ int iIndex = m_aData.FindKey(CTreeItem(hItem, (CTreeViewCtrlEx*)this));
+ 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

Mime Type
text/x-diff
Expires
Tue, Jan 7, 12:51 AM (2 h, 53 m)
Storage Engine
local-disk
Storage Format
Raw Data
Storage Handle
b3/7c/0ca186bad138d36f077b1e29d5cd

Event Timeline