Skip to content

Commit 6f35e5b

Browse files
BeneficialCodeBeneficialCode
committed
PE dump bugs fixed.
1 parent a9dd48b commit 6f35e5b

File tree

7 files changed

+293
-57
lines changed

7 files changed

+293
-57
lines changed

‎PEParser/PEParser.cpp‎

Lines changed: 59 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,35 @@
99
#pragma comment(lib,"imagehlp")
1010

1111

12-
PEParser::PEParser(const wchar_t* path) :_path(path) {
13-
_hFile = ::CreateFile(path, GENERIC_READ,
14-
FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
15-
if (_hFile == INVALID_HANDLE_VALUE)
16-
return;
17-
::GetFileSizeEx(_hFile, &_fileSize);
18-
_hMemMap = ::CreateFileMapping(_hFile, nullptr, PAGE_READONLY, 0, 0, nullptr);
19-
if (!_hMemMap)
20-
return;
21-
22-
_address = (PBYTE)::MapViewOfFile(_hMemMap, FILE_MAP_READ, 0, 0, 0);
23-
if (!_address)
24-
return;
12+
PEParser::PEParser(const wchar_t* path,bool isScylla) :_path(path) {
13+
if (isScylla) {
14+
_hFile = ::CreateFile(path, GENERIC_READ|GENERIC_WRITE,
15+
FILE_SHARE_READ| FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, 0, nullptr);
16+
if (_hFile == INVALID_HANDLE_VALUE)
17+
return;
18+
::GetFileSizeEx(_hFile, &_fileSize);
19+
_hMemMap = ::CreateFileMapping(_hFile, nullptr, PAGE_READWRITE, 0, 0, nullptr);
20+
if (!_hMemMap)
21+
return;
22+
23+
_address = (PBYTE)::MapViewOfFile(_hMemMap, FILE_MAP_READ|FILE_MAP_WRITE, 0, 0, 0);
24+
if (!_address)
25+
return;
26+
}
27+
else {
28+
_hFile = ::CreateFile(path, GENERIC_READ,
29+
FILE_SHARE_READ, nullptr, OPEN_EXISTING, 0, nullptr);
30+
if (_hFile == INVALID_HANDLE_VALUE)
31+
return;
32+
::GetFileSizeEx(_hFile, &_fileSize);
33+
_hMemMap = ::CreateFileMapping(_hFile, nullptr, PAGE_READONLY, 0, 0, nullptr);
34+
if (!_hMemMap)
35+
return;
36+
37+
_address = (PBYTE)::MapViewOfFile(_hMemMap, FILE_MAP_READ, 0, 0, 0);
38+
if (!_address)
39+
return;
40+
}
2541

2642
CheckValidity();
2743
if (IsValid()) {
@@ -665,16 +681,22 @@ void PEParser::AlignAllSectionHeaders() {
665681
DWORD fileAlignment = GetFileAlignment();
666682
DWORD newFileSize = 0;
667683

684+
685+
std::sort(_PESections.begin(), _PESections.end(), [](const PEFileSection& d1, const PEFileSection& d2) {
686+
return d1._sectionHeader.PointerToRawData < d2._sectionHeader.PointerToRawData;
687+
});
688+
668689
newFileSize = _dosHeader->e_lfanew + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER) +
669-
_ntHeader->FileHeader.SizeOfOptionalHeader * sizeof(IMAGE_SECTION_HEADER);
690+
_ntHeader->FileHeader.SizeOfOptionalHeader + GetSectionCount() * sizeof(IMAGE_SECTION_HEADER);
670691

671692
for (int i = 0; i < GetSectionCount(); ++i) {
672-
sections[i].VirtualAddress = AlignValue(sections[i].VirtualAddress, sectionAlignment);
673-
sections[i].Misc.VirtualSize = AlignValue(sections[i].Misc.VirtualSize, sectionAlignment);
693+
_PESections[i]._sectionHeader.VirtualAddress = AlignValue(sections[i].VirtualAddress, sectionAlignment);
694+
_PESections[i]._sectionHeader.Misc.VirtualSize = AlignValue(sections[i].Misc.VirtualSize, sectionAlignment);
674695

675-
sections[i].PointerToRawData = AlignValue(newFileSize, fileAlignment);
696+
_PESections[i]._sectionHeader.PointerToRawData = AlignValue(newFileSize, fileAlignment);
697+
_PESections[i]._sectionHeader.SizeOfRawData = AlignValue(_PESections[i]._dataSize, fileAlignment);
676698

677-
newFileSize = sections[i].PointerToRawData + sections[i].SizeOfRawData;
699+
newFileSize = _PESections[i]._sectionHeader.PointerToRawData + _PESections[i]._sectionHeader.SizeOfRawData;
678700
}
679701

680702
std::sort(_PESections.begin(), _PESections.end(), [](const PEFileSection& d1, const PEFileSection& d2) {
@@ -717,6 +739,8 @@ void PEParser::FixPEHeader() {
717739
_opt64->NumberOfRvaAndSizes = IMAGE_NUMBEROF_DIRECTORY_ENTRIES;
718740
_fileHeader->SizeOfOptionalHeader = sizeof(IMAGE_OPTIONAL_HEADER64);
719741

742+
_opt64->SizeOfImage = GetSectionHeaderBasedSizeOfImage();
743+
720744
if (_moduleBase) {
721745
_opt64->ImageBase = _moduleBase;
722746
}
@@ -847,6 +871,7 @@ void PEParser::GetPESections() {
847871

848872
for (WORD i = 0; i < count; i++) {
849873
_PESections[i]._normalSize = _sections[i].Misc.VirtualSize;
874+
_PESections[i]._dataSize = _sections[i].Misc.VirtualSize;
850875
offset = _sections[i].VirtualAddress;
851876
GetPESectionData(offset, _PESections[i]);
852877
}
@@ -938,8 +963,7 @@ bool PEParser::SavePEFileToDisk(const WCHAR* pNewFile) {
938963
writeSize = sizeof(IMAGE_SECTION_HEADER);
939964
auto sections = GetSections();
940965
for (WORD i = 0; i < GetSectionCount(); i++) {
941-
auto section = sections[i];
942-
if (!WriteMemoryToFile(_hFile, fileOffset, writeSize, &section)) {
966+
if (!WriteMemoryToFile(_hFile, fileOffset, writeSize, &_PESections[i]._sectionHeader)) {
943967
ret = false;
944968
break;
945969
}
@@ -948,12 +972,11 @@ bool PEParser::SavePEFileToDisk(const WCHAR* pNewFile) {
948972

949973

950974
for (WORD i = 0; i < GetSectionCount(); i++) {
951-
auto section = sections[i];
952-
if (!section.PointerToRawData)
975+
if (!_PESections[i]._sectionHeader.PointerToRawData)
953976
continue;
954977

955-
if (section.PointerToRawData > fileOffset) {
956-
writeSize = section.PointerToRawData - fileOffset;
978+
if (_PESections[i]._sectionHeader.PointerToRawData > fileOffset) {
979+
writeSize = _PESections[i]._sectionHeader.PointerToRawData - fileOffset;
957980

958981
if (!WriteZeroMemoryToFile(_hFile, fileOffset, writeSize)) {
959982
ret = false;
@@ -962,15 +985,24 @@ bool PEParser::SavePEFileToDisk(const WCHAR* pNewFile) {
962985
fileOffset += writeSize;
963986
}
964987

965-
writeSize = section.SizeOfRawData;
988+
writeSize = _PESections[i]._sectionHeader.SizeOfRawData;
966989

967990
if (writeSize) {
968-
BYTE* pData = GetFileAddress(section.PointerToRawData);
969-
if (!WriteMemoryToFile(_hFile, section.PointerToRawData, writeSize, pData)) {
991+
if (!WriteMemoryToFile(_hFile, _PESections[i]._sectionHeader.PointerToRawData,
992+
writeSize, _PESections[i]._pData)) {
970993
ret = false;
971994
break;
972995
}
973996
fileOffset += writeSize;
997+
998+
if (_PESections[i]._dataSize < _PESections[i]._sectionHeader.SizeOfRawData) {
999+
writeSize = _PESections[i]._sectionHeader.SizeOfRawData - _PESections[i]._dataSize;
1000+
if (!WriteZeroMemoryToFile(_hFile, fileOffset, writeSize)) {
1001+
ret = false;
1002+
break;
1003+
}
1004+
fileOffset += writeSize;
1005+
}
9741006
}
9751007
}
9761008

‎PEParser/PEParser.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ struct RelocInfo {
166166

167167
class PEParser {
168168
public:
169-
explicit PEParser(const wchar_t* path);
169+
explicit PEParser(const wchar_t* path,bool isScylla = false);
170170
~PEParser();
171171
PEParser(void* base);
172172

‎WinArk/ImportRebuilder.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
class ImportRebuilder: public PEParser{
1010
public:
11-
ImportRebuilder(const WCHAR* file): PEParser(file) {
11+
ImportRebuilder(const WCHAR* file): PEParser(file,true) {
1212
}
1313
bool RebuildImportTable(const WCHAR* newFilePath, std::map<DWORD_PTR, ImportModuleThunk>& moduleThunkMap);
1414
void EnableOFTSupport();

‎WinArk/ProcessAccessHelper.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,8 +362,12 @@ bool ProcessAccessHelper::GetProcessModules(HANDLE hProcess, std::vector<ModuleI
362362

363363
for (auto& info : infos) {
364364
std::wstring path = info.get()->Path;
365+
std::wstring name = info.get()->Name;
365366
if (path.empty())
366367
continue;
368+
if (_wcsicmp(L"kernelbase.dll", name.c_str()) == 0)
369+
continue;
370+
367371

368372
module._modBaseAddr = (DWORD_PTR)info.get()->Base;
369373
module._modBaseSize = info.get()->ModuleSize;

‎WinArk/ScyllaDlg.cpp‎

Lines changed: 176 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "Architecture.h"
55
#include "IATSearcher.h"
66
#include <PEParser.h>
7+
#include "ImportRebuilder.h"
78

89

910
CScyllaDlg::CScyllaDlg(const WinSys::ProcessManager& pm, ProcessInfoEx& px)
@@ -347,13 +348,15 @@ bool CScyllaDlg::GetCurrentDefaultDumpFilename(WCHAR* pBuffer, size_t bufSize) {
347348
wcscpy_s(fullPath, m_px.GetExecutablePath().c_str());
348349
}
349350

350-
WCHAR* pSlash = wcsrchr(fullPath, L'\\');
351-
if (pSlash) {
352-
pSlash++;
351+
WCHAR* pTemp = wcsrchr(fullPath, L'\\');
352+
if (pTemp) {
353+
pTemp++;
354+
355+
wcscpy_s(pBuffer, bufSize, pTemp);
353356

354-
wcscpy_s(pBuffer, bufSize, pSlash);
355-
if (pSlash) {
356-
*pSlash = 0;
357+
pTemp = wcsrchr(pBuffer, L'.');
358+
if (pTemp) {
359+
*pTemp = 0;
357360
if (ProcessAccessHelper::_pSelectedModule) {
358361
wcscat_s(pBuffer, bufSize, L"_dump.dll");
359362
}
@@ -396,4 +399,171 @@ bool CScyllaDlg::ShowFileDialog(WCHAR* pSelectedFile, bool save, const WCHAR* pD
396399
return 0 != GetSaveFileName(&ofn);
397400
else
398401
return 0 != GetOpenFileName(&ofn);
402+
}
403+
404+
void CScyllaDlg::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl) {
405+
FixDumpHandler();
406+
}
407+
408+
void CScyllaDlg::FixDumpHandler() {
409+
if (_treeImports.GetCount() < 2) {
410+
return;
411+
}
412+
413+
WCHAR newFilePath[MAX_PATH] = { 0 };
414+
WCHAR selectedFilePath[MAX_PATH] = { 0 };
415+
const WCHAR* pFileFilter = nullptr;
416+
DWORD_PTR modBase = 0;
417+
DWORD_PTR entryPoint = _oepAddress.GetValue();
418+
419+
if (ProcessAccessHelper::_pSelectedModule) {
420+
modBase = ProcessAccessHelper::_pSelectedModule->_modBaseAddr;
421+
pFileFilter = s_FilterDll;
422+
}
423+
else {
424+
modBase = ProcessAccessHelper::_targetImageBase;
425+
pFileFilter = s_FilterExe;
426+
}
427+
428+
GetCurrentModulePath(_text, _countof(_text));
429+
if (ShowFileDialog(selectedFilePath, false, nullptr, pFileFilter, nullptr, _text)) {
430+
wcscpy_s(newFilePath, selectedFilePath);
431+
432+
const WCHAR* pExtension = nullptr;
433+
WCHAR* pDot = wcsrchr(newFilePath, L'.');
434+
if (pDot) {
435+
*pDot = L'\0';
436+
pExtension = selectedFilePath + (pDot - newFilePath);
437+
}
438+
439+
wcscat_s(newFilePath, L"_SCY");
440+
441+
if (pExtension) {
442+
wcscat_s(newFilePath, pExtension);
443+
}
444+
445+
ImportRebuilder rebuilder(selectedFilePath);
446+
447+
rebuilder.RebuildImportTable(newFilePath, _importsHandling.m_ModuleMap);
448+
}
449+
}
450+
451+
void CScyllaDlg::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl) {
452+
PERebuildHandler();
453+
}
454+
455+
void CScyllaDlg::PERebuildHandler() {
456+
DWORD newSize = 0;
457+
WCHAR selectedFilePath[MAX_PATH] = { 0 };
458+
459+
GetCurrentModulePath(_text, _countof(_text));
460+
if (ShowFileDialog(selectedFilePath, false, nullptr, s_FilterExeDll, nullptr, _text)) {
461+
DWORD fileSize = ProcessAccessHelper::GetFileSize(selectedFilePath);
462+
463+
PEParser parser(selectedFilePath, true);
464+
465+
if (!parser.IsValid()) {
466+
return;
467+
}
468+
469+
parser.GetPESections();
470+
parser.SetDefaultFileAligment();
471+
parser.AlignAllSectionHeaders();
472+
parser.FixPEHeader();
473+
parser.SavePEFileToDisk(selectedFilePath);
474+
}
475+
}
476+
477+
CTreeItem CScyllaDlg::FindTreeItem(CPoint pt, bool screenCoordinates) {
478+
if (screenCoordinates) {
479+
_treeImports.ScreenToClient(&pt);
480+
}
481+
UINT flags;
482+
CTreeItem over = _treeImports.HitTest(pt, &flags);
483+
if (over) {
484+
if (!(flags & TVHT_ONITEM))
485+
{
486+
over.m_hTreeItem = NULL;
487+
}
488+
}
489+
return over;
490+
}
491+
492+
LRESULT CScyllaDlg::OnTreeImportsDoubleClick(const NMHDR* pnmh) {
493+
if (_treeImports.GetCount() < 1)
494+
return 0;
495+
496+
CTreeItem over = FindTreeItem(CPoint(GetMessagePos()), true);
497+
if (over && _importsHandling.IsImport(over)) {
498+
// todo: Pick API Handler
499+
500+
}
501+
502+
return 0;
503+
}
504+
505+
LRESULT CScyllaDlg::OnTreeImportsKeyDown(const NMHDR* pnmh) {
506+
const NMTVKEYDOWN* tkd = (NMTVKEYDOWN*)pnmh;
507+
switch (tkd->wVKey)
508+
{
509+
case VK_RETURN:
510+
{
511+
CTreeItem selected = _treeImports.GetFocusItem();
512+
if (!selected.IsNull() && _importsHandling.IsImport(selected))
513+
{
514+
// to do: Pick API handler
515+
}
516+
}
517+
return 1;
518+
case VK_DELETE:
519+
DeleteSelectedImportsHandler();
520+
return 1;
521+
}
522+
523+
SetMsgHandled(FALSE);
524+
return 0;
525+
}
526+
527+
void CScyllaDlg::DeleteSelectedImportsHandler() {
528+
CTreeItem selected = _treeImports.GetFirstSelectedItem();
529+
while (!selected.IsNull())
530+
{
531+
if (_importsHandling.IsModule(selected))
532+
{
533+
_importsHandling.CutModule(selected);
534+
}
535+
else
536+
{
537+
_importsHandling.CutImport(selected);
538+
}
539+
selected = _treeImports.GetNextSelectedItem(selected);
540+
}
541+
UpdateStatusBar();
542+
}
543+
544+
void CScyllaDlg::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl) {
545+
ShowInvalidImportsHandler();
546+
}
547+
548+
void CScyllaDlg::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl) {
549+
ShowSuspectImportsHandler();
550+
}
551+
552+
void CScyllaDlg::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl) {
553+
ClearImportsHandler();
554+
}
555+
556+
void CScyllaDlg::ShowInvalidImportsHandler() {
557+
_importsHandling.SelectImports(true, false);
558+
GotoDlgCtrl(_treeImports);
559+
}
560+
561+
void CScyllaDlg::ShowSuspectImportsHandler() {
562+
_importsHandling.SelectImports(false, true);
563+
GotoDlgCtrl(_treeImports);
564+
}
565+
566+
void CScyllaDlg::ClearImportsHandler() {
567+
_importsHandling.ClearAllImports();
568+
UpdateStatusBar();
399569
}

0 commit comments

Comments
 (0)