From 15b446b632059bdba7fc8228c2a0e54b8f7f2b76 Mon Sep 17 00:00:00 2001
From: hasherezade <hasherezade@gmail.com>
Date: Fri, 3 Sep 2021 17:46:12 +0200
Subject: [PATCH] [FEATURE] If new IATs detected, recreate full table (Issue
 #89)

---
 postprocessors/imp_rec/imp_reconstructor.cpp | 16 +++++++++++++++-
 postprocessors/imp_rec/imp_reconstructor.h   |  6 +++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/postprocessors/imp_rec/imp_reconstructor.cpp b/postprocessors/imp_rec/imp_reconstructor.cpp
index 522325f95..f1e61beaf 100644
--- 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
index bab83972e..53191db6c 100644
--- 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);