diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp index 7ee228e..79deb91 100644 --- a/Scylla/MainGui.cpp +++ b/Scylla/MainGui.cpp @@ -1,1248 +1,1244 @@ #include "MainGui.h" #include "Architecture.h" //#include "PluginLoader.h" //#include "ConfigurationHolder.h" #include "PeDump.h" #include "PeRebuild.h" #include "DllInjectionPlugin.h" #include "DisassemblerGui.h" #include "PickApiGui.h" //#include "NativeWinApi.h" #include "ImportRebuild.h" #include "SystemInformation.h" #include "Scylla.h" #include "AboutGui.h" #include "OptionsGui.h" #include "TreeImportExport.h" extern CAppModule _Module; // o_O const WCHAR MainGui::filterExe[] = L"Executable (*.exe)\0*.exe\0All files\0*.*\0"; const WCHAR MainGui::filterDll[] = L"Dynamic Link Library (*.dll)\0*.dll\0All files\0*.*\0"; const WCHAR MainGui::filterExeDll[] = L"Executable (*.exe)\0*.exe\0Dynamic Link Library (*.dll)\0*.dll\0All files\0*.*\0"; const WCHAR MainGui::filterTxt[] = L"Text file (*.txt)\0*.txt\0All files\0*.*\0"; const WCHAR MainGui::filterXml[] = L"XML file (*.xml)\0*.xml\0All files\0*.*\0"; MainGui::MainGui() : selectedProcess(0), importsHandling(TreeImports), TreeImportsSubclass(this, IDC_TREE_IMPORTS) { /* Logger::getDebugLogFilePath(); ConfigurationHolder::loadConfiguration(); PluginLoader::findAllPlugins(); NativeWinApi::initialize(); SystemInformation::getSystemInformation(); if(ConfigurationHolder::getConfigObject(DEBUG_PRIVILEGE)->isTrue()) { processLister.setDebugPrivileges(); } ProcessAccessHelp::getProcessModules(GetCurrentProcessId(), ProcessAccessHelp::ownModuleList); */ Scylla::init(); hIcon.LoadIcon(IDI_ICON_SCYLLA); hMenuImports.LoadMenu(IDR_MENU_IMPORTS); hMenuLog.LoadMenu(IDR_MENU_LOG); accelerators.LoadAccelerators(IDR_ACCELERATOR_MAIN); hIconCheck.LoadIcon(IDI_ICON_CHECK, 16, 16); hIconWarning.LoadIcon(IDI_ICON_WARNING, 16, 16); hIconError.LoadIcon(IDI_ICON_ERROR, 16, 16); } BOOL MainGui::PreTranslateMessage(MSG* pMsg) { if(accelerators.TranslateAccelerator(m_hWnd, pMsg)) { return TRUE; // handled keyboard shortcuts } else if(IsDialogMessage(pMsg)) { return TRUE; // handled dialog messages } return FALSE; } BOOL MainGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam) { if (SystemInformation::currenOS == UNKNOWN_OS) { if(IDCANCEL == MessageBox(L"Operating System is not supported\r\nContinue anyway?", L"Scylla", MB_ICONWARNING | MB_OKCANCEL)) { SendMessage(WM_CLOSE); return FALSE; } } // register ourselves to receive PreTranslateMessage CMessageLoop* pLoop = _Module.GetMessageLoop(); pLoop->AddMessageFilter(this); setupStatusBar(); DoDataExchange(); // attach controls DlgResize_Init(true, true); // init CDialogResize Scylla::windowLog.setWindow(ListLog); appendPluginListToMenu(hMenuImports.GetSubMenu(0)); appendPluginListToMenu(CMenuHandle(GetMenu()).GetSubMenu(MenuImportsOffsetTrace)); enableDialogControls(FALSE); setIconAndDialogCaption(); return TRUE; } void MainGui::OnDestroy() { PostQuitMessage(0); } void MainGui::OnSize(UINT nType, CSize size) { StatusBar.SendMessage(WM_SIZE); SetMsgHandled(FALSE); } void MainGui::OnContextMenu(CWindow wnd, CPoint point) { switch(wnd.GetDlgCtrlID()) { case IDC_TREE_IMPORTS: DisplayContextMenuImports(wnd, point); return; case IDC_LIST_LOG: DisplayContextMenuLog(wnd, point); return; } SetMsgHandled(FALSE); } void MainGui::OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl) { // Handle plugin trace menu selection if(uNotifyCode == 0 && !wndCtl.IsWindow()) // make sure it's a menu { if ((nID >= PLUGIN_MENU_BASE_ID) && (nID <= (int)(Scylla::plugins.getScyllaPluginList().size() + Scylla::plugins.getImprecPluginList().size() + PLUGIN_MENU_BASE_ID))) { pluginActionHandler(nID); return; } } SetMsgHandled(FALSE); } LRESULT MainGui::OnTreeImportsDoubleClick(const NMHDR* pnmh) { if(TreeImports.GetCount() < 1) return 0; // Get item under cursor CTreeItem over = findTreeItem(CPoint(GetMessagePos()), true); if(over && importsHandling.isImport(over)) { pickApiActionHandler(over); } return 0; } LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh) { const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh; switch(tkd->wVKey) { case VK_RETURN: { CTreeItem selected = TreeImports.GetFocusItem(); if(!selected.IsNull() && importsHandling.isImport(selected)) { pickApiActionHandler(selected); } } return 1; case VK_DELETE: deleteSelectedImportsActionHandler(); return 1; } SetMsgHandled(FALSE); return 0; } UINT MainGui::OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg) { if(lpMsg) { switch(lpMsg->wParam) { case VK_RETURN: return DLGC_WANTMESSAGE; } } SetMsgHandled(FALSE); return 0; } void MainGui::OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags) { switch(nChar) { case VK_RETURN: break; default: SetMsgHandled(FALSE); break; } } void MainGui::OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl) { fillProcessListComboBox(ComboProcessList); } void MainGui::OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { processSelectedActionHandler(ComboProcessList.GetCurSel()); } void MainGui::OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl) { pickDllActionHandler(); } void MainGui::OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl) { optionsActionHandler(); } void MainGui::OnDump(UINT uNotifyCode, int nID, CWindow wndCtl) { dumpActionHandler(); } void MainGui::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl) { dumpFixActionHandler(); } void MainGui::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl) { peRebuildActionHandler(); } void MainGui::OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl) { dllInjectActionHandler(); } void MainGui::OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl) { iatAutosearchActionHandler(); } void MainGui::OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl) { getImportsActionHandler(); } void MainGui::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl) { showInvalidImportsActionHandler(); } void MainGui::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl) { showSuspectImportsActionHandler(); } void MainGui::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl) { clearImportsActionHandler(); } void MainGui::OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { invalidateSelectedImportsActionHandler(); } void MainGui::OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl) { deleteSelectedImportsActionHandler(); } void MainGui::OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl) { saveTreeActionHandler(); } void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl) { loadTreeActionHandler(); } void MainGui::OnAutotrace(UINT uNotifyCode, int nID, CWindow wndCtl) { // TODO } void MainGui::OnExit(UINT uNotifyCode, int nID, CWindow wndCtl) { DestroyWindow(); } void MainGui::OnAbout(UINT uNotifyCode, int nID, CWindow wndCtl) { showAboutDialog(); } void MainGui::setupStatusBar() { StatusBar.Create(m_hWnd, NULL, L"", WS_CHILD | WS_VISIBLE | WS_CLIPCHILDREN | WS_CLIPSIBLINGS | SBARS_TOOLTIPS, NULL, IDC_STATUS_BAR); CRect rcMain, rcStatus; GetClientRect(&rcMain); StatusBar.GetWindowRect(&rcStatus); const int PARTS = 4; int widths[PARTS]; widths[PART_COUNT] = rcMain.Width() / 5; widths[PART_INVALID] = widths[PART_COUNT] + rcMain.Width() / 5; widths[PART_IMAGEBASE] = widths[PART_INVALID] + rcMain.Width() / 3; widths[PART_MODULE] = -1; StatusBar.SetParts(PARTS, widths); ResizeClient(rcMain.Width(), rcMain.Height() + rcStatus.Height(), FALSE); } void MainGui::updateStatusBar() { // Rewrite ImportsHandling so we get these easily unsigned int totalImports = importsHandling.thunkCount(); unsigned int invalidImports = importsHandling.invalidThunkCount(); // \t = center, \t\t = right-align swprintf_s(stringBuffer, L"\tImports: %u", totalImports); StatusBar.SetText(PART_COUNT, stringBuffer); if(invalidImports > 0) { StatusBar.SetIcon(PART_INVALID, hIconError); } else { StatusBar.SetIcon(PART_INVALID, hIconCheck); } swprintf_s(stringBuffer, L"\tInvalid: %u", invalidImports); StatusBar.SetText(PART_INVALID, stringBuffer); if(selectedProcess) { DWORD_PTR imageBase = 0; const WCHAR * fileName = 0; if(ProcessAccessHelp::selectedModule) { imageBase = ProcessAccessHelp::selectedModule->modBaseAddr; fileName = ProcessAccessHelp::selectedModule->getFilename(); } else { imageBase = selectedProcess->imageBase; fileName = selectedProcess->filename; } swprintf_s(stringBuffer, L"\tImagebase: " PRINTF_DWORD_PTR_FULL, imageBase); StatusBar.SetText(PART_IMAGEBASE, stringBuffer); StatusBar.SetText(PART_MODULE, fileName); StatusBar.SetTipText(PART_MODULE, fileName); } else { StatusBar.SetText(PART_IMAGEBASE, L""); StatusBar.SetText(PART_MODULE, L""); } } bool MainGui::showFileDialog(WCHAR * selectedFile, bool save, const WCHAR * defFileName, const WCHAR * filter, const WCHAR * defExtension, const WCHAR * directory) { OPENFILENAME ofn = {0}; // WTL doesn't support new explorer styles on Vista and up // This is because it uses a custom hook, we could remove it or derive // from CFileDialog but this solution is easier and allows more control anyway (e.g. initial dir) if(defFileName) { wcscpy_s(selectedFile, MAX_PATH, defFileName); } else { selectedFile[0] = L'\0'; } ofn.lStructSize = sizeof(ofn); ofn.hwndOwner = m_hWnd; ofn.lpstrFilter = filter; ofn.lpstrDefExt = defExtension; // only first 3 chars are used, no dots! ofn.lpstrFile = selectedFile; ofn.lpstrInitialDir = directory; ofn.nMaxFile = MAX_PATH; ofn.Flags = OFN_PATHMUSTEXIST | OFN_HIDEREADONLY; /* *OFN_EXPLORER is automatically used, it only has to be specified *if using a custom hook *OFN_LONGNAMES is automatically used by explorer-style dialogs */ if(save) ofn.Flags |= OFN_OVERWRITEPROMPT; else ofn.Flags |= OFN_FILEMUSTEXIST; if(save) return 0 != GetSaveFileName(&ofn); else return 0 != GetOpenFileName(&ofn); } void MainGui::setIconAndDialogCaption() { SetIcon(hIcon, TRUE); SetIcon(hIcon, FALSE); SetWindowText(APPNAME L" " ARCHITECTURE L" " APPVERSION); } void MainGui::pickDllActionHandler() { if(!selectedProcess) return; PickDllGui dlgPickDll(ProcessAccessHelp::moduleList); if(dlgPickDll.DoModal()) { //get selected module ProcessAccessHelp::selectedModule = dlgPickDll.getSelectedModule(); ProcessAccessHelp::targetImageBase = ProcessAccessHelp::selectedModule->modBaseAddr; Scylla::windowLog.log(L"->>> Module %s selected.", ProcessAccessHelp::selectedModule->getFilename()); Scylla::windowLog.log(L"Imagebase: " PRINTF_DWORD_PTR_FULL L" Size: %08X", ProcessAccessHelp::selectedModule->modBaseAddr, ProcessAccessHelp::selectedModule->modBaseSize); } else { ProcessAccessHelp::selectedModule = 0; } updateStatusBar(); } void MainGui::pickApiActionHandler(CTreeItem item) { if(!importsHandling.isImport(item)) return; // TODO: new node when user picked an API from another DLL? PickApiGui dlgPickApi(ProcessAccessHelp::moduleList); if(dlgPickApi.DoModal()) { const ApiInfo* api = dlgPickApi.getSelectedApi(); if(api && api->module) { importsHandling.setImport(item, api->module->getFilename(), api->name, api->ordinal, api->hint, true, api->isForwarded); } } updateStatusBar(); } void MainGui::startDisassemblerGui(CTreeItem selectedTreeNode) { if(!selectedProcess) return; DWORD_PTR address = importsHandling.getApiAddressByNode(selectedTreeNode); if (address) { BYTE test; if(!ProcessAccessHelp::readMemoryFromProcess(address, sizeof(test), &test)) { swprintf_s(stringBuffer, L"Can't read memory at " PRINTF_DWORD_PTR_FULL, address); MessageBox(stringBuffer, L"Failure", MB_ICONERROR); } else { DisassemblerGui dlgDisassembler(address); dlgDisassembler.DoModal(); } } } void MainGui::processSelectedActionHandler(int index) { std::vector& processList = Scylla::processLister.getProcessList(); Process &process = processList.at(index); selectedProcess = 0; clearImportsActionHandler(); Scylla::windowLog.log(L"Analyzing %s", process.fullPath); if (ProcessAccessHelp::hProcess != 0) { ProcessAccessHelp::closeProcessHandle(); apiReader.clearAll(); } if (!ProcessAccessHelp::openProcessHandle(process.PID)) { enableDialogControls(FALSE); Scylla::windowLog.log(L"Error: Cannot open process handle."); updateStatusBar(); return; } ProcessAccessHelp::getProcessModules(process.PID, ProcessAccessHelp::moduleList); apiReader.readApisFromModuleList(); Scylla::windowLog.log(L"Loading modules done."); //TODO improve ProcessAccessHelp::selectedModule = 0; ProcessAccessHelp::targetSizeOfImage = process.imageSize; ProcessAccessHelp::targetImageBase = process.imageBase; ProcessAccessHelp::getSizeOfImageCurrentProcess(); process.imageSize = (DWORD)ProcessAccessHelp::targetSizeOfImage; Scylla::windowLog.log(L"Imagebase: " PRINTF_DWORD_PTR_FULL L" Size: %08X", process.imageBase, process.imageSize); process.entryPoint = ProcessAccessHelp::getEntryPointFromFile(process.fullPath); EditOEPAddress.SetValue(process.entryPoint + process.imageBase); selectedProcess = &process; enableDialogControls(TRUE); updateStatusBar(); } void MainGui::fillProcessListComboBox(CComboBox& hCombo) { hCombo.ResetContent(); std::vector& processList = Scylla::processLister.getProcessListSnapshot(); for (size_t i = 0; i < processList.size(); i++) { swprintf_s(stringBuffer, L"0x%04X - %s - %s", processList[i].PID, processList[i].filename, processList[i].fullPath); hCombo.AddString(stringBuffer); } } /* void MainGui::addTextToOutputLog(const WCHAR * text) { if (m_hWnd) { ListLog.SetCurSel(ListLog.AddString(text)); } } */ void MainGui::clearOutputLog() { if (m_hWnd) { ListLog.ResetContent(); } } bool MainGui::saveLogToFile(const WCHAR * file) { const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian const WCHAR newLine[] = L"\r\n"; bool success = true; HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0); if(hFile != INVALID_HANDLE_VALUE) { ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM); WCHAR * buffer = 0; size_t bufsize = 0; for(int i = 0; i < ListLog.GetCount(); i++) { size_t size = ListLog.GetTextLen(i); size += _countof(newLine)-1; if(size+1 > bufsize) { bufsize = size+1; delete[] buffer; try { buffer = new WCHAR[bufsize]; } catch(std::bad_alloc&) { buffer = 0; success = false; break; } } ListLog.GetText(i, buffer); wcscat_s(buffer, bufsize, newLine); ProcessAccessHelp::writeMemoryToFileEnd(hFile, (DWORD)(size * sizeof(WCHAR)), buffer); } delete[] buffer; CloseHandle(hFile); } return success; } void MainGui::showInvalidImportsActionHandler() { importsHandling.selectImports(true, false); GotoDlgCtrl(TreeImports); } void MainGui::showSuspectImportsActionHandler() { importsHandling.selectImports(false, true); GotoDlgCtrl(TreeImports); } void MainGui::deleteSelectedImportsActionHandler() { CTreeItem selected = TreeImports.GetFirstSelectedItem(); while(!selected.IsNull()) { if(importsHandling.isModule(selected)) { importsHandling.cutModule(selected); } else { importsHandling.cutImport(selected); } selected = TreeImports.GetNextSelectedItem(selected); } updateStatusBar(); } void MainGui::invalidateSelectedImportsActionHandler() { CTreeItem selected = TreeImports.GetFirstSelectedItem(); while(!selected.IsNull()) { if(importsHandling.isImport(selected)) { importsHandling.invalidateImport(selected); } selected = TreeImports.GetNextSelectedItem(selected); } updateStatusBar(); } void MainGui::loadTreeActionHandler() { if(!selectedProcess) return; WCHAR selectedFilePath[MAX_PATH]; - TreeImportExport treeIO; - DWORD_PTR addrOEP = 0; - DWORD_PTR addrIAT = 0; - DWORD sizeIAT = 0; - getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if(showFileDialog(selectedFilePath, false, NULL, filterXml, NULL, stringBuffer)) { - if(!treeIO.importTreeList(selectedFilePath, importsHandling.moduleList, &addrOEP, &addrIAT, &sizeIAT)) + TreeImportExport treeIO(selectedFilePath); + DWORD_PTR addrOEP = 0; + DWORD_PTR addrIAT = 0; + DWORD sizeIAT = 0; + + if(!treeIO.importTreeList(importsHandling.moduleList, &addrOEP, &addrIAT, &sizeIAT)) { Scylla::windowLog.log(L"Loading tree file failed %s", selectedFilePath); MessageBox(L"Loading tree file failed.", L"Failure", MB_ICONERROR); } else { EditOEPAddress.SetValue(addrOEP); EditIATAddress.SetValue(addrIAT); EditIATSize.SetValue(sizeIAT); importsHandling.displayAllImports(); updateStatusBar(); Scylla::windowLog.log(L"Loaded tree file %s", selectedFilePath); Scylla::windowLog.log(L"-> OEP: " PRINTF_DWORD_PTR_FULL, addrOEP); Scylla::windowLog.log(L"-> IAT: " PRINTF_DWORD_PTR_FULL L" Size: " PRINTF_DWORD_PTR, addrIAT, sizeIAT); } } } void MainGui::saveTreeActionHandler() { if(!selectedProcess) return; WCHAR selectedFilePath[MAX_PATH]; - TreeImportExport treeIO; - DWORD_PTR addrOEP; - DWORD_PTR addrIAT; - DWORD sizeIAT; - getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if(showFileDialog(selectedFilePath, true, NULL, filterXml, L"xml", stringBuffer)) { - addrOEP = EditOEPAddress.GetValue(); - addrIAT = EditIATAddress.GetValue(); - sizeIAT = EditIATSize.GetValue(); + TreeImportExport treeIO(selectedFilePath); + DWORD_PTR addrOEP = EditOEPAddress.GetValue(); + DWORD_PTR addrIAT = EditIATAddress.GetValue(); + DWORD sizeIAT = EditIATSize.GetValue(); - if(!treeIO.exportTreeList(selectedFilePath, importsHandling.moduleList, selectedProcess, addrOEP, addrIAT, sizeIAT)) + if(!treeIO.exportTreeList(importsHandling.moduleList, selectedProcess, addrOEP, addrIAT, sizeIAT)) { Scylla::windowLog.log(L"Saving tree file failed %s", selectedFilePath); MessageBox(L"Saving tree file failed.", L"Failure", MB_ICONERROR); } else { Scylla::windowLog.log(L"Saved tree file %s", selectedFilePath); } } } void MainGui::iatAutosearchActionHandler() { DWORD_PTR searchAddress = 0; DWORD_PTR addressIAT = 0; DWORD sizeIAT = 0; IATSearch iatSearch; if(!selectedProcess) return; if(EditOEPAddress.GetWindowTextLength() > 0) { searchAddress = EditOEPAddress.GetValue(); if (searchAddress) { if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT)) { Scylla::windowLog.log(L"IAT found at VA " PRINTF_DWORD_PTR_FULL L" RVA " PRINTF_DWORD_PTR_FULL L" Size 0x%04X (%d)", addressIAT, addressIAT - ProcessAccessHelp::targetImageBase, sizeIAT, sizeIAT); EditIATAddress.SetValue(addressIAT); EditIATSize.SetValue(sizeIAT); swprintf_s(stringBuffer, L"IAT found:\r\n\r\nStart: " PRINTF_DWORD_PTR_FULL L"\r\nSize: 0x%04X (%d) ", addressIAT, sizeIAT, sizeIAT); MessageBox(stringBuffer, L"IAT found", MB_ICONINFORMATION); } else { Scylla::windowLog.log(L"IAT not found at OEP " PRINTF_DWORD_PTR_FULL L"!", searchAddress); } } } } void MainGui::getImportsActionHandler() { if(!selectedProcess) return; DWORD_PTR addressIAT = EditIATAddress.GetValue(); DWORD sizeIAT = EditIATSize.GetValue(); if (addressIAT && sizeIAT) { apiReader.readAndParseIAT(addressIAT, sizeIAT, importsHandling.moduleList); importsHandling.displayAllImports(); updateStatusBar(); } } void MainGui::SetupImportsMenuItems(CTreeItem item) { bool isItem, isImport = false; isItem = !item.IsNull(); if(isItem) { isImport = importsHandling.isImport(item); } CMenuHandle hSub = hMenuImports.GetSubMenu(0); UINT itemOnly = isItem ? MF_ENABLED : MF_GRAYED; UINT importOnly = isImport ? MF_ENABLED : MF_GRAYED; hSub.EnableMenuItem(ID__INVALIDATE, itemOnly); hSub.EnableMenuItem(ID__DISASSEMBLE, importOnly); hSub.EnableMenuItem(ID__CUTTHUNK, importOnly); hSub.EnableMenuItem(ID__DELETETREENODE, itemOnly); } void MainGui::DisplayContextMenuImports(CWindow hwnd, CPoint pt) { if(TreeImports.GetCount() < 1) return; CTreeItem over, parent; if(pt.x == -1 && pt.y == -1) // invoked by keyboard { CRect pos; over = TreeImports.GetFocusItem(); if(over) { over.EnsureVisible(); over.GetRect(&pos, TRUE); TreeImports.ClientToScreen(&pos); } else { TreeImports.GetWindowRect(&pos); } pt = pos.TopLeft(); } else { // Get item under cursor over = findTreeItem(pt, true); } SetupImportsMenuItems(over); CMenuHandle hSub = hMenuImports.GetSubMenu(0); BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD, pt.x, pt.y, hwnd); if (menuItem) { if ((menuItem >= PLUGIN_MENU_BASE_ID) && (menuItem <= (int)(Scylla::plugins.getScyllaPluginList().size() + Scylla::plugins.getImprecPluginList().size() + PLUGIN_MENU_BASE_ID))) { //wsprintf(stringBuffer, L"%d %s\n",menuItem,pluginList[menuItem - PLUGIN_MENU_BASE_ID].pluginName); //MessageBox(stringBuffer, L"plugin selection"); pluginActionHandler(menuItem); return; } switch (menuItem) { case ID__INVALIDATE: if(importsHandling.isModule(over)) importsHandling.invalidateModule(over); else importsHandling.invalidateImport(over); break; case ID__DISASSEMBLE: startDisassemblerGui(over); break; case ID__EXPANDALLNODES: importsHandling.expandAllTreeNodes(); break; case ID__COLLAPSEALLNODES: importsHandling.collapseAllTreeNodes(); break; case ID__CUTTHUNK: importsHandling.cutImport(over); break; case ID__DELETETREENODE: importsHandling.cutModule(importsHandling.isImport(over) ? over.GetParent() : over); break; } } updateStatusBar(); } void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt) { if(pt.x == -1 && pt.y == -1) // invoked by keyboard { CRect pos; ListLog.GetWindowRect(&pos); pt = pos.TopLeft(); } CMenuHandle hSub = hMenuLog.GetSubMenu(0); BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD, pt.x, pt.y, hwnd); if (menuItem) { switch (menuItem) { case ID__SAVE: WCHAR selectedFilePath[MAX_PATH]; getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt", stringBuffer)) { saveLogToFile(selectedFilePath); } break; case ID__CLEAR: clearOutputLog(); break; } } } void MainGui::appendPluginListToMenu(CMenuHandle hMenu) { std::vector &scyllaPluginList = Scylla::plugins.getScyllaPluginList(); std::vector &imprecPluginList = Scylla::plugins.getImprecPluginList(); if (scyllaPluginList.size() > 0) { CMenuHandle newMenu; newMenu.CreatePopupMenu(); for (size_t i = 0; i < scyllaPluginList.size(); i++) { newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName); } hMenu.AppendMenu(MF_MENUBARBREAK); hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins"); } if (imprecPluginList.size() > 0) { CMenuHandle newMenu; newMenu.CreatePopupMenu(); for (size_t i = 0; i < imprecPluginList.size(); i++) { newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName); } hMenu.AppendMenu(MF_MENUBARBREAK); hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins"); } } void MainGui::dumpActionHandler() { if(!selectedProcess) return; WCHAR selectedFilePath[MAX_PATH]; const WCHAR * fileFilter; const WCHAR * defExtension; PeDump peDump; if (ProcessAccessHelp::selectedModule) { fileFilter = filterDll; defExtension = L"dll"; } else { fileFilter = filterExe; defExtension = L"exe"; } getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension, stringBuffer)) { if (ProcessAccessHelp::selectedModule) { //dump DLL peDump.imageBase = ProcessAccessHelp::selectedModule->modBaseAddr; peDump.sizeOfImage = ProcessAccessHelp::selectedModule->modBaseSize; //get it from gui peDump.entryPoint = EditOEPAddress.GetValue(); wcscpy_s(peDump.fullpath, ProcessAccessHelp::selectedModule->fullPath); } else { peDump.imageBase = ProcessAccessHelp::targetImageBase; peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage; //get it from gui peDump.entryPoint = EditOEPAddress.GetValue(); wcscpy_s(peDump.fullpath, selectedProcess->fullPath); } peDump.useHeaderFromDisk = Scylla::config[USE_PE_HEADER_FROM_DISK].isTrue(); if (peDump.dumpCompleteProcessToDisk(selectedFilePath)) { Scylla::windowLog.log(L"Dump success %s", selectedFilePath); } else { Scylla::windowLog.log(L"Error: Cannot dump image."); MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR); } } } void MainGui::peRebuildActionHandler() { DWORD newSize = 0; WCHAR selectedFilePath[MAX_PATH]; PeRebuild peRebuild; getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if(showFileDialog(selectedFilePath, false, NULL, filterExeDll, NULL, stringBuffer)) { if (Scylla::config[CREATE_BACKUP].isTrue()) { if (!ProcessAccessHelp::createBackupFile(selectedFilePath)) { Scylla::windowLog.log(L"Creating backup file failed %s", selectedFilePath); } } LONGLONG fileSize = ProcessAccessHelp::getFileSize(selectedFilePath); LPVOID mapped = peRebuild.createFileMappingViewFull(selectedFilePath); newSize = peRebuild.realignPE(mapped, (DWORD)fileSize); peRebuild.closeAllMappingHandles(); if (newSize < 10) { Scylla::windowLog.log(L"Rebuild failed %s", selectedFilePath); MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR); } else { peRebuild.truncateFile(selectedFilePath, newSize); Scylla::windowLog.log(L"Rebuild success %s", selectedFilePath); Scylla::windowLog.log(L"-> Old file size 0x%08X new file size 0x%08X (%d %%)", (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) ); } } } void MainGui::dumpFixActionHandler() { if(!selectedProcess) return; if (TreeImports.GetCount() < 2) { Scylla::windowLog.log(L"Nothing to rebuild"); return; } WCHAR newFilePath[MAX_PATH]; WCHAR selectedFilePath[MAX_PATH]; const WCHAR * fileFilter; if (ProcessAccessHelp::selectedModule) { fileFilter = filterDll; } else { fileFilter = filterExe; } getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if (showFileDialog(selectedFilePath, false, NULL, fileFilter, NULL, stringBuffer)) { wcscpy_s(newFilePath, selectedFilePath); const WCHAR * extension = 0; WCHAR* dot = wcsrchr(newFilePath, L'.'); if (dot) { *dot = L'\0'; extension = selectedFilePath + (dot - newFilePath); //wcsrchr(selectedFilePath, L'.'); } wcscat_s(newFilePath, L"_SCY"); if(extension) { wcscat_s(newFilePath, extension); } ImportRebuild importRebuild; if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList)) { Scylla::windowLog.log(L"Import Rebuild success %s", newFilePath); } else { Scylla::windowLog.log(L"Import Rebuild failed %s", selectedFilePath); MessageBox(L"Import Rebuild failed", L"Failure", MB_ICONERROR); } } } void MainGui::enableDialogControls(BOOL value) { BOOL valButton = value ? TRUE : FALSE; GetDlgItem(IDC_BTN_PICKDLL).EnableWindow(valButton); GetDlgItem(IDC_BTN_DUMP).EnableWindow(valButton); GetDlgItem(IDC_BTN_FIXDUMP).EnableWindow(valButton); GetDlgItem(IDC_BTN_IATAUTOSEARCH).EnableWindow(valButton); GetDlgItem(IDC_BTN_GETIMPORTS).EnableWindow(valButton); GetDlgItem(IDC_BTN_SUSPECTIMPORTS).EnableWindow(valButton); GetDlgItem(IDC_BTN_INVALIDIMPORTS).EnableWindow(valButton); GetDlgItem(IDC_BTN_CLEARIMPORTS).EnableWindow(valButton); CMenuHandle menu = GetMenu(); UINT valMenu = value ? MF_ENABLED : MF_GRAYED; menu.EnableMenuItem(ID_FILE_DUMP, valMenu); menu.EnableMenuItem(ID_FILE_FIXDUMP, valMenu); menu.EnableMenuItem(ID_IMPORTS_INVALIDATESELECTED, valMenu); menu.EnableMenuItem(ID_IMPORTS_CUTSELECTED, valMenu); menu.EnableMenuItem(ID_IMPORTS_SAVETREE, valMenu); menu.EnableMenuItem(ID_IMPORTS_LOADTREE, valMenu); menu.EnableMenuItem(ID_MISC_DLLINJECTION, valMenu); menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetScylla, MF_BYPOSITION | valMenu); menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetImpRec, MF_BYPOSITION | valMenu); //not yet implemented GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE); menu.EnableMenuItem(ID_TRACE_AUTOTRACE, MF_GRAYED); } CTreeItem MainGui::findTreeItem(CPoint pt, bool screenCoordinates) { if(screenCoordinates) { TreeImports.ScreenToClient(&pt); } UINT flags; CTreeItem over = TreeImports.HitTest(pt, &flags); if(over) { if(!(flags & TVHT_ONITEM)) { over.m_hTreeItem = NULL; } } return over; } void MainGui::showAboutDialog() { AboutGui dlgAbout; dlgAbout.DoModal(); } void MainGui::dllInjectActionHandler() { if(!selectedProcess) return; WCHAR selectedFilePath[MAX_PATH]; HMODULE hMod = 0; DllInjection dllInjection; getCurrentModulePath(stringBuffer, _countof(stringBuffer)); if (showFileDialog(selectedFilePath, false, NULL, filterDll, NULL, stringBuffer)) { hMod = dllInjection.dllInjection(ProcessAccessHelp::hProcess, selectedFilePath); if (hMod && Scylla::config[DLL_INJECTION_AUTO_UNLOAD].isTrue()) { if (!dllInjection.unloadDllInProcess(ProcessAccessHelp::hProcess, hMod)) { Scylla::windowLog.log(L"DLL unloading failed, target %s", selectedFilePath); } } if (hMod) { Scylla::windowLog.log(L"DLL Injection was successful, target %s", selectedFilePath); } else { Scylla::windowLog.log(L"DLL Injection failed, target %s", selectedFilePath); } } } void MainGui::optionsActionHandler() { OptionsGui dlgOptions; dlgOptions.DoModal(); } void MainGui::clearImportsActionHandler() { importsHandling.clearAllImports(); updateStatusBar(); } void MainGui::pluginActionHandler( int menuItem ) { if(!selectedProcess) return; DllInjectionPlugin dllInjectionPlugin; std::vector &scyllaPluginList = Scylla::plugins.getScyllaPluginList(); std::vector &imprecPluginList = Scylla::plugins.getImprecPluginList(); menuItem -= PLUGIN_MENU_BASE_ID; dllInjectionPlugin.hProcess = ProcessAccessHelp::hProcess; dllInjectionPlugin.apiReader = &apiReader; if (menuItem < (int)scyllaPluginList.size()) { //scylla plugin dllInjectionPlugin.injectPlugin(scyllaPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); } else { #ifndef _WIN64 menuItem -= (int)scyllaPluginList.size(); //imprec plugin dllInjectionPlugin.injectImprecPlugin(imprecPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize); #endif } importsHandling.scanAndFixModuleList(); importsHandling.displayAllImports(); updateStatusBar(); } bool MainGui::getCurrentModulePath(WCHAR * buffer, size_t bufferSize) { if(!selectedProcess) return false; if(ProcessAccessHelp::selectedModule) { wcscpy_s(buffer, bufferSize, ProcessAccessHelp::selectedModule->fullPath); } else { wcscpy_s(buffer, bufferSize, selectedProcess->fullPath); } WCHAR * slash = wcsrchr(buffer, L'\\'); if(slash) { *(slash+1) = L'\0'; } return true; } diff --git a/Scylla/TreeImportExport.cpp b/Scylla/TreeImportExport.cpp index e17c027..0dc54c2 100644 --- a/Scylla/TreeImportExport.cpp +++ b/Scylla/TreeImportExport.cpp @@ -1,337 +1,296 @@ #include "TreeImportExport.h" #include "Architecture.h" #include "Scylla.h" #include "StringConversion.h" -bool TreeImportExport::exportTreeList(const WCHAR * targetXmlFile, std::map & moduleList, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT) +TreeImportExport::TreeImportExport(const WCHAR * targetXmlFile) +{ + wcscpy_s(xmlPath, targetXmlFile); +} + +bool TreeImportExport::exportTreeList(const std::map & moduleList, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT) { TiXmlDocument doc; - TiXmlDeclaration * decl = new TiXmlDeclaration( "1.0", "", ""); + TiXmlDeclaration * decl = new TiXmlDeclaration("1.0", "", ""); doc.LinkEndChild(decl); TiXmlElement * rootElement = new TiXmlElement("target"); - setTargetInformation(rootElement, process,addressOEP,addressIAT,sizeIAT); + setTargetInformation(rootElement, process, addressOEP, addressIAT, sizeIAT); addModuleListToRootElement(rootElement, moduleList); doc.LinkEndChild(rootElement); - return saveXmlToFile(doc,targetXmlFile); + return saveXmlToFile(doc, xmlPath); } -bool TreeImportExport::importTreeList(const WCHAR * targetXmlFile, std::map & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT) +bool TreeImportExport::importTreeList(std::map & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT) { - TiXmlElement * targetElement; - TiXmlDocument doc; - char * buffer = readXmlFile(targetXmlFile); - int count = 0; - moduleList.clear(); + *addressOEP = *addressIAT = 0; + *sizeIAT = 0; - if (buffer) + TiXmlDocument doc; + if(!readXmlFile(doc, xmlPath)) { - doc.Parse(buffer); - if (doc.Error()) - { - Scylla::windowLog.log(L"Load Tree :: Error parsing xml %S: %S\r\n", doc.Value(), doc.ErrorDesc()); - delete [] buffer; - return false; - } - - targetElement = doc.FirstChildElement(); - - *addressOEP = ConvertStringToDwordPtr(targetElement->Attribute("oep_va")); - *addressIAT = ConvertStringToDwordPtr(targetElement->Attribute("iat_va")); - *sizeIAT = (DWORD)ConvertStringToDwordPtr(targetElement->Attribute("iat_size")); + Scylla::windowLog.log(L"Load Tree :: Error parsing xml %S: %S\r\n", doc.Value(), doc.ErrorDesc()); + return false; + } - //test = targetElement->Attribute("filename"); + TiXmlElement * targetElement = doc.FirstChildElement(); - parseAllElementModules(targetElement, moduleList); + *addressOEP = ConvertStringToDwordPtr(targetElement->Attribute("oep_va")); + *addressIAT = ConvertStringToDwordPtr(targetElement->Attribute("iat_va")); + *sizeIAT = (DWORD)ConvertStringToDwordPtr(targetElement->Attribute("iat_size")); - delete [] buffer; - } + parseAllElementModules(targetElement, moduleList); return true; } -void TreeImportExport::setTargetInformation(TiXmlElement * rootElement, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT) +void TreeImportExport::setTargetInformation(TiXmlElement * rootElement, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT) { StringConversion::ToASCII(process->filename, xmlStringBuffer, _countof(xmlStringBuffer)); rootElement->SetAttribute("filename", xmlStringBuffer); ConvertDwordPtrToString(addressOEP); rootElement->SetAttribute("oep_va", xmlStringBuffer); ConvertDwordPtrToString(addressIAT); rootElement->SetAttribute("iat_va", xmlStringBuffer); ConvertDwordPtrToString(sizeIAT); rootElement->SetAttribute("iat_size", xmlStringBuffer); } -char * TreeImportExport::readXmlFile(const WCHAR * xmlFilePath) +bool TreeImportExport::readXmlFile(TiXmlDocument& doc, const WCHAR * xmlFilePath) { - FILE * pFile = 0; - long lSize = 0; - char * buffer = 0; + bool success = false; - if (_wfopen_s(&pFile, xmlFilePath, L"r") == NULL) + FILE * pFile = 0; + if (_wfopen_s(&pFile, xmlFilePath, L"rb") == 0) { - fseek(pFile, 0, SEEK_END); - lSize = ftell(pFile); - fseek(pFile, 0, SEEK_SET); - - if (lSize > 2) - { - buffer = new char[lSize + sizeof(char)]; - - ZeroMemory(buffer, lSize + sizeof(char)); - - fread(buffer, sizeof(char), lSize, pFile); - - if (!feof(pFile) || ferror(pFile)) - { - delete [] buffer; - buffer = 0; - } - } - + success = doc.LoadFile(pFile); fclose (pFile); - return buffer; - } - else - { - return 0; } + + return success; } -bool TreeImportExport::saveXmlToFile(TiXmlDocument doc, const WCHAR * xmlFilePath) +bool TreeImportExport::saveXmlToFile(const TiXmlDocument& doc, const WCHAR * xmlFilePath) { FILE * pFile = 0; - - if (_wfopen_s(&pFile, xmlFilePath, L"w") == NULL) + if (_wfopen_s(&pFile, xmlFilePath, L"wb") == 0) { doc.Print(pFile); - fclose (pFile); + fclose(pFile); return true; } else { return false; } } -void TreeImportExport::addModuleListToRootElement( TiXmlElement * rootElement, std::map & moduleList ) +void TreeImportExport::addModuleListToRootElement(TiXmlElement * rootElement, const std::map & moduleList) { - std::map::iterator mapIt; - std::map::iterator mapIt2; - ImportModuleThunk * importModuleThunk = 0; - ImportThunk * importThunk = 0; - - TiXmlElement * moduleElement; - TiXmlElement * importElement; - - for ( mapIt = moduleList.begin() ; mapIt != moduleList.end(); mapIt++ ) + std::map::const_iterator it_mod; + for(it_mod = moduleList.begin(); it_mod != moduleList.end(); it_mod++) { - importModuleThunk = &((*mapIt).second); + const ImportModuleThunk& importModuleThunk = it_mod->second; - moduleElement = getModuleXmlElement(importModuleThunk); + TiXmlElement* moduleElement = getModuleXmlElement(&importModuleThunk); - for ( mapIt2 = (*mapIt).second.thunkList.begin() ; mapIt2 != (*mapIt).second.thunkList.end(); mapIt2++ ) + std::map::const_iterator it_thunk; + for(it_thunk = importModuleThunk.thunkList.begin(); it_thunk != importModuleThunk.thunkList.end(); it_thunk++) { - importThunk = &((*mapIt2).second); + const ImportThunk& importThunk = it_thunk->second; + + TiXmlElement* importElement = getImportXmlElement(&importThunk); - importElement = getImportXmlElement(importThunk); moduleElement->LinkEndChild(importElement); } rootElement->LinkEndChild(moduleElement); } } TiXmlElement * TreeImportExport::getModuleXmlElement(const ImportModuleThunk * importModuleThunk) { TiXmlElement * moduleElement = new TiXmlElement("module"); StringConversion::ToASCII(importModuleThunk->moduleName, xmlStringBuffer, _countof(xmlStringBuffer)); moduleElement->SetAttribute("filename", xmlStringBuffer); ConvertDwordPtrToString(importModuleThunk->getFirstThunk()); - moduleElement->SetAttribute("first_thunk_rva",xmlStringBuffer); + moduleElement->SetAttribute("first_thunk_rva", xmlStringBuffer); return moduleElement; } TiXmlElement * TreeImportExport::getImportXmlElement(const ImportThunk * importThunk) { TiXmlElement * importElement = 0; if (importThunk->valid) { importElement = new TiXmlElement("import_valid"); if(importThunk->name[0] != '\0') { - importElement->SetAttribute("name",importThunk->name); + importElement->SetAttribute("name", importThunk->name); } ConvertWordToString(importThunk->ordinal); importElement->SetAttribute("ordinal",xmlStringBuffer); ConvertWordToString(importThunk->hint); importElement->SetAttribute("hint",xmlStringBuffer); ConvertBoolToString(importThunk->suspect); importElement->SetAttribute("suspect", xmlStringBuffer); } else { importElement = new TiXmlElement("import_invalid"); } ConvertDwordPtrToString(importThunk->rva); importElement->SetAttribute("iat_rva", xmlStringBuffer); ConvertDwordPtrToString(importThunk->apiAddressVA); - importElement->SetAttribute("address_va",xmlStringBuffer); + importElement->SetAttribute("address_va", xmlStringBuffer); return importElement; } void TreeImportExport::ConvertBoolToString(const bool boolValue) { if (boolValue) { strcpy_s(xmlStringBuffer, "1"); } else { strcpy_s(xmlStringBuffer, "0"); } } bool TreeImportExport::ConvertStringToBool(const char * strValue) { if (strValue) { if (strValue[0] == '1') { return true; } } return false; } void TreeImportExport::ConvertDwordPtrToString(const DWORD_PTR dwValue) { sprintf_s(xmlStringBuffer, PRINTF_DWORD_PTR_FULL_S, dwValue); } DWORD_PTR TreeImportExport::ConvertStringToDwordPtr(const char * strValue) { DWORD_PTR result = 0; if (strValue) { #ifdef _WIN64 result = _strtoi64(strValue, NULL, 16); #else result = strtoul(strValue, NULL, 16); #endif } return result; } void TreeImportExport::ConvertWordToString(const WORD dwValue) { sprintf_s(xmlStringBuffer, "%04X", dwValue); } WORD TreeImportExport::ConvertStringToWord(const char * strValue) { WORD result = 0; if (strValue) { result = (WORD)strtoul(strValue, NULL, 16); } return result; } -void TreeImportExport::parseAllElementModules( TiXmlElement * targetElement, std::map & moduleList ) +void TreeImportExport::parseAllElementModules(TiXmlElement * targetElement, std::map & moduleList) { - TiXmlElement * moduleElement = 0; ImportModuleThunk importModuleThunk; - const char * filename = 0; - for(moduleElement = targetElement->FirstChildElement(); moduleElement; moduleElement = moduleElement->NextSiblingElement()) + for(TiXmlElement * moduleElement = targetElement->FirstChildElement(); moduleElement; moduleElement = moduleElement->NextSiblingElement()) { - filename = moduleElement->Attribute("filename"); - + const char * filename = moduleElement->Attribute("filename"); if (filename) { StringConversion::ToUTF16(filename, importModuleThunk.moduleName, _countof(importModuleThunk.moduleName)); importModuleThunk.firstThunk = ConvertStringToDwordPtr(moduleElement->Attribute("first_thunk_rva")); importModuleThunk.thunkList.clear(); - parseAllElementImports(moduleElement, &importModuleThunk); - moduleList.insert(std::pair(importModuleThunk.firstThunk, importModuleThunk)); - + moduleList[importModuleThunk.firstThunk] = importModuleThunk; } } } -void TreeImportExport::parseAllElementImports( TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk ) +void TreeImportExport::parseAllElementImports(TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk) { - TiXmlElement * importElement = 0; ImportThunk importThunk; - const char * temp = 0; - for(importElement = moduleElement->FirstChildElement(); importElement; importElement = importElement->NextSiblingElement()) + for(TiXmlElement * importElement = moduleElement->FirstChildElement(); importElement; importElement = importElement->NextSiblingElement()) { - temp = importElement->Value(); + const char * temp = importElement->Value(); if (!strcmp(temp, "import_valid")) { temp = importElement->Attribute("name"); if (temp) { strcpy_s(importThunk.name, temp); } else { importThunk.name[0] = 0; } wcscpy_s(importThunk.moduleName, importModuleThunk->moduleName); importThunk.suspect = ConvertStringToBool(importElement->Attribute("suspect")); importThunk.ordinal = ConvertStringToWord(importElement->Attribute("ordinal")); importThunk.hint = ConvertStringToWord(importElement->Attribute("hint")); importThunk.valid = true; } else { importThunk.valid = false; importThunk.suspect = true; } importThunk.apiAddressVA = ConvertStringToDwordPtr(importElement->Attribute("address_va")); importThunk.rva = ConvertStringToDwordPtr(importElement->Attribute("iat_rva")); if (importThunk.rva != 0) { - importModuleThunk->thunkList.insert(std::pair(importThunk.rva, importThunk)); + importModuleThunk->thunkList[importThunk.rva] = importThunk; } } } diff --git a/Scylla/TreeImportExport.h b/Scylla/TreeImportExport.h index 443807a..61362d1 100644 --- a/Scylla/TreeImportExport.h +++ b/Scylla/TreeImportExport.h @@ -1,38 +1,42 @@ #pragma once #include #include "ProcessLister.h" #include "Thunks.h" #include class TreeImportExport { public: - bool exportTreeList(const WCHAR * targetXmlFile, std::map & moduleList, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT); - bool importTreeList(const WCHAR * targetXmlFile, std::map & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT); + TreeImportExport(const WCHAR * targetXmlFile); + + bool exportTreeList(const std::map & moduleList, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT) ; + bool importTreeList(std::map & moduleList, DWORD_PTR * addressOEP, DWORD_PTR * addressIAT, DWORD * sizeIAT); private: - char xmlStringBuffer[100]; + WCHAR xmlPath[MAX_PATH]; + + char xmlStringBuffer[MAX_PATH]; + + void setTargetInformation(TiXmlElement * rootElement, const Process * process, DWORD_PTR addressOEP, DWORD_PTR addressIAT, DWORD sizeIAT); + void addModuleListToRootElement(TiXmlElement * rootElement, const std::map & moduleList); + + void parseAllElementModules(TiXmlElement * targetElement, std::map & moduleList); + void parseAllElementImports(TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk); - void addModuleListToRootElement( TiXmlElement * rootElement, std::map & moduleList ); TiXmlElement * getModuleXmlElement(const ImportModuleThunk * importModuleThunk); TiXmlElement * getImportXmlElement(const ImportThunk * importThunk); - bool saveXmlToFile(TiXmlDocument doc, const WCHAR * xmlFilePath); - char * readXmlFile(const WCHAR * xmlFilePath); - - void setTargetInformation(TiXmlElement * rootElement, const Process * process, const DWORD_PTR addressOEP, const DWORD_PTR addressIAT, const DWORD sizeIAT); + bool saveXmlToFile(const TiXmlDocument& doc, const WCHAR * xmlFilePath); + bool readXmlFile(TiXmlDocument& doc, const WCHAR * xmlFilePath); void ConvertBoolToString(const bool boolValue); void ConvertWordToString(const WORD dwValue); void ConvertDwordPtrToString(const DWORD_PTR dwValue); DWORD_PTR ConvertStringToDwordPtr(const char * strValue); WORD ConvertStringToWord(const char * strValue); bool ConvertStringToBool(const char * strValue); - - void parseAllElementModules( TiXmlElement * targetElement, std::map & moduleList ); - void parseAllElementImports( TiXmlElement * moduleElement, ImportModuleThunk * importModuleThunk ); };