From 56c778b37229606818487b87e96d74b1db858817 Mon Sep 17 00:00:00 2001
From: Aliz Hammond <ahammond@countercept.com>
Date: Thu, 19 Sep 2019 19:14:21 +0800
Subject: [PATCH] Fix unicode conversion bug in section names; convert
 statistics to MB; bail if symsrv.dll is not found

---
 UI/UI.cpp                             | 23 +++++++++++++++--------
 UI/findPFNDatabase.cpp                | 22 ++++++++++++++++++----
 injectionUtils/moduleManipulation.cpp |  1 +
 3 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/UI/UI.cpp b/UI/UI.cpp
index d950c9d..a518c14 100644
--- a/UI/UI.cpp
+++ b/UI/UI.cpp
@@ -2,7 +2,8 @@
 #include <windows.h>
 #include <psapi.h>
 #include <tlhelp32.h>
-#include <iostream>
+#include <locale>
+#include <codecvt>
 
 #include <vector>
 #include <map>
@@ -39,10 +40,12 @@ public:
 modifiedPage::modifiedPage(DWORD newProcessID, wchar_t* newModuleName, void* newPageBase, BYTE newSectionName[8], unsigned long long newSectionOffset)
 	: processID(newProcessID), moduleName(newModuleName), pageBase((unsigned long long)newPageBase), sectionName(L""), sectionOffset(newSectionOffset)
 {
-	wchar_t sectionNameCleaned[9];
-	memset(sectionNameCleaned, 0, 9 * sizeof(wchar_t));
-	wsprintf(sectionNameCleaned, L"%.8s", newSectionName);
-	sectionName.append(sectionNameCleaned);
+	std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
+	std::wstring secName = std::wstring(converter.from_bytes((char*)newSectionName, (char*)&newSectionName[7]));
+	size_t nullPos = secName.find(L'\0');
+	if (nullPos != secName.npos)
+		secName = secName.substr(0, nullPos);
+ 	sectionName.append(secName);
 }
 
 BOOL EnableDebugPrivilege(BOOL bEnable)
@@ -272,12 +275,16 @@ int main()
 
 	printf("Scan took %dms\n", (end - start));
 
-	// Print some stats and the results.
-	printf("Scanned %d pages, ignored %d NX pages (total %d). Found %d modified pages.\n", stat.scannedPages, stat.ignoredPagesNX, stat.ignoredPagesNX + stat.scannedPages, stat.modifiedPages);
+	// Print some stats and the results. Each page is 4KB, so 256 pages is 1MB.
+	printf("Scanned %d pages (%d MB), ignored %d NX pages (%d MB), totalling %d pages  (%d MB). Found %d modified pages (%.02f MB).\n", 
+		stat.scannedPages, stat.scannedPages / 256,
+		stat.ignoredPagesNX, stat.scannedPages / 256,
+		stat.ignoredPagesNX + stat.scannedPages, (stat.ignoredPagesNX + stat.scannedPages) / 256,
+		stat.modifiedPages, ((float)stat.modifiedPages) / 256);
 	for (unsigned int n = 0; n < results.size(); n++)
 	{
 		modifiedPage thisModifiedPage = results[n];
-		printf("PID %04d module '%ls', section %S, offset 0x%08llux\n", thisModifiedPage.processID, thisModifiedPage.moduleName.c_str(), thisModifiedPage.sectionName.c_str(), thisModifiedPage.sectionOffset);
+		printf("PID %04d module '%ls', section %S, offset 0x%08llx\n", thisModifiedPage.processID, thisModifiedPage.moduleName.c_str(), thisModifiedPage.sectionName.c_str(), thisModifiedPage.sectionOffset);
 	/*
 		HANDLE toScanHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, thisModifiedPage.processID);
 		if (toScanHandle == NULL)
diff --git a/UI/findPFNDatabase.cpp b/UI/findPFNDatabase.cpp
index 8eab203..18c8c22 100644
--- a/UI/findPFNDatabase.cpp
+++ b/UI/findPFNDatabase.cpp
@@ -24,7 +24,7 @@ unsigned long long findPFNDatabase()
 	wchar_t curPath[MAX_PATH];
 	GetCurrentDirectory(MAX_PATH, curPath);
 	std::wstring symPath(L"");
-	symPath.append(L"symsrv*symsrv.dll*");
+	symPath.append(L"srv*");
 	symPath.append(curPath);
 	symPath.append(L"*http://msdl.microsoft.com/download/symbols");
 
@@ -40,13 +40,27 @@ unsigned long long findPFNDatabase()
 		return false;
 	}
 
+	HMODULE symSrv = LoadLibrary(L"symsrv.dll");
+	if (symSrv == NULL)
+	{
+		printf("symsrv.dll not found. Please install a copy in your PATH.\n");
+		return false;
+	}
+
 	printf("Loading PDBs..\n");
 	hr = g_pDiaDataSource->loadDataForExe(exeFilename.c_str(), symPath.c_str(), NULL);
-	if (FAILED(hr)) 
+	if (FAILED(hr))
 	{
-		printf("loadDataForExe failed for file '%ls' - HRESULT is %08X\n", exeFilename.c_str(), hr);
 		if (hr == E_PDB_NOT_FOUND)
-			printf("This is E_PDB_NOT_FOUND. Check that you have internet connectivity, and the correct symbol server configured.\n");
+		{
+			printf("DIA returned E_PDB_NOT_FOUND. Check that you have internet connectivity, and the correct symbol server configured.\n");
+			printf("The symbol path used was '%S'.\n", symPath.c_str());
+		}
+		else
+		{
+			printf("loadDataForExe failed for file '%ls' - HRESULT is %08X\n", exeFilename.c_str(), hr);
+		}
+
 		return false;
 	}
 	printf("Loading PDBs complete.\n");
diff --git a/injectionUtils/moduleManipulation.cpp b/injectionUtils/moduleManipulation.cpp
index 7f7c462..6bd7bf9 100644
--- a/injectionUtils/moduleManipulation.cpp
+++ b/injectionUtils/moduleManipulation.cpp
@@ -4,6 +4,7 @@
 #include <sstream>
 #include <vector>
 #include <map>
+#include <iostream>
 
 #include "moduleManipulation.h"
 #include "public.h"