[FEATURE] If new IATs detected, recreate full table (Issue #89)

hasherezade · Sep 3, 2021 · 15b446b · 15b446b
1 parent fb9c3d8
commit 15b446b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/postprocessors/imp_rec/imp_reconstructor.cpp b/postprocessors/imp_rec/imp_reconstructor.cpp
@@ -35,6 +35,20 @@ BYTE* pesieve::ImportTableBuffer::getDllSpaceAt(const DWORD rva, size_t required
 
 //---
 
+bool pesieve::ImpReconstructor::hasNewImportTables()
+{
+	bool has_new_table = false;
+	std::map<DWORD, IATBlock*>::const_iterator iats_itr;
+	for (iats_itr = foundIATs.cbegin(); iats_itr != foundIATs.cend(); ++iats_itr) {
+		const IATBlock* iblock = iats_itr->second;
+		if (iblock->isTerminated && !iblock->isMain) {
+			has_new_table = true;
+			break;
+		}
+	}
+	return has_new_table;
+}
+
 pesieve::ImpReconstructor::t_imprec_res pesieve::ImpReconstructor::rebuildImportTable(const IN peconv::ExportsMapper* exportsMap, IN const pesieve::t_imprec_mode &imprec_mode)
 {
 	if (!exportsMap || imprec_mode == pesieve::PE_IMPREC_NONE) {
@@ -55,7 +69,7 @@ pesieve::ImpReconstructor::t_imprec_res pesieve::ImpReconstructor::rebuildImport
 		return IMP_RECOVERY_NOT_APPLICABLE;
 	}
 
-	if (imprec_mode == PE_IMPREC_UNERASE || imprec_mode == PE_IMPREC_AUTO) {
+	if (imprec_mode == PE_IMPREC_UNERASE || (imprec_mode == PE_IMPREC_AUTO && !hasNewImportTables())) {
 
 		if (this->isDefaultImportValid(exportsMap)) {
 			// Valid Import Table already set

diff --git a/postprocessors/imp_rec/imp_reconstructor.h b/postprocessors/imp_rec/imp_reconstructor.h
@@ -190,7 +190,11 @@ namespace pesieve {
 
 	private:
 
-		IATBlock * findIAT(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)
+		bool hasNewImportTables();
+
 		bool findImportTable(IN const peconv::ExportsMapper* exportsMap);
 		size_t collectIATs(IN const peconv::ExportsMapper* exportsMap);