Add Process::Sections method

- I will need these soon...
visuve · Apr 1, 2024 · dc968b2 · dc968b2
1 parent 9f04ba5
commit dc968b2
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 10 deletions.
diff --git a/HackLib/Process.cpp b/HackLib/Process.cpp
@@ -144,7 +144,7 @@ MODULEENTRY32W Process::FindModuleEntry(std::wstring_view name) const
 	return System::ModuleEntryByName(_pid, name);
 }
 
-IMAGE_NT_HEADERS Process::NtHeader() const
+std::pair<IMAGE_DOS_HEADER, IMAGE_NT_HEADERS> Process::Headers() const
 {
 	const auto dosHeader = Read<IMAGE_DOS_HEADER>(_baseAddress);
 
@@ -153,33 +153,51 @@ IMAGE_NT_HEADERS Process::NtHeader() const
 		throw LogicException("Invalid DOS header");
 	}
 
-	const auto ntHeaders = Read<IMAGE_NT_HEADERS>(_baseAddress + dosHeader.e_lfanew);
+	const auto ntHeader = Read<IMAGE_NT_HEADERS>(_baseAddress + dosHeader.e_lfanew);
 
-	if (ntHeaders.Signature != IMAGE_NT_SIGNATURE) 
+	if (ntHeader.Signature != IMAGE_NT_SIGNATURE) 
 	{
 		throw LogicException("Invalid NT headers");
 	}
 
 #ifdef _WIN64
-	if (ntHeaders.OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC)
+	if (ntHeader.OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC)
 #else
-	if (ntHeaders.OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC)
+	if (ntHeader.OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC)
 #endif
 	{
 		throw LogicException("Invalid magic");
 	}
 
-	if (!(ntHeaders.FileHeader.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE))
+	if (!(ntHeader.FileHeader.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE))
 	{	
 		throw LogicException("Not an executable");
 	}
 
-	return ntHeaders;
+	return { dosHeader, ntHeader };
+}
+
+std::vector<IMAGE_SECTION_HEADER> Process::Sections() const
+{
+	const auto [dosHeader, ntHeader] = Headers();
+
+	Pointer sectionStart = _baseAddress + dosHeader.e_lfanew + sizeof(IMAGE_NT_HEADERS);
+
+	std::vector<IMAGE_SECTION_HEADER> sections(ntHeader.FileHeader.NumberOfSections);
+
+	if (sections.empty())
+	{
+		throw LogicException("No sections");
+	}
+
+	Read(sectionStart, sections.data(), sections.size() * sizeof(IMAGE_SECTION_HEADER));
+
+	return sections;
 }
 
 IMAGE_IMPORT_DESCRIPTOR Process::FindImportDescriptor(std::string_view moduleName) const
 {
-	IMAGE_NT_HEADERS ntHeader = NtHeader();
+	const auto [_, ntHeader] = Headers();
 
 	IMAGE_DATA_DIRECTORY importEntry = ntHeader.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
 

diff --git a/HackLib/Process.hpp b/HackLib/Process.hpp
@@ -215,7 +215,8 @@ class Process
 
 	MODULEENTRY32W FindModuleEntry(std::wstring_view name) const;
 
-	IMAGE_NT_HEADERS NtHeader() const;
+	std::pair<IMAGE_DOS_HEADER, IMAGE_NT_HEADERS> Headers() const;
+	std::vector<IMAGE_SECTION_HEADER> Sections() const;
 
 	IMAGE_IMPORT_DESCRIPTOR FindImportDescriptor(std::string_view moduleName) const;
 	Pointer FindImportEntry(IMAGE_IMPORT_DESCRIPTOR iid, std::string_view functionName) const;

diff --git a/HackLibTests/ProcessTests.cpp b/HackLibTests/ProcessTests.cpp
@@ -5,7 +5,36 @@ TEST(ProcessTests, Header)
 	DWORD pid = GetCurrentProcessId();
 	Process current(pid);
 
-	EXPECT_EQ(current.NtHeader().OptionalHeader.Subsystem, WORD(IMAGE_SUBSYSTEM_WINDOWS_CUI));
+	const auto [dosHeader, ntHeader] = current.Headers();
+
+	EXPECT_NE(ntHeader.OptionalHeader.Subsystem, WORD(IMAGE_SUBSYSTEM_WINDOWS_GUI));
+	EXPECT_EQ(ntHeader.OptionalHeader.Subsystem, WORD(IMAGE_SUBSYSTEM_WINDOWS_CUI));
+}
+
+TEST(ProcessTests, Sections)
+{
+	DWORD pid = GetCurrentProcessId();
+	Process current(pid);
+
+	auto sections = current.Sections();
+
+#ifdef _WIN64
+	EXPECT_EQ(sections.size(), size_t(7));
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[0].Name), ".text");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[1].Name), ".rdata");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[2].Name), ".data");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[3].Name), ".pdata");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[4].Name), "_RDATA");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[5].Name), ".rsrc");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[6].Name), ".reloc");
+#else
+	EXPECT_EQ(sections.size(), size_t(5));
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[0].Name), ".text");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[1].Name), ".rdata");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[2].Name), ".data");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[3].Name), ".rsrc");
+	EXPECT_STREQ(reinterpret_cast<char*>(sections[4].Name), ".reloc");
+#endif
 }
 
 TEST(ProcessTests, ModuleNotFound)