Fix alignment issue when removing a PE section

Author: romainthomas

api/python/PE/objects/pyBinary.cpp

@@ -194,6 +194,11 @@ void create<Binary>(py::module& m) {
 
     .def("remove_all_relocations", &Binary::remove_all_relocations)
 
+    .def("remove",
+        static_cast<void(Binary::*)(const Section&, bool)>(&Binary::remove),
+        "Remove the " RST_CLASS_REF(lief.PE.Section) " given in first parameter",
+        "section"_a, "clear"_a = false)
+
     .def_property_readonly("data_directories",
         static_cast<no_const_getter<it_data_directories>>(&Binary::data_directories),
         "Return an iterator on the " RST_CLASS_REF(lief.PE.DataDirectory) "",

include/LIEF/PE/Binary.hpp

@@ -202,6 +202,9 @@ class LIEF_API Binary : public LIEF::Binary {
   //! @param[in] name Name of section to delete
   virtual void remove_section(const std::string& name, bool clear = false) override;
 
+  //! Remove the given section
+  void remove(const Section& section, bool clear = false);
+
   //! @brief Add a section to the binary and return the section added.
   Section& add_section(
       const Section& section,

src/PE/Binary.cpp

@@ -500,27 +500,59 @@ uint32_t Binary::sizeof_headers(void) const {
 
 void Binary::remove_section(const std::string& name, bool clear) {
 
-  this->header().numberof_sections(this->header().numberof_sections() - 1);
-
-  this->optional_header().sizeof_headers(this->sizeof_headers());
-  this->optional_header().sizeof_image(static_cast<uint32_t>(this->virtual_size()));
 
   auto&& it_section = std::find_if(
       std::begin(this->sections_),
       std::end(this->sections_),
       [&name] (const Section* section) {
         return section->name() == name;
       });
+
   if (it_section == std::end(this->sections_)) {
     LOG(ERROR) << "Unable to find section: '" << name << "'" << std::endl;
+    return;
+  }
+
+  return this->remove(**it_section, clear);
+}
+
+
+void Binary::remove(const Section& section, bool clear) {
+  auto&& it_section = std::find_if(
+      std::begin(this->sections_),
+      std::end(this->sections_),
+      [&section] (const Section* s) {
+        return *s == section;
+      });
+  if (it_section == std::end(this->sections_)) {
+    LOG(ERROR) << "Unable to find section: '" << section.name() << "'" << std::endl;
+    return;
   }
 
   Section* to_remove = *it_section;
+  const size_t section_index = std::distance(std::begin(this->sections_), it_section);
+
+  if (section_index < (this->sections_.size() - 1) and section_index > 0) {
+    Section* previous = this->sections_[section_index - 1];
+    const size_t raw_size_gap = (to_remove->offset() + to_remove->size()) - (previous->offset() + previous->size());
+    previous->size(previous->size() + raw_size_gap);
+
+    const size_t vsize_size_gap = (to_remove->virtual_address() + to_remove->virtual_size()) - (previous->virtual_address() + previous->virtual_size());
+    previous->virtual_size(previous->virtual_size() + vsize_size_gap);
+  }
+
+
   if (clear) {
     to_remove->clear(0);
   }
+
   delete to_remove;
   this->sections_.erase(it_section);
+
+  this->header().numberof_sections(this->header().numberof_sections() - 1);
+
+  this->optional_header().sizeof_headers(this->sizeof_headers());
+  this->optional_header().sizeof_image(static_cast<uint32_t>(this->virtual_size()));
 }
 
 void Binary::make_space_for_new_section(void) {

tests/pe/test_pe.py

@@ -14,6 +14,7 @@
 import ctypes
 import zipfile
 import json
+import re
 
 from subprocess import Popen
 
@@ -25,6 +26,14 @@ def setUp(self):
         self.logger = logging.getLogger(__name__)
         self.maxDiff = None
 
+        self.tmp_dir = tempfile.mkdtemp(suffix='_lief_tests')
+        self.logger.debug("temp dir: {}".format(self.tmp_dir))
+
+
+        if sys.platform.startswith("win"):
+            SEM_NOGPFAULTERRORBOX = 0x0002 # From MSDN
+            ctypes.windll.kernel32.SetErrorMode(SEM_NOGPFAULTERRORBOX);
+
 
     def test_code_view_pdb(self):
         path = get_sample('PE/PE64_x86-64_binary_ConsoleApplication1.exe')
@@ -61,6 +70,39 @@ def test_code_view_pdb(self):
             'type': 'CODEVIEW'
         })
 
+    def test_remove_section(self):
+        path = get_sample('PE/PE64_x86-64_remove_section.exe')
+        sample = lief.parse(path)
+
+        output = os.path.join(self.tmp_dir, "section_removed.exe")
+
+        sample.remove_section("lief")
+        sample.write(output)
+
+        st = os.stat(output)
+        os.chmod(output, st.st_mode | stat.S_IEXEC)
+
+        if sys.platform.startswith("win"):
+            subprocess_flags = 0x8000000 # win32con.CREATE_NO_WINDOW?
+            p = Popen([output], shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, creationflags=subprocess_flags)
+
+            stdout, _ = p.communicate()
+            self.logger.debug(stdout.decode("utf8"))
+            self.assertIn("Hello World", stdout)
+
+
+    def tearDown(self):
+        # Delete it
+        try:
+            if os.path.isdir(self.tmp_dir):
+                shutil.rmtree(self.tmp_dir)
+        except Exception as e:
+            self.logger.error(e)
+
+
+
+
+