diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp index dd04beb..f664bc1 100644 --- a/Scylla/MainGui.cpp +++ b/Scylla/MainGui.cpp @@ -1,1197 +1,1197 @@ #include "MainGui.h" #include // WTL common dialogs #include "definitions.h" #include "PluginLoader.h" #include "ConfigurationHolder.h" #include "PeDump.h" #include "PeRebuild.h" #include "DllInjectionPlugin.h" #include "DisassemblerGui.h" #include "PickApiGui.h" #include "NativeWinApi.h" #include "ImportRebuild.h" #include "SystemInformation.h" #include "AboutGui.h" #include "OptionsGui.h" 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); hIconWarning.LoadIcon(IDI_ICON_WARNING); hIconError.LoadIcon(IDI_ICON_ERROR); 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(); DoDataExchange(); // attach controls DlgResize_Init(true, true); EditOEPAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); EditIATAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); EditIATSize.LimitText(MAX_HEX_VALUE_EDIT_LENGTH); appendPluginListToMenu(CMenuHandle(GetMenu()).GetSubMenu(MenuImportsOffsetTrace)); enableDialogControls(FALSE); setIconAndDialogCaption(); return TRUE; } void MainGui::OnDestroy() { PostQuitMessage(0); } void MainGui::OnSize(UINT nType, CSize size) { StatusBar.SendMessage(WM_SIZE); SetMsgHandled(FALSE); } void MainGui::OnLButtonDown(UINT nFlags, CPoint point) { SetMsgHandled(FALSE); } void MainGui::OnContextMenu(CWindow wnd, CPoint point) { // point = -1, -1 for keyboard invoked shortcut! switch(wnd.GetDlgCtrlID()) { case IDC_TREE_IMPORTS: DisplayContextMenuImports(wnd, point); return; case IDC_LIST_LOG: DisplayContextMenuLog(wnd, point); return; //default: // wnd == m_hWnd? // DisplayContextMenu(wnd, point); // return; } SetMsgHandled(FALSE); } void MainGui::OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl) { // Make sure it's a menu if(uNotifyCode == 0 && !wndCtl.IsWindow()) { if ((nID >= PLUGIN_MENU_BASE_ID) && (nID <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID))) { pluginActionHandler(nID); return; } } SetMsgHandled(FALSE); } LRESULT MainGui::OnTreeImportsClick(const NMHDR* pnmh) { SetMsgHandled(FALSE); return 0; } LRESULT MainGui::OnTreeImportsDoubleClick(const NMHDR* pnmh) { if(TreeImports.GetCount() < 1) return 0; // Get item under cursor CPoint client = GetMessagePos(); TreeImports.ScreenToClient(&client); UINT flags; CTreeItem over = TreeImports.HitTest(client, &flags); CTreeItem parent; if(over) { if(!(flags & TVHT_ONITEM)) { over = NULL; } else { parent = over.GetParent(); } } if(!over.IsNull() && !parent.IsNull()) { pickApiActionHandler(over); } return 0; } LRESULT MainGui::OnTreeImportsRightClick(const NMHDR* pnmh) { SetMsgHandled(FALSE); return 0; } LRESULT MainGui::OnTreeImportsRightDoubleClick(const NMHDR* pnmh) { SetMsgHandled(FALSE); return 0; } LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh) { const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh; switch(tkd->wVKey) { case VK_RETURN: { CTreeItem selected = TreeImports.GetFirstSelectedItem(); if(!selected.IsNull() && !selected.GetParent().IsNull()) { pickApiActionHandler(selected); } } return 1; case VK_DELETE: { CTreeItem selected = TreeImports.GetFirstSelectedItem(); while(!selected.IsNull()) { if(selected.GetParent().IsNull()) { importsHandling.deleteTreeNode(selected); } else { importsHandling.cutThunk(selected); } selected = TreeImports.GetNextSelectedItem(selected); } } return 1; } SetMsgHandled(FALSE); return 0; } UINT MainGui::OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg) { if(lpMsg) { switch(lpMsg->wParam) { case VK_RETURN: return DLGC_WANTMESSAGE; } } SetMsgHandled(FALSE); return 0; } void MainGui::OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags) { switch(nChar) { case VK_RETURN: break; default: SetMsgHandled(FALSE); break; } } void MainGui::OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl) { fillProcessListComboBox(ComboProcessList); } void MainGui::OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { processSelectedActionHandler(ComboProcessList.GetCurSel()); } void MainGui::OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl) { pickDllActionHandler(); } void MainGui::OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl) { optionsActionHandler(); } void MainGui::OnDump(UINT uNotifyCode, int nID, CWindow wndCtl) { dumpActionHandler(); } void MainGui::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl) { dumpFixActionHandler(); } void MainGui::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl) { peRebuildActionHandler(); } void MainGui::OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl) { dllInjectActionHandler(); } void MainGui::OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl) { iatAutosearchActionHandler(); } void MainGui::OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl) { getImportsActionHandler(); } void MainGui::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl) { showInvalidImportsActionHandler(); } void MainGui::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl) { showSuspectImportsActionHandler(); } void MainGui::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl) { clearImportsActionHandler(); } void MainGui::OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { // 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(); } void MainGui::setupStatusBar() { StatusBar = CreateStatusWindow(WS_CHILD | WS_VISIBLE | WS_CLIPCHILDREN | WS_CLIPSIBLINGS, L"", m_hWnd, ATL_IDW_STATUS_BAR); CRect rcMain, rcStatus; GetClientRect(&rcMain); StatusBar.GetClientRect(&rcStatus); const int widthIcon = 16 + 2; int widthTexts = (rcMain.Width() - widthIcon) / 3; const int PARTS = 4; int widths[PARTS]; widths[PART_ICON] = widthIcon; widths[PART_COUNT] = widths[PART_ICON] + widthTexts; widths[PART_INVALID] = widths[PART_COUNT] + widthTexts; widths[PART_MODULE] = -1; StatusBar.SetParts(PARTS, widths); StatusBar.SetMinHeight(widthIcon); ResizeClient(rcMain.Width(), rcMain.Height() + rcStatus.Height(), FALSE); } 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); StatusBar.SetText(PART_MODULE, processAccessHelp.selectedModule->getFilename()); } 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::iterator iterator1; std::map::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.hint = (WORD)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& 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); StatusBar.SetText(PART_MODULE, process.filename); } void MainGui::fillProcessListComboBox(CComboBox& hCombo) { hCombo.ResetContent(); std::vector& processList = processLister.getProcessListSnapshot(); for (size_t i = 0; i < processList.size(); i++) { swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("0x%04X - %s - %s"),processList[i].PID,processList[i].filename,processList[i].fullPath); hCombo.AddString(stringBuffer); } } void MainGui::addTextToOutputLog(const WCHAR * text) { if (m_hWnd) { ListLog.SetCurSel(ListLog.AddString(text)); } } void MainGui::clearOutputLog() { if (m_hWnd) { ListLog.ResetContent(); } } bool MainGui::saveLogToFile(const WCHAR * file) { const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian const WCHAR newLine[] = L"\r\n"; bool success = true; HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0); if(hFile != INVALID_HANDLE_VALUE) { ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM); WCHAR * buffer = 0; int bufsize = 0; 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.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 &scyllaPluginList = PluginLoader::getScyllaPluginList(); std::vector &imprecPluginList = PluginLoader::getImprecPluginList(); if (scyllaPluginList.size() > 0) { CMenuHandle newMenu; newMenu.CreatePopupMenu(); for (size_t i = 0; i < scyllaPluginList.size(); i++) { newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName); } hMenu.AppendMenu(MF_MENUBARBREAK); hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins"); } if (imprecPluginList.size() > 0) { CMenuHandle newMenu; newMenu.CreatePopupMenu(); for (size_t i = 0; i < imprecPluginList.size(); i++) { newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName); } hMenu.AppendMenu(MF_MENUBARBREAK); hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins"); } } void MainGui::dumpActionHandler() { if(!selectedProcess) return; WCHAR selectedFilePath[MAX_PATH]; const WCHAR * fileFilter; const WCHAR * defExtension; PeDump peDump; if (processAccessHelp.selectedModule) { fileFilter = filterDll; defExtension = L"dll"; } else { fileFilter = filterExe; defExtension = L"exe"; } if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension)) { if (processAccessHelp.selectedModule) { //dump DLL peDump.imageBase = processAccessHelp.selectedModule->modBaseAddr; peDump.sizeOfImage = processAccessHelp.selectedModule->modBaseSize; //get it from gui peDump.entryPoint = getOEPFromGui(); 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 &scyllaPluginList = PluginLoader::getScyllaPluginList(); std::vector &imprecPluginList = PluginLoader::getImprecPluginList(); menuItem -= PLUGIN_MENU_BASE_ID; dllInjectionPlugin.hProcess = ProcessAccessHelp::hProcess; dllInjectionPlugin.apiReader = &apiReader; if (menuItem < (int)scyllaPluginList.size()) { //scylla plugin dllInjectionPlugin.injectPlugin(scyllaPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); } else { #ifndef _WIN64 menuItem -= (int)scyllaPluginList.size(); //imprec plugin dllInjectionPlugin.injectImprecPlugin(imprecPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); #endif } importsHandling.scanAndFixModuleList(); importsHandling.displayAllImports(); } diff --git a/Scylla/OptionsGui.cpp b/Scylla/OptionsGui.cpp index fb035ca..8394d3e 100644 --- a/Scylla/OptionsGui.cpp +++ b/Scylla/OptionsGui.cpp @@ -1,89 +1,89 @@ #include "OptionsGui.h" #include "ConfigurationHolder.h" BOOL OptionsGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) { loadOptions(); DoDataExchange(DDX_LOAD); // show settings EditSectionName.LimitText(IMAGE_SIZEOF_SHORT_NAME); CenterWindow(); return TRUE; } void OptionsGui::OnOK(UINT uNotifyCode, int nID, CWindow wndCtl) { DoDataExchange(DDX_SAVE); saveOptions(); ConfigurationHolder::saveConfiguration(); EndDialog(0); } void OptionsGui::OnCancel(UINT uNotifyCode, int nID, CWindow wndCtl) { EndDialog(0); } void OptionsGui::saveOptions() { std::map::iterator mapIter; for (mapIter = ConfigurationHolder::getConfigList().begin() ; mapIter != ConfigurationHolder::getConfigList().end(); mapIter++) { switch(mapIter->first) { case USE_PE_HEADER_FROM_DISK: usePEHeaderFromDisk ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case DEBUG_PRIVILEGE: debugPrivilege ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case CREATE_BACKUP: createBackup ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case DLL_INJECTION_AUTO_UNLOAD: dllInjectionAutoUnload ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case UPDATE_HEADER_CHECKSUM: updateHeaderChecksum ? mapIter->second.setTrue() : mapIter->second.setFalse(); break; case IAT_SECTION_NAME: - wcscpy_s(mapIter->second.valueString, _countof(mapIter->second.valueString), iatSectionName); + wcscpy_s(mapIter->second.valueString, 100, iatSectionName); break; } } } void OptionsGui::loadOptions() { std::map::iterator mapIter; for (mapIter = ConfigurationHolder::getConfigList().begin() ; mapIter != ConfigurationHolder::getConfigList().end(); mapIter++) { switch(mapIter->first) { case USE_PE_HEADER_FROM_DISK: usePEHeaderFromDisk = mapIter->second.isTrue(); break; case DEBUG_PRIVILEGE: debugPrivilege = mapIter->second.isTrue(); break; case CREATE_BACKUP: createBackup = mapIter->second.isTrue(); break; case DLL_INJECTION_AUTO_UNLOAD: dllInjectionAutoUnload = mapIter->second.isTrue(); break; case UPDATE_HEADER_CHECKSUM: updateHeaderChecksum = mapIter->second.isTrue(); break; case IAT_SECTION_NAME: - wcsncpy_s(iatSectionName, _countof(iatSectionName), mapIter->second.valueString, _countof(iatSectionName)-1); + wcsncpy_s(iatSectionName,9, mapIter->second.valueString, _countof(iatSectionName)-1); iatSectionName[_countof(iatSectionName)-1] = L'\0'; break; } } } diff --git a/Scylla/PickDllGui.cpp b/Scylla/PickDllGui.cpp index acb1b21..5fc07b3 100644 --- a/Scylla/PickDllGui.cpp +++ b/Scylla/PickDllGui.cpp @@ -1,145 +1,145 @@ #include "PickDllGui.h" PickDllGui::PickDllGui(std::vector &moduleList) : moduleList(moduleList) { selectedModule = 0; prevColumn = -1; ascending = true; } BOOL PickDllGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) { DoDataExchange(); // attach controls DlgResize_Init(true, true); addColumnsToModuleList(ListDLLSelect); displayModuleList(ListDLLSelect); CenterWindow(); return TRUE; } LRESULT PickDllGui::OnListDllColumnClicked(NMHDR* pnmh) { NMLISTVIEW* list = (NMLISTVIEW*)pnmh; int column = list->iSubItem; if(column == prevColumn) { ascending = !ascending; } else { prevColumn = column; ascending = true; } // lo-byte: column, hi-byte: sort-order ListDLLSelect.SortItems(&listviewCompareFunc, MAKEWORD(column, ascending)); return 0; } LRESULT PickDllGui::OnListDllDoubleClick(NMHDR* pnmh) { NMITEMACTIVATE* ia = (NMITEMACTIVATE*)pnmh; LVHITTESTINFO hti; hti.pt = ia->ptAction; int clicked = ListDLLSelect.HitTest(&hti); if(clicked != -1) { selectedModule = (ModuleInfo *)ListDLLSelect.GetItemData(clicked); EndDialog(1); } return 0; } void PickDllGui::OnOK(UINT uNotifyCode, int nID, CWindow wndCtl) { int index = ListDLLSelect.GetSelectionMark(); if (index != -1) { selectedModule = (ModuleInfo *)ListDLLSelect.GetItemData(index); EndDialog(1); } } void PickDllGui::OnCancel(UINT uNotifyCode, int nID, CWindow wndCtl) { EndDialog(0); } void PickDllGui::addColumnsToModuleList(CListViewCtrl& list) { list.SetExtendedListViewStyle(LVS_EX_FULLROWSELECT, LVS_EX_FULLROWSELECT); list.InsertColumn(COL_NAME, L"Name", LVCFMT_LEFT); list.InsertColumn(COL_IMAGEBASE, L"ImageBase", LVCFMT_CENTER); list.InsertColumn(COL_IMAGESIZE, L"ImageSize", LVCFMT_CENTER); list.InsertColumn(COL_PATH, L"Path", LVCFMT_LEFT); } void PickDllGui::displayModuleList(CListViewCtrl& list) { WCHAR temp[20]; list.DeleteAllItems(); std::vector::const_iterator iter; int count = 0; for( iter = moduleList.begin(); iter != moduleList.end(); iter++ , count++) { list.InsertItem(count, iter->getFilename()); #ifdef _WIN64 swprintf_s(temp, _countof(temp), L"%016I64X", iter->modBaseAddr); #else swprintf_s(temp, _countof(temp), L"%08X", iter->modBaseAddr); #endif list.SetItemText(count, COL_IMAGEBASE, temp); swprintf_s(temp, _countof(temp),L"%08X",iter->modBaseSize); list.SetItemText(count, COL_IMAGESIZE, temp); list.SetItemText(count, COL_PATH, iter->fullPath); list.SetItemData(count, (DWORD_PTR)&(*iter)); } list.SetColumnWidth(COL_NAME, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_IMAGEBASE, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_IMAGESIZE, LVSCW_AUTOSIZE_USEHEADER); list.SetColumnWidth(COL_PATH, LVSCW_AUTOSIZE_USEHEADER); } // lParamSort - lo-byte: column, hi-byte: sort-order int PickDllGui::listviewCompareFunc(LPARAM lParam1, LPARAM lParam2, LPARAM lParamSort) { const ModuleInfo * module1 = (ModuleInfo *)lParam1; const ModuleInfo * module2 = (ModuleInfo *)lParam2; int column = LOBYTE(lParamSort); - bool ascending = HIBYTE(lParamSort) != 0; + bool ascending = (HIBYTE(lParamSort) == TRUE); int diff = 0; switch(column) { case COL_NAME: diff = _wcsicmp(module1->getFilename(), module2->getFilename()); break; case COL_IMAGEBASE: diff = module1->modBaseAddr < module2->modBaseAddr ? -1 : 1; break; case COL_IMAGESIZE: diff = module1->modBaseSize < module2->modBaseSize ? -1 : 1; break; case COL_PATH: diff = _wcsicmp(module1->fullPath, module2->fullPath); break; } return ascending ? diff : -diff; } diff --git a/Scylla/definitions.h b/Scylla/definitions.h index 065b0a0..ac6cb76 100644 --- a/Scylla/definitions.h +++ b/Scylla/definitions.h @@ -1,28 +1,28 @@ #pragma once #define APPNAME "Scylla" #ifdef _WIN64 #define ARCHITECTURE "x64" #define PRINTF_DWORD_PTR "%I64X" #define PRINTF_DWORD_PTR_FULL "%016I64X" #define PRINTF_DWORD_PTR_HALF "%08I64X" #define PRINTF_INTEGER "%I64u" #define MAX_HEX_VALUE_EDIT_LENGTH 16 #else #define ARCHITECTURE "x86" #define PRINTF_DWORD_PTR "%X" #define PRINTF_DWORD_PTR_FULL "%08X" #define PRINTF_DWORD_PTR_HALF "%08X" #define PRINTF_INTEGER "%u" #define MAX_HEX_VALUE_EDIT_LENGTH 8 #endif -#define APPVERSION "v0.4" +#define APPVERSION "v0.5 Beta" #define PLUGIN_MENU_BASE_ID 0x10 \ No newline at end of file diff --git a/Scylla/multitree.h b/Scylla/multitree.h index d7cb27e..95ad7da 100644 --- a/Scylla/multitree.h +++ b/Scylla/multitree.h @@ -1,580 +1,580 @@ #pragma once ///////////////////////////////////////////////////////////////////////////// // MultiSelectTree - A tree control with multi-select capabilities // // Written by Bjarke Viksoe (bjarke@viksoe.dk) // Copyright (c) 2005 Bjarke Viksoe. // // Add the following macro to the parent's message map: // REFLECT_NOTIFICATIONS() // // This code may be used in compiled form in any way you desire. This // source file may be redistributed by any means PROVIDING it is // not sold for profit without the authors written consent, and // providing that this notice and the authors name is included. // // This file is provided "as is" with no expressed or implied warranty. // The author accepts no liability if it causes any damage to you or your // computer whatsoever. It's free, so don't hassle me about it. // // Beware of bugs. // #ifndef __cplusplus #error WTL requires C++ compilation (use a .cpp suffix) #endif #ifndef __ATLMISC_H__ #error multitree.h requires atlmisc.h to be included first #endif #ifndef __ATLCRACK_H__ #error multitree.h requires atlcrack.h to be included first #endif #ifndef __ATLCTRLS_H__ #error multitree.h requires atlctrls.h to be included first #endif // Extended MultiSelectTree styles static const DWORD MTVS_EX_NOMARQUEE = 0x00000001; // New control notifications static const UINT TVN_ITEMSELECTING = 0x0001; static const UINT TVN_ITEMSELECTED = 0x0002; static bool operator==(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 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 == TRUE); } } 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) 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; + return (int)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 > { public: DECLARE_WND_SUPERCLASS(_T("WTL_MultiSelectTree"), GetWndClassName()) };