Enable DOS stub rebuilding

  Python API:

  * lief.PE.Binary.dos_stub property
  * lief.PE.Builder.build_dos_stub method

  C++ API:

  * LIEF::PE::Binary::dos_stub setter/getter
  * LIEF::PE::Builder::build_dos_stub method

Author: romainthomas

api/python/PE/objects/pyBinary.cpp

@@ -27,6 +27,12 @@ using no_const_func = T (Binary::*)(P);
 template<class T>
 using no_const_getter = T (Binary::*)(void);
 
+template<class T>
+using getter_t = T (Binary::*)(void) const;
+
+template<class T>
+using setter_t = void (Binary::*)(T);
+
 void init_PE_Binary_class(py::module& m) {
   py::class_<Binary, LIEF::Binary>(m, "Binary")
     .def(py::init<const std::string &, PE_TYPE>())
@@ -204,6 +210,11 @@ void init_PE_Binary_class(py::module& m) {
         "Return the overlay content",
         py::return_value_policy::reference)
 
+    .def_property("dos_stub",
+        static_cast<getter_t<const std::vector<uint8_t>&>>(&Binary::dos_stub),
+        static_cast<setter_t<const std::vector<uint8_t>&>>(&Binary::dos_stub),
+        "DOS stub content")
+
     .def("add_import_function",
         &Binary::add_import_function,
         "Add a function to the given " RST_CLASS_REF(lief.PE.Import) " name",

api/python/PE/objects/pyBuilder.cpp

@@ -63,6 +63,11 @@ void init_PE_Builder_class(py::module& m) {
         "Rebuild the binary's overlay",
         py::return_value_policy::reference)
 
+    .def("build_dos_stub",
+        static_cast<Builder& (Builder::*)(bool)>(&Builder::build_dos_stub),
+        "Rebuild the DOS stub",
+        py::return_value_policy::reference)
+
     .def("write",
         &Builder::write,
         "Write the build result into the ``output`` file",

include/LIEF/PE/Binary.hpp

@@ -240,6 +240,16 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
     const std::vector<uint8_t>& overlay(void) const;
     std::vector<uint8_t>&       overlay(void);
 
+    // ========
+    // DOS Stub
+    // ========
+
+    //! @brief Return the DOS stub content
+    const std::vector<uint8_t>& dos_stub(void) const;
+    std::vector<uint8_t>&       dos_stub(void);
+
+    //! @brief Update the DOS stub content
+    void dos_stub(const std::vector<uint8_t>& content);
 
     // =========================
     // Methods to manage Imports
@@ -377,6 +387,7 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
     Export               export_;
     Debug                debug_;
     std::vector<uint8_t> overlay_;
+    std::vector<uint8_t> dos_stub_;
 
     std::map<std::string, std::map<std::string, uint64_t>> hooks_;
 };

include/LIEF/PE/Builder.hpp

@@ -81,6 +81,9 @@ class DLL_PUBLIC Builder
     //! @brief Rebuild the binary's overlay
     Builder& build_overlay(bool flag);
 
+    //! @brief Rebuild the DOS stub content
+    Builder& build_dos_stub(bool flag);
+
     //! @brief Return the build result
     const std::vector<uint8_t>& get_build(void);
 
@@ -113,6 +116,7 @@ class DLL_PUBLIC Builder
     void build_relocation(void);
     void build_resources(void);
     void build_overlay(void);
+    void build_dos_stub(void);
 
     void compute_resources_size(
         ResourceNode *node,
@@ -139,6 +143,7 @@ class DLL_PUBLIC Builder
     bool build_tls_;
     bool build_resources_;
     bool build_overlay_;
+    bool build_dos_stub_;
 
 };
 

include/LIEF/PE/Parser.hpp

@@ -84,14 +84,15 @@ class DLL_PUBLIC Parser : public LIEF::Parser {
     void build_symbols(void);
     void build_signature(void);
     void build_overlay(void);
+    void build_dos_stub(void);
 
     ResourceNode* build_resource_node(
         const pe_resource_directory_table *directoryTable, uint32_t baseOffset);
 
 
     std::unique_ptr<VectorStream> stream_;
     Binary*                       binary_;
-    PE_TYPE           type_;
+    PE_TYPE                       type_;
 };
 
 

src/PE/Binary.cpp

@@ -994,6 +994,22 @@ std::vector<uint8_t>& Binary::overlay(void) {
   return const_cast<std::vector<uint8_t>&>(static_cast<const Binary*>(this)->overlay());
 }
 
+// Dos stub
+// ========
+
+const std::vector<uint8_t>& Binary::dos_stub(void) const {
+  return this->dos_stub_;
+}
+
+std::vector<uint8_t>& Binary::dos_stub(void) {
+  return const_cast<std::vector<uint8_t>&>(static_cast<const Binary*>(this)->dos_stub());
+}
+
+
+void Binary::dos_stub(const std::vector<uint8_t>& content) {
+  this->dos_stub_ = content;
+}
+
 // Resource manager
 // ===============
 

src/PE/Builder.cpp

@@ -43,7 +43,8 @@ Builder::Builder(Binary* binary) :
   build_relocations_{false},
   build_tls_{false},
   build_resources_{false},
-  build_overlay_{true}
+  build_overlay_{true},
+  build_dos_stub_{true}
 {}
 
 
@@ -76,6 +77,11 @@ Builder& Builder::build_overlay(bool flag) {
   return *this;
 }
 
+Builder& Builder::build_dos_stub(bool flag) {
+  this->build_dos_stub_ = flag;
+  return *this;
+}
+
 
 void Builder::write(const std::string& filename) const {
   std::ofstream output_file{filename, std::ios::out | std::ios::binary | std::ios::trunc};
@@ -459,6 +465,14 @@ Builder& Builder::operator<<(const DosHeader& dos_header) {
 
   this->ios_.seekp(0);
   this->ios_.write(reinterpret_cast<const uint8_t*>(&dosHeader), sizeof(pe_dos_header));
+  if (this->binary_->dos_stub().size() > 0 and this->build_dos_stub_) {
+
+    if (sizeof(pe_dos_header) + this->binary_->dos_stub().size() > dos_header.addressof_new_exeheader()) {
+      LOG(WARNING) << "Inconsistent 'addressof_new_exeheader' (0x" << std::hex << dos_header.addressof_new_exeheader();
+    }
+    this->ios_.write(this->binary_->dos_stub());
+  }
+
   return *this;
 }
 
@@ -542,12 +556,13 @@ Builder& Builder::operator<<(const Section& section) {
 std::ostream& operator<<(std::ostream& os, const Builder& b) {
   os << std::left;
   os << std::boolalpha;
-  os << std::setw(20) << "Builde imports:"     << b.build_imports_     << std::endl;
-  os << std::setw(20) << "Patch imports:"      << b.patch_imports_     << std::endl;
-  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;
+  os << std::setw(20) << "Build imports:"     << b.build_imports_     << std::endl;
+  os << std::setw(20) << "Patch imports:"     << b.patch_imports_     << std::endl;
+  os << std::setw(20) << "Build relocations:" << b.build_relocations_ << std::endl;
+  os << std::setw(20) << "Build TLS:"         << b.build_tls_         << std::endl;
+  os << std::setw(20) << "Build resources:"   << b.build_resources_   << std::endl;
+  os << std::setw(20) << "Build overlay:"     << b.build_overlay_     << std::endl;
+  os << std::setw(20) << "Build dos stub:"    << b.build_dos_stub_    << std::endl;
   return os;
 }
 

src/PE/Parser.cpp

@@ -83,6 +83,15 @@ void Parser::init(const std::string& name) {
 
 }
 
+void Parser::build_dos_stub(void) {
+  const DosHeader& dos_header = this->binary_->dos_header();
+  const uint64_t sizeof_dos_stub = dos_header.addressof_new_exeheader() - sizeof(pe_dos_header);
+
+  const uint8_t* ptr_to_dos_stub = reinterpret_cast<const uint8_t*>(this->stream_->read(
+        sizeof(pe_dos_header),
+        sizeof_dos_stub));
+  this->binary_->dos_stub_ = {ptr_to_dos_stub, ptr_to_dos_stub + sizeof_dos_stub};
+}
 
 
 

src/PE/Parser.tcc

@@ -27,6 +27,10 @@ void Parser::build(void) {
     LOG(WARNING) << e.what();
   }
 
+  LOG(DEBUG) << "[+] Retreive Dos stub";
+
+  this->build_dos_stub();
+
   LOG(DEBUG) << "[+] Decomposing Sections";
 
   try {