Merge branch 'imp_rec_refact' (Issue #91)

postprocessors/imp_rec/iat_finder.cpp

@@ -1,20 +1 @@
 #include "iat_finder.h"
-
-pesieve::IATBlock* pesieve::find_iat_block(
-	IN bool is64bit,
-	IN BYTE* vBuf,
-	IN size_t vBufSize,
-	IN const peconv::ExportsMapper* exportsMap,
-	IN OPTIONAL size_t search_offset
-)
-{
-	IATBlock* iat_block = nullptr;
-	if (is64bit) {
-		iat_block = find_iat<ULONGLONG>(true, vBuf, vBufSize, exportsMap, search_offset);
-	}
-	else {
-		iat_block = find_iat<DWORD>(false, vBuf, vBufSize, exportsMap, search_offset);
-	}
-	return iat_block;
-}
-

postprocessors/imp_rec/iat_finder.h

@@ -11,16 +11,21 @@
 
 namespace pesieve {
 
-	IATBlock* find_iat_block(
-		IN bool is64bit,
-		IN BYTE* vBuf,
-		IN size_t vBufSize,
-		IN const peconv::ExportsMapper* exportsMap,
-		IN OPTIONAL size_t search_offset
-	);
+	class ThunkFoundCallback
+	{
+	public:
+		ThunkFoundCallback()
+		{
+		}
+
+		virtual bool shouldProcessVA(ULONGLONG va) = 0;
+
+		virtual bool shouldAcceptExport(ULONGLONG va, const peconv::ExportedFunc &exp) = 0;
+	};
+	//---
 
 	template <typename FIELD_T>
-	size_t fill_iat(BYTE* vBuf, size_t vBufSize, IN const peconv::ExportsMapper* exportsMap, IN OUT IATBlock &iat)
+	size_t fill_iat(BYTE* vBuf, size_t vBufSize, IN const peconv::ExportsMapper* exportsMap, IN OUT IATBlock &iat, IN ThunkFoundCallback *callback)
 	{
 		if (!vBuf || !exportsMap || !iat.iatOffset) return 0;
 
@@ -44,16 +49,28 @@ namespace pesieve {
 				}
 				continue;
 			}
-			exp = exportsMap->find_export_by_va(*imp);
+
+			const FIELD_T imp_va = (*imp);
+
+			if (callback) {
+				if (!callback->shouldProcessVA(imp_va)) break;
+			}
+
+			//std::cout << "checking: " << std::hex << possible_rva << std::endl;
+			const peconv::ExportedFunc *exp = exportsMap->find_export_by_va(imp_va);
 			if (!exp) break;
 
+			if (callback) {
+				if (!callback->shouldAcceptExport(imp_va, *exp)) break;
+			}
+
 			is_terminated = false;
 			DWORD offset = MASK_TO_DWORD((BYTE*)imp - vBuf);
-			iat.append(offset, *imp, exp);
+			iat.append(offset, imp_va, exp);
 
 			if (!series) series = new IATThunksSeries(offset);
 			if (series) {
-				series->insert(offset, *imp);
+				series->insert(offset, imp_va);
 			}
 		}
 		if (series) {
@@ -71,9 +88,11 @@ namespace pesieve {
 	}
 
 	template <typename FIELD_T>
-	IATBlock* find_iat(bool is64bit, BYTE* vBuf, size_t vBufSize, IN const peconv::ExportsMapper* exportsMap, IN OPTIONAL size_t search_offset = 0)
+	IATBlock* find_iat(BYTE* vBuf, size_t vBufSize, IN const peconv::ExportsMapper* exportsMap, IN size_t search_offset, IN ThunkFoundCallback *callback)
 	{
-		if (!vBuf || !exportsMap) return nullptr;
+		if (!vBuf || !vBufSize ||  !exportsMap) return nullptr;
+
+		const bool is64bit = (sizeof(FIELD_T) == sizeof(DWORD)) ? false : true;
 
 		size_t max_check = vBufSize - sizeof(FIELD_T);
 		if (search_offset > vBufSize || max_check < sizeof(FIELD_T)) {
@@ -83,16 +102,25 @@ namespace pesieve {
 		for (BYTE* ptr = vBuf + search_offset; ptr < vBuf + max_check; ptr++) {
 			FIELD_T *to_check = (FIELD_T*)ptr;
 			if (!peconv::validate_ptr(vBuf, vBufSize, to_check, sizeof(FIELD_T))) break;
-			FIELD_T possible_rva = (*to_check);
-			if (possible_rva == 0) continue;
+			FIELD_T possible_va = (*to_check);
+			if (possible_va == 0) continue;
+
+			if (callback) {
+				if (!callback->shouldProcessVA(possible_va)) continue;
+			}
+
 			//std::cout << "checking: " << std::hex << possible_rva << std::endl;
-			const peconv::ExportedFunc *exp = exportsMap->find_export_by_va(possible_rva);
+			const peconv::ExportedFunc *exp = exportsMap->find_export_by_va(possible_va);
 			if (!exp) continue;
 
+			if (callback) {
+				if (!callback->shouldAcceptExport(possible_va, *exp)) continue;
+			}
+
 			DWORD iat_offset = DWORD(ptr - vBuf);
 			IATBlock *iat_block = new IATBlock(is64bit, iat_offset);
 			//validate IAT:
-			size_t _iat_size = fill_iat<FIELD_T>(vBuf, vBufSize, exportsMap, *iat_block);
+			size_t _iat_size = fill_iat<FIELD_T>(vBuf, vBufSize, exportsMap, *iat_block, callback);
 			if (_iat_size > 0) {
 				iat_block->iatSize = _iat_size;
 				return iat_block;

postprocessors/imp_rec/imp_reconstructor.cpp

@@ -152,7 +152,7 @@ bool pesieve::ImpReconstructor::isDefaultImportValid(IN const peconv::ExportsMap
 	// verify if the Import Table that is currently set is fine:
 
 	DWORD iat_offset = iat_dir->VirtualAddress;
-	IATBlock* iat_block = find_iat_block(is64bit, vBuf, vBufSize, exportsMap, iat_offset);
+	IATBlock* iat_block = this->findIATBlock(exportsMap, iat_offset);
 	if (!iat_block) {
 		//could not find any IAT Block at this IAT offset. The IAT offset may be incorrect.
 		return false;
@@ -181,13 +181,58 @@ bool pesieve::ImpReconstructor::isDefaultImportValid(IN const peconv::ExportsMap
 	return false;
 }
 
+IATBlock* pesieve::ImpReconstructor::findIATBlock(IN const peconv::ExportsMapper* exportsMap, size_t start_offset)
+{
+	if (!exportsMap) return nullptr;
+
+	// filter calls to the own exports
+	class ThunkFilterSelfCallback : public ThunkFoundCallback
+	{
+	public:
+		ThunkFilterSelfCallback(const ULONGLONG mod_start, size_t mod_size)
+			: startAddr(mod_start), endAddr(mod_start + mod_size)
+		{
+		}
+
+		virtual bool shouldProcessVA(ULONGLONG va)
+		{
+			if (va >= startAddr && va < endAddr) {
+				// the address is in the current module: this may be a call to module's own function
+				return false;
+			}
+			return true;
+		}
+
+		virtual bool shouldAcceptExport(ULONGLONG va, const peconv::ExportedFunc &exp)
+		{
+			// accept any
+			return true;
+		}
+
+	protected:
+		const ULONGLONG startAddr;
+		const ULONGLONG endAddr;
+	};
+	//---
+	ThunkFilterSelfCallback filter = ThunkFilterSelfCallback(peBuffer.moduleBase, peBuffer.getBufferSize());
+
+	IATBlock* iat_block = nullptr;
+	if (this->is64bit) {
+		iat_block = find_iat<ULONGLONG>(this->peBuffer.vBuf, this->peBuffer.vBufSize, exportsMap, start_offset, &filter);
+	}
+	else {
+		iat_block = find_iat<DWORD>(this->peBuffer.vBuf, this->peBuffer.vBufSize, exportsMap, start_offset, &filter);
+	}
+	return iat_block;
+}
+
 IATBlock* pesieve::ImpReconstructor::findIAT(IN const peconv::ExportsMapper* exportsMap, size_t start_offset)
 {
 	BYTE *vBuf = this->peBuffer.vBuf;
 	const size_t vBufSize = this->peBuffer.vBufSize;
 	if (!vBuf) return nullptr;
 
-	IATBlock* iat_block = find_iat_block(is64bit, vBuf, vBufSize, exportsMap, start_offset);;
+	IATBlock* iat_block = findIATBlock(exportsMap, start_offset);
 	if (!iat_block) {
 		return nullptr;
 	}

postprocessors/imp_rec/imp_reconstructor.h

@@ -189,7 +189,7 @@ namespace pesieve {
 		bool printFoundIATs(std::string reportPath);
 
 	private:
-
+		IATBlock* findIATBlock(IN const peconv::ExportsMapper* exportsMap, size_t start_offset);
 		IATBlock* findIAT(IN const peconv::ExportsMapper* exportsMap, size_t start_offset);
 
 		//!  has more IATs that the main one (dynamically loaded)

postprocessors/pe_buffer.cpp

@@ -3,7 +3,7 @@
 #include <iostream>
 #include "../scanners/artefact_scanner.h"
 
-size_t pesieve::PeBuffer::calcRemoteImgSize(ULONGLONG modBaseAddr)
+size_t pesieve::PeBuffer::calcRemoteImgSize(ULONGLONG modBaseAddr) const
 {
 	const size_t hdr_buffer_size = PAGE_SIZE;
 	BYTE hdr_buffer[hdr_buffer_size] = { 0 };

postprocessors/pe_buffer.h

@@ -32,7 +32,7 @@ namespace pesieve {
 		}
 
 		// Returns the size of the internal buffer
-		size_t getBufferSize()
+		size_t getBufferSize() const
 		{
 			return vBufSize;
 		}
@@ -58,12 +58,12 @@ namespace pesieve {
 
 		bool dumpToFile(IN std::string dumpFileName);
 
-		ULONGLONG getModuleBase()
+		ULONGLONG getModuleBase() const
 		{
 			return moduleBase;
 		}
 
-		ULONGLONG getRelocBase()
+		ULONGLONG getRelocBase() const
 		{
 			return relocBase;
 		}
@@ -76,7 +76,7 @@ namespace pesieve {
 	protected:
 		bool _readRemote(ULONGLONG module_base, size_t pe_vsize);
 
-		size_t calcRemoteImgSize(ULONGLONG module_base);
+		size_t calcRemoteImgSize(ULONGLONG module_base) const;
 
 		bool allocBuffer(const size_t pe_vsize)
 		{