Fix #20

Author: romainthomas

api/python/PE/objects/pyBinary.cpp

@@ -171,6 +171,11 @@ void init_PE_Binary_class(py::module& m) {
     .def_property_readonly("resources_manager",
         static_cast<no_const_getter<ResourcesManager>>(&Binary::get_resources_manager))
 
+    .def_property_readonly("overlay",
+        static_cast<no_const_getter<std::vector<uint8_t>&>>(&Binary::overlay),
+        "Return the overlay content",
+        py::return_value_policy::reference)
+
     .def("add_import_function",
         &Binary::add_import_function,
         py::return_value_policy::reference)

api/python/PE/objects/pyBuilder.cpp

@@ -50,6 +50,10 @@ void init_PE_Builder_class(py::module& m) {
         &Builder::build_imports,
         py::return_value_policy::reference)
 
+    .def("build_overlay",
+        static_cast<Builder& (Builder::*)(bool)>(&Builder::build_overlay),
+        py::return_value_policy::reference)
+
     .def("write",
         &Builder::write)
 

include/LIEF/PE/Binary.hpp

@@ -215,6 +215,15 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
     Debug&       get_debug(void);
     const Debug& get_debug(void) const;
 
+    // =======
+    // Overlay
+    // =======
+
+    //! @brief Return the overlay content
+    const std::vector<uint8_t>& overlay(void) const;
+    std::vector<uint8_t>&       overlay(void);
+
+
     // =========================
     // Methods to manage Imports
     // =========================
@@ -293,32 +302,33 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
     void update_lookup_address_table_offset(void);
     void update_iat(void);
 
-    PE_TYPE type_;
-    DosHeader          dos_header_;
-    Header             header_;
-    OptionalHeader     optional_header_;
-
-    bool               has_tls_;
-    bool               has_imports_;
-    bool               has_signature_;
-    bool               has_exports_;
-    bool               has_resources_;
-    bool               has_exceptions_;
-    bool               has_relocations_;
-    bool               has_debug_;
-    bool               has_configuration_;
-
-    Signature          signature_;
-    TLS                tls_;
-    sections_t         sections_;
-    data_directories_t data_directories_;
-    symbols_t          symbols_;
-    strings_table_t    strings_table_;
-    relocations_t      relocations_;
-    ResourceNode*      resources_;
-    imports_t          imports_;
-    Export             export_;
-    Debug              debug_;
+    PE_TYPE              type_;
+    DosHeader            dos_header_;
+    Header               header_;
+    OptionalHeader       optional_header_;
+
+    bool                 has_tls_;
+    bool                 has_imports_;
+    bool                 has_signature_;
+    bool                 has_exports_;
+    bool                 has_resources_;
+    bool                 has_exceptions_;
+    bool                 has_relocations_;
+    bool                 has_debug_;
+    bool                 has_configuration_;
+
+    Signature            signature_;
+    TLS                  tls_;
+    sections_t           sections_;
+    data_directories_t   data_directories_;
+    symbols_t            symbols_;
+    strings_table_t      strings_table_;
+    relocations_t        relocations_;
+    ResourceNode*        resources_;
+    imports_t            imports_;
+    Export               export_;
+    Debug                debug_;
+    std::vector<uint8_t> overlay_;
 
 };
 

include/LIEF/PE/Builder.hpp

@@ -54,6 +54,7 @@ class DLL_PUBLIC Builder
     Builder& build_relocations(bool flag = true);
     Builder& build_tls(bool flag = true);
     Builder& build_resources(bool flag);
+    Builder& build_overlay(bool flag);
 
     const std::vector<uint8_t>& get_build(void);
     void write(const std::string& filename) const;
@@ -83,6 +84,7 @@ class DLL_PUBLIC Builder
     void build_string_table(void);
     void build_relocation(void);
     void build_resources(void);
+    void build_overlay(void);
 
     void compute_resources_size(
         ResourceNode *node,
@@ -108,6 +110,7 @@ class DLL_PUBLIC Builder
     bool build_relocations_;
     bool build_tls_;
     bool build_resources_;
+    bool build_overlay_;
 
 };
 

include/LIEF/PE/Parser.hpp

@@ -83,6 +83,7 @@ class DLL_PUBLIC Parser : public LIEF::Parser {
     void build_string_table(void);
     void build_symbols(void);
     void build_signature(void);
+    void build_overlay(void);
 
     ResourceNode* build_resource_node(
         const pe_resource_directory_table *directoryTable, uint32_t baseOffset);

src/PE/Binary.cpp

@@ -925,6 +925,17 @@ std::vector<uint8_t> Binary::get_content_from_virtual_address(uint64_t virtual_a
 
 }
 
+// Overlay
+// =======
+
+const std::vector<uint8_t>& Binary::overlay(void) const {
+  return this->overlay_;
+}
+
+std::vector<uint8_t>& Binary::overlay(void) {
+  return const_cast<std::vector<uint8_t>&>(static_cast<const Binary*>(this)->overlay());
+}
+
 // Resource manager
 // ===============
 

src/PE/Builder.cpp

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <fstream>
 #include <iterator>
+#include <numeric>
 
 #include "easylogging++.h"
 
@@ -41,7 +42,8 @@ Builder::Builder(Binary* binary) :
   patch_imports_{false},
   build_relocations_{false},
   build_tls_{false},
-  build_resources_{false}
+  build_resources_{false},
+  build_overlay_{true}
 {}
 
 
@@ -69,7 +71,10 @@ Builder& Builder::build_resources(bool flag) {
   return *this;
 }
 
-
+Builder& Builder::build_overlay(bool flag) {
+  this->build_overlay_ = flag;
+  return *this;
+}
 
 
 void Builder::write(const std::string& filename) const {
@@ -85,8 +90,6 @@ void Builder::write(const std::string& filename) const {
 }
 
 
-
-
 void Builder::build(void) {
 
   LOG(DEBUG) << "Rebuilding" << std::endl;
@@ -151,6 +154,12 @@ void Builder::build(void) {
 
   LOG(DEBUG) << "[+] Rebuilding string table" << std::endl;
   //this->build_string_table();
+
+  if (this->binary_->overlay().size() > 0 and this->build_overlay_) {
+    this->build_overlay();
+  }
+
+
 }
 
 const std::vector<uint8_t>& Builder::get_build(void) {
@@ -402,6 +411,25 @@ void Builder::build_string_table(void) {
   //TODO
 }
 
+void Builder::build_overlay(void) {
+  LOG(DEBUG) << "Building overlay";
+
+  const uint64_t last_section_offset = std::accumulate(
+      std::begin(this->binary_->sections_),
+      std::end(this->binary_->sections_), 0,
+      [] (uint64_t offset, const Section* section) {
+        return std::max<uint64_t>(section->offset() + section->size(), offset);
+      });
+
+  LOG(DEBUG) << "Overlay offset: 0x" << std::hex << last_section_offset;
+  LOG(DEBUG) << "Overlay size: " << std::dec << this->binary_->overlay().size();
+
+  const size_t saved_offset = this->ios_.tellp();
+  this->ios_.seekp(last_section_offset);
+  this->ios_.write(this->binary_->overlay());
+  this->ios_.seekp(saved_offset);
+}
+
 Builder& Builder::operator<<(const DosHeader& dos_header) {
 
   pe_dos_header dosHeader;
@@ -519,6 +547,7 @@ std::ostream& operator<<(std::ostream& os, const Builder& b) {
   os << std::setw(20) << "Builde relocations:" << b.build_relocations_ << std::endl;
   os << std::setw(20) << "Builde TLS:"         << b.build_tls_         << std::endl;
   os << std::setw(20) << "Builder resources:"  << b.build_resources_   << std::endl;
+  os << std::setw(20) << "Builder overlay:"    << b.build_overlay_     << std::endl;
   return os;
 }
 

src/PE/Parser.cpp

@@ -17,6 +17,7 @@
 #include <iterator>
 #include <iostream>
 #include <string>
+#include <numeric>
 
 #include <mbedtls/platform.h>
 #include <mbedtls/oid.h>
@@ -1018,6 +1019,33 @@ void Parser::build_signature(void) {
   this->binary_->has_signature_ = true;
 }
 
+
+void Parser::build_overlay(void) {
+  LOG(DEBUG) << "Parsing Overlay";
+  const uint64_t last_section_offset = std::accumulate(
+      std::begin(this->binary_->sections_),
+      std::end(this->binary_->sections_), 0,
+      [] (uint64_t offset, const Section* section) {
+        return std::max<uint64_t>(section->offset() + section->size(), offset);
+      });
+
+  LOG(DEBUG) << "Overlay offset: 0x" << std::hex << last_section_offset;
+
+  const uint64_t overlay_size = this->stream_->size() - last_section_offset;
+
+  LOG(DEBUG) << "Overlay size: " << std::dec << overlay_size;
+
+  const uint8_t* ptr_to_overlay = reinterpret_cast<const uint8_t*>(this->stream_->read(
+        last_section_offset,
+        overlay_size));
+
+  this->binary_->overlay_ = {
+        ptr_to_overlay,
+        ptr_to_overlay + overlay_size
+      };
+
+}
+
 //
 // Return the Binary constructed
 //

src/PE/Parser.tcc

@@ -47,6 +47,8 @@ void Parser::build(void) {
   } catch (const corrupted& e) {
     LOG(WARNING) << e.what();
   }
+
+  this->build_overlay();
 }
 
 template<typename PE_T>