From e161a5a23368b3a5138af0685d846a66259fcb67 Mon Sep 17 00:00:00 2001
From: rukaimi <rukaimi@985867f9-ca9c-e1f6-822d-e8a4186388af>
Date: Wed, 12 Sep 2012 20:09:13 +0000
Subject: [PATCH] 1. Added new option for import rebuilder 2. Minor bugfixes
---
pe_lib/pe_32_64.cpp | 176 +++++++++++++++++++++++++++++++++++-------
pe_lib/pe_base.cpp | 27 +++++--
pe_lib/pe_base.h | 15 +++-
pe_lib/pe_exception.h | 2 +-
4 files changed, 183 insertions(+), 37 deletions(-)
diff --git a/pe_lib/pe_32_64.cpp b/pe_lib/pe_32_64.cpp
index 64fa7de..cbcf194 100644
--- a/pe_lib/pe_32_64.cpp
+++ b/pe_lib/pe_32_64.cpp
@@ -671,7 +671,7 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
//Check if import_section is last one. If it's not, check if there's enough place for import data
if(&import_section != &*(sections_.end() - 1) &&
(import_section.empty() || align_up(import_section.get_size_of_raw_data(), get_file_alignment()) < needed_size + import_settings.get_offset_from_section_start()))
- throw pe_exception("Insufficient space for import directory", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for import directory", pe_exception::insufficient_space);
std::string& raw_data = import_section.get_raw_data();
@@ -695,9 +695,20 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
IMAGE_IMPORT_DESCRIPTOR descr = {0};
descr.TimeDateStamp = (*it).get_timestamp(); //Restore timestamp
descr.Name = rva_from_section_offset(import_section, current_string_pointer); //Library name RVA
-
- //If we should save IATs for current import descriptor
+
+ //If we should save IAT for current import descriptor
bool save_iats_for_this_descriptor = import_settings.save_iat_and_original_iat_rvas() && (*it).get_rva_to_iat() != 0;
+ //If we should write original IAT
+ bool write_original_iat = (!save_iats_for_this_descriptor && import_settings.build_original_iat()) || import_settings.fill_missing_original_iats();
+
+ //If we should rewrite saved original IAT for current import descriptor (without changing its position)
+ bool rewrite_saved_original_iat = save_iats_for_this_descriptor && import_settings.rewrite_iat_and_original_iat_contents() && import_settings.build_original_iat();
+ //If we should rewrite saved IAT for current import descriptor (without changing its position)
+ bool rewrite_saved_iat = save_iats_for_this_descriptor && import_settings.rewrite_iat_and_original_iat_contents() && (*it).get_rva_to_iat() != 0;
+
+ //Helper values if we're rewriting existing IAT or orig.IAT
+ DWORD original_first_thunk = 0;
+ DWORD first_thunk = 0;
if(save_iats_for_this_descriptor)
{
@@ -708,6 +719,20 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
descr.OriginalFirstThunk = import_settings.build_original_iat() ? (*it).get_rva_to_original_iat() : 0;
descr.FirstThunk = (*it).get_rva_to_iat();
+
+ original_first_thunk = descr.OriginalFirstThunk;
+ first_thunk = descr.FirstThunk;
+
+ if(rewrite_saved_original_iat)
+ {
+ if((*it).get_rva_to_original_iat())
+ write_original_iat = true;
+ else
+ rewrite_saved_original_iat = false;
+ }
+
+ if(rewrite_saved_iat)
+ save_iats_for_this_descriptor = false;
}
else
{
@@ -728,9 +753,6 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
const import_library::imported_list& funcs = (*it).get_imported_functions();
for(import_library::imported_list::const_iterator f = funcs.begin(); f != funcs.end(); ++f)
{
- //If we must write original IAT
- bool write_original_iat = (!save_iats_for_this_descriptor && import_settings.build_original_iat()) || import_settings.fill_missing_original_iats();
-
if((*f).has_name()) //If function is imported by name
{
//Get RVA of IMAGE_IMPORT_BY_NAME
@@ -743,22 +765,59 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
//We're creating original IATs - so we can write to IAT saved VA (because IMAGE_IMPORT_BY_NAME will be read
//by PE loader from original IAT)
typename PEClassType::BaseSize iat_value = static_cast<typename PEClassType::BaseSize>((*f).get_iat_va());
- memcpy(&raw_data[current_pos_for_iat], &iat_value, sizeof(iat_value));
+
+ if(rewrite_saved_iat)
+ {
+ if(section_data_length_from_rva(first_thunk, first_thunk, section_data_raw, true) <= sizeof(iat_value))
+ throw pe_exception("Insufficient space inside initial IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(first_thunk, true), &iat_value, sizeof(iat_value));
+
+ first_thunk += sizeof(iat_value);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_iat], &iat_value, sizeof(iat_value));
+ current_pos_for_iat += sizeof(rva_of_named_import);
+ }
}
else
{
//Else - write to IAT RVA of IMAGE_IMPORT_BY_NAME
- memcpy(&raw_data[current_pos_for_iat], &rva_of_named_import, sizeof(rva_of_named_import));
+ if(rewrite_saved_iat)
+ {
+ if(section_data_length_from_rva(first_thunk, first_thunk, section_data_raw, true) <= sizeof(rva_of_named_import))
+ throw pe_exception("Insufficient space inside initial IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(first_thunk, true), &rva_of_named_import, sizeof(rva_of_named_import));
+
+ first_thunk += sizeof(rva_of_named_import);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_iat], &rva_of_named_import, sizeof(rva_of_named_import));
+ current_pos_for_iat += sizeof(rva_of_named_import);
+ }
}
-
- current_pos_for_iat += sizeof(rva_of_named_import);
}
if(write_original_iat)
{
- //We're creating original IATs
- memcpy(&raw_data[current_pos_for_original_iat], &rva_of_named_import, sizeof(rva_of_named_import));
- current_pos_for_original_iat += sizeof(rva_of_named_import);
+ if(rewrite_saved_original_iat)
+ {
+ if(section_data_length_from_rva(original_first_thunk, original_first_thunk, section_data_raw, true) <= sizeof(rva_of_named_import))
+ throw pe_exception("Insufficient space inside initial original IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(original_first_thunk, true), &rva_of_named_import, sizeof(rva_of_named_import));
+
+ original_first_thunk += sizeof(rva_of_named_import);
+ }
+ else
+ {
+ //We're creating original IATs
+ memcpy(&raw_data[current_pos_for_original_iat], &rva_of_named_import, sizeof(rva_of_named_import));
+ current_pos_for_original_iat += sizeof(rva_of_named_import);
+ }
}
//Write IMAGE_IMPORT_BY_NAME (WORD hint + string function name)
@@ -780,34 +839,97 @@ const pe_base::image_directory pe<PEClassType>::rebuild_imports(const imported_f
//We're creating original IATs - so we can wtire to IAT saved VA (because ordinal will be read
//by PE loader from original IAT)
typename PEClassType::BaseSize iat_value = static_cast<typename PEClassType::BaseSize>((*f).get_iat_va());
- memcpy(&raw_data[current_pos_for_iat], &iat_value, sizeof(iat_value));
+ if(rewrite_saved_iat)
+ {
+ if(section_data_length_from_rva(first_thunk, first_thunk, section_data_raw, true) <= sizeof(iat_value))
+ throw pe_exception("Insufficient space inside initial IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(first_thunk, true), &iat_value, sizeof(iat_value));
+
+ first_thunk += sizeof(iat_value);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_iat], &iat_value, sizeof(iat_value));
+ current_pos_for_iat += sizeof(thunk_value);
+ }
}
else
{
//Else - write ordinal to IAT
- memcpy(&raw_data[current_pos_for_iat], &thunk_value, sizeof(thunk_value));
+ if(rewrite_saved_iat)
+ {
+ if(section_data_length_from_rva(first_thunk, first_thunk, section_data_raw, true) <= sizeof(thunk_value))
+ throw pe_exception("Insufficient space inside initial IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(first_thunk, true), &thunk_value, sizeof(thunk_value));
+
+ first_thunk += sizeof(thunk_value);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_iat], &thunk_value, sizeof(thunk_value));
+ }
}
-
- current_pos_for_iat += sizeof(thunk_value);
}
//We're writing ordinal to original IAT slot
if(write_original_iat)
{
- memcpy(&raw_data[current_pos_for_original_iat], &thunk_value, sizeof(thunk_value));
- current_pos_for_original_iat += sizeof(thunk_value);
+ if(rewrite_saved_original_iat)
+ {
+ if(section_data_length_from_rva(original_first_thunk, original_first_thunk, section_data_raw, true) <= sizeof(thunk_value))
+ throw pe_exception("Insufficient space inside initial original IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(original_first_thunk, true), &thunk_value, sizeof(thunk_value));
+
+ original_first_thunk += sizeof(thunk_value);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_original_iat], &thunk_value, sizeof(thunk_value));
+ current_pos_for_original_iat += sizeof(thunk_value);
+ }
}
}
}
- if((!save_iats_for_this_descriptor && import_settings.build_original_iat()) || import_settings.fill_missing_original_iats())
+ if(!save_iats_for_this_descriptor)
+ {
+ //Ending null thunks
+ typename PEClassType::BaseSize thunk_value = 0;
+
+ if(rewrite_saved_iat)
+ {
+ if(section_data_length_from_rva(first_thunk, first_thunk, section_data_raw, true) <= sizeof(thunk_value))
+ throw pe_exception("Insufficient space inside initial IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(first_thunk, true), &thunk_value, sizeof(thunk_value));
+
+ first_thunk += sizeof(thunk_value);
+ }
+ else
+ {
+ memcpy(&raw_data[current_pos_for_iat], &thunk_value, sizeof(thunk_value));
+ current_pos_for_iat += sizeof(thunk_value);
+ }
+ }
+
+ if(write_original_iat)
{
//Ending null thunks
typename PEClassType::BaseSize thunk_value = 0;
- memcpy(&raw_data[current_pos_for_iat], &thunk_value, sizeof(thunk_value));
- current_pos_for_iat += sizeof(thunk_value);
- if(import_settings.build_original_iat())
+ if(rewrite_saved_original_iat)
+ {
+ if(section_data_length_from_rva(original_first_thunk, original_first_thunk, section_data_raw, true) <= sizeof(thunk_value))
+ throw pe_exception("Insufficient space inside initial original IAT", pe_exception::insufficient_space);
+
+ memcpy(section_data_from_rva(original_first_thunk, true), &thunk_value, sizeof(thunk_value));
+
+ original_first_thunk += sizeof(thunk_value);
+ }
+ else
{
memcpy(&raw_data[current_pos_for_original_iat], &thunk_value, sizeof(thunk_value));
current_pos_for_original_iat += sizeof(thunk_value);
@@ -936,7 +1058,7 @@ const pe_base::image_directory pe<PEClassType>::rebuild_tls(const tls_info& info
//Check if tls_section is last one. If it's not, check if there's enough place for TLS data
if(&tls_section != &*(sections_.end() - 1) &&
(tls_section.empty() || align_up(tls_section.get_size_of_raw_data(), get_file_alignment()) < needed_size + offset_from_section_start))
- throw pe_exception("Insufficient space for TLS directory", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for TLS directory", pe_exception::insufficient_space);
//Check raw data positions
if(info.get_raw_data_end_rva() < info.get_raw_data_start_rva() || info.get_index_rva() == 0)
@@ -996,7 +1118,7 @@ const pe_base::image_directory pe<PEClassType>::rebuild_tls(const tls_info& info
//Check if there's enough virtual space for it...
if(section_data_length_from_rva(info.get_raw_data_start_rva(), info.get_raw_data_start_rva(), section_data_virtual, true)
< info.get_raw_data_end_rva() - info.get_raw_data_start_rva())
- throw pe_exception("Insufficient space for TLS raw data", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for TLS raw data", pe_exception::insufficient_space);
else
write_raw_data_size = available_raw_length; //We'll write just a part of full raw data
}
@@ -1025,11 +1147,11 @@ const pe_base::image_directory pe<PEClassType>::rebuild_tls(const tls_info& info
//Check if there's enough space to write callbacks TLS data...
if(section_data_length_from_rva(info.get_callbacks_rva(), info.get_callbacks_rva(), section_data_raw, true)
< needed_callback_size - sizeof(typename PEClassType::BaseSize) /* last zero element can be virtual only */)
- throw pe_exception("Insufficient space for TLS callbacks data", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for TLS callbacks data", pe_exception::insufficient_space);
if(section_data_length_from_rva(info.get_callbacks_rva(), info.get_callbacks_rva(), section_data_virtual, true)
< needed_callback_size /* check here full virtual data length available */)
- throw pe_exception("Insufficient space for TLS callbacks data", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for TLS callbacks data", pe_exception::insufficient_space);
std::vector<typename PEClassType::BaseSize> callbacks_virtual_addresses;
callbacks_virtual_addresses.reserve(info.get_tls_callbacks().size() + 1 /* last null element */);
diff --git a/pe_lib/pe_base.cpp b/pe_lib/pe_base.cpp
index cef3238..1057fdf 100644
--- a/pe_lib/pe_base.cpp
+++ b/pe_lib/pe_base.cpp
@@ -1044,7 +1044,7 @@ void pe_base::read_pe(std::istream& file, bool read_bound_import_raw_data, bool
//If section raw data size is greater than virtual, fix it
last_raw_size = s.header_.SizeOfRawData;
- if(s.header_.SizeOfRawData > s.header_.Misc.VirtualSize)
+ if(align_up(s.header_.SizeOfRawData, get_file_alignment()) > align_up(s.header_.Misc.VirtualSize, get_section_alignment()))
s.header_.SizeOfRawData = s.header_.Misc.VirtualSize;
//Check virtual and raw section sizes and addresses
@@ -2115,7 +2115,7 @@ const pe_base::image_directory pe_base::rebuild_exports(const export_info& info,
//Check if exports_section is last one. If it's not, check if there's enough place for exports data
if(&exports_section != &*(sections_.end() - 1) &&
(exports_section.empty() || align_up(exports_section.get_size_of_raw_data(), get_file_alignment()) < needed_size + offset_from_section_start))
- throw pe_exception("Insufficient space for export directory", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for export directory", pe_exception::insufficient_space);
std::string& raw_data = exports_section.get_raw_data();
@@ -2254,7 +2254,8 @@ pe_base::import_rebuilder_settings::import_rebuilder_settings(bool set_to_pe_hea
save_iat_and_original_iat_rvas_(true),
fill_missing_original_iats_(false),
set_to_pe_headers_(set_to_pe_headers),
- zero_directory_entry_iat_(auto_zero_directory_entry_iat)
+ zero_directory_entry_iat_(auto_zero_directory_entry_iat),
+ rewrite_iat_and_original_iat_contents_(false)
{}
//Returns offset from section start where import directory data will be placed
@@ -2276,6 +2277,14 @@ bool pe_base::import_rebuilder_settings::save_iat_and_original_iat_rvas() const
return save_iat_and_original_iat_rvas_;
}
+//Returns true if Original import address and import address tables contents will be rewritten
+//works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
+//and save_iat_and_original_iat_rvas is true
+bool pe_base::import_rebuilder_settings::rewrite_iat_and_original_iat_contents() const
+{
+ return rewrite_iat_and_original_iat_contents_;
+}
+
//Returns true if original missing IATs will be rebuilt
//(only if IATs are saved)
bool pe_base::import_rebuilder_settings::fill_missing_original_iats() const
@@ -2309,9 +2318,13 @@ void pe_base::import_rebuilder_settings::build_original_iat(bool enable)
//Sets if Original import address and import address tables will not be rebuilt,
//works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
-void pe_base::import_rebuilder_settings::save_iat_and_original_iat_rvas(bool enable)
+void pe_base::import_rebuilder_settings::save_iat_and_original_iat_rvas(bool enable, bool enable_rewrite_iat_and_original_iat_contents)
{
save_iat_and_original_iat_rvas_ = enable;
+ if(save_iat_and_original_iat_rvas_)
+ rewrite_iat_and_original_iat_contents_ = enable_rewrite_iat_and_original_iat_contents;
+ else
+ rewrite_iat_and_original_iat_contents_ = false;
}
//Sets if original missing IATs will be rebuilt
@@ -2642,7 +2655,7 @@ const pe_base::image_directory pe_base::rebuild_relocations(const relocation_tab
//Check if reloc_section is last one. If it's not, check if there's enough place for relocations data
if(&reloc_section != &*(sections_.end() - 1) &&
(reloc_section.empty() || align_up(reloc_section.get_size_of_raw_data(), get_file_alignment()) < needed_size + offset_from_section_start))
- throw pe_exception("Insufficient space for relocations directory", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for relocations directory", pe_exception::insufficient_space);
std::string& raw_data = reloc_section.get_raw_data();
@@ -3695,7 +3708,7 @@ const pe_base::image_directory pe_base::rebuild_resources(resource_directory& in
//Check if exports_section is last one. If it's not, check if there's enough place for resource data
if(&resources_section != &*(sections_.end() - 1) &&
(resources_section.empty() || align_up(resources_section.get_size_of_raw_data(), get_file_alignment()) < needed_size + offset_from_section_start))
- throw pe_exception("Insufficient space for resource directory", pe_exception::insuffisient_space);
+ throw pe_exception("Insufficient space for resource directory", pe_exception::insufficient_space);
std::string& raw_data = resources_section.get_raw_data();
@@ -5106,7 +5119,7 @@ void pe_base::recalculate_section_sizes(section& s)
}
//Can occur only for last section
- if(align_up(s.get_virtual_size(), get_section_alignment()) < s.get_size_of_raw_data())
+ if(align_up(s.get_virtual_size(), get_section_alignment()) < align_up(s.get_size_of_raw_data(), get_file_alignment()))
set_section_virtual_size(s, align_up(s.get_size_of_raw_data(), get_section_alignment())); //Recalculate section virtual size
}
diff --git a/pe_lib/pe_base.h b/pe_lib/pe_base.h
index 35e5336..f218be0 100644
--- a/pe_lib/pe_base.h
+++ b/pe_lib/pe_base.h
@@ -12,7 +12,7 @@
//Please don't remove this information from header
//PE Library (c) DX 2011 - 2012, http://kaimi.ru
-//Version: 0.1.2
+//Version: 0.1.3
//Free to use, modify and distribute
// == more important ==
@@ -838,9 +838,15 @@ public: //IMPORTS
DWORD get_offset_from_section_start() const;
//Returns true if Original import address table (IAT) will be rebuilt
bool build_original_iat() const;
+
//Returns true if Original import address and import address tables will not be rebuilt,
//works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
bool save_iat_and_original_iat_rvas() const;
+ //Returns true if Original import address and import address tables contents will be rewritten
+ //works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
+ //and save_iat_and_original_iat_rvas is true
+ bool rewrite_iat_and_original_iat_contents() const;
+
//Returns true if original missing IATs will be rebuilt
//(only if IATs are saved)
bool fill_missing_original_iats() const;
@@ -849,6 +855,7 @@ public: //IMPORTS
//Returns true if IMAGE_DIRECTORY_ENTRY_IAT must be zeroed, works only if auto_set_to_pe_headers = true
bool zero_directory_entry_iat() const;
+
public: //Setters
//Sets offset from section start where import directory data will be placed
void set_offset_from_section_start(DWORD offset);
@@ -856,7 +863,10 @@ public: //IMPORTS
void build_original_iat(bool enable);
//Sets if Original import address and import address tables will not be rebuilt,
//works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
- void save_iat_and_original_iat_rvas(bool enable);
+ //enable_rewrite_iat_and_original_iat_contents sets if Original import address and import address tables contents will be rewritten
+ //works only if import descriptor IAT (and orig.IAT, if present) RVAs are not zero
+ //and save_iat_and_original_iat_rvas is true
+ void save_iat_and_original_iat_rvas(bool enable, bool enable_rewrite_iat_and_original_iat_contents = false);
//Sets if original missing IATs will be rebuilt
//(only if IATs are saved)
void fill_missing_original_iats(bool enable);
@@ -872,6 +882,7 @@ public: //IMPORTS
bool fill_missing_original_iats_;
bool set_to_pe_headers_;
bool zero_directory_entry_iat_;
+ bool rewrite_iat_and_original_iat_contents_;
};
//You can get all image imports with get_imported_functions() function
diff --git a/pe_lib/pe_exception.h b/pe_lib/pe_exception.h
index ddc9da6..7b05071 100644
--- a/pe_lib/pe_exception.h
+++ b/pe_lib/pe_exception.h
@@ -67,7 +67,7 @@ public:
stream_is_bad,
section_is_not_attached,
- insuffisient_space,
+ insufficient_space,
cannot_rebase_relocations,
--
GitLab