Enable to remove sections using the abstract layer

romainthomas · romainthomas · commit 918438c6bee5 · 2018-07-09T15:08:39.000+02:00
diff --git a/api/python/Abstract/objects/pyBinary.cpp b/api/python/Abstract/objects/pyBinary.cpp
@@ -71,6 +71,11 @@ void init_LIEF_Binary_class(py::module& m) {
         &Binary::entrypoint,
         "Binary's entrypoint")
 
+    .def("remove_section",
+        static_cast<void (Binary::*)(const std::string&, bool)>(&Binary::remove_section),
+        "Remove the section with the given name",
+        "name"_a, "clear"_a = false)
+
     .def_property_readonly("sections",
         static_cast<it_t<it_sections>>(&Binary::sections),
         "Return a list in **read only** of binary's abstract " RST_CLASS_REF(lief.Section) "",
diff --git a/api/python/ELF/objects/pyBinary.cpp b/api/python/ELF/objects/pyBinary.cpp
@@ -403,11 +403,6 @@ void create<Binary>(py::module& m) {
         "Remove the given library",
         "library_name"_a)
 
-    .def("remove_section",
-        &Binary::remove_section,
-        "Remove the given section from its name",
-        "section_name"_a, "clear"_a = false)
-
     .def("get_library",
         static_cast<no_const_func<DynamicEntryLibrary&, const std::string&>>(&Binary::get_library),
         "Return the " RST_CLASS_REF(lief.ELF.DynamicEntryLibrary) " with the given ``name``",
diff --git a/api/python/MachO/objects/pySection.cpp b/api/python/MachO/objects/pySection.cpp
@@ -97,6 +97,15 @@ void create<Section>(py::module& m) {
         static_cast<getter_t<Section::flag_list_t>>(&Section::flags_list),
         py::return_value_policy::reference_internal)
 
+    .def_property_readonly("segment",
+        static_cast<SegmentCommand& (Section::*)(void)>(&Section::segment),
+        "" RST_CLASS_REF(lief.MachO.SegmentCommand) " segment associated with the section",
+        py::return_value_policy::reference)
+
+    .def_property_readonly("has_segment",
+        &Section::has_segment,
+        "True if the current section has a segment associated with")
+
     .def("has",
         static_cast<bool(Section::*)(MACHO_SECTION_FLAGS) const>(&Section::has),
         "Check if the section has the given " RST_CLASS_REF(lief.MachO.SECTION_FLAGS) "",
diff --git a/api/python/PE/objects/pyBinary.cpp b/api/python/PE/objects/pyBinary.cpp
@@ -178,7 +178,6 @@ void create<Binary>(py::module& m) {
         "section"_a, py::arg("type") = PE_SECTION_TYPES::UNKNOWN,
         py::return_value_policy::reference)
 
-    //.def("delete_section", (void (Binary::*)(const std::string&)) &Binary::delete_section)
     //.def("get_import_section",
     //    static_cast<no_const_getter<Section&>>(&Binary::get_import_section),
     //    py::return_value_policy::reference_internal)
diff --git a/include/LIEF/Abstract/Binary.hpp b/include/LIEF/Abstract/Binary.hpp
@@ -77,6 +77,8 @@ class LIEF_API Binary : public Object {
   it_sections       sections(void);
   it_const_sections sections(void) const;
 
+  virtual void remove_section(const std::string& name, bool clear = false) = 0;
+
   //! @brief Returns binary's relocations
   it_relocations       relocations(void);
   it_const_relocations relocations(void) const;
diff --git a/include/LIEF/ELF/Binary.hpp b/include/LIEF/ELF/Binary.hpp
@@ -374,7 +374,7 @@ class LIEF_API Binary : public LIEF::Binary {
   //!
   //! We clear data used by this section and it's removed from
   //! section table
-  void remove_section(const std::string& name, bool clear = false);
+  virtual void remove_section(const std::string& name, bool clear = false) override;
 
   //! @brief Reconstruct the binary object and write it in `filename`
   //! @param filename Path to write the reconstructed binary
diff --git a/include/LIEF/MachO/Binary.hpp b/include/LIEF/MachO/Binary.hpp
@@ -154,6 +154,8 @@ class LIEF_API Binary : public LIEF::Binary  {
   //! or the last one
   Section* add_section(const SegmentCommand& segment, const Section& section);
 
+  virtual void remove_section(const std::string& name, bool clear = false) override;
+
   //! Remove the given command
   bool remove(const LoadCommand& command);
 
diff --git a/include/LIEF/MachO/Section.hpp b/include/LIEF/MachO/Section.hpp
@@ -86,6 +86,12 @@ class LIEF_API Section : public LIEF::Section {
   flag_list_t             flags_list(void) const;
   uint32_t                raw_flags(void) const;
 
+  bool                    has_segment(void) const;
+  SegmentCommand&         segment(void);
+  const SegmentCommand&   segment(void) const;
+
+  void clear(uint8_t v);
+
   it_relocations relocations(void);
   it_const_relocations relocations(void) const;
 
diff --git a/include/LIEF/PE/Binary.hpp b/include/LIEF/PE/Binary.hpp
@@ -200,7 +200,7 @@ class LIEF_API Binary : public LIEF::Binary {
   //! @brief Delete the section with the given name
   //!
   //! @param[in] name Name of section to delete
-  void  delete_section(const std::string& name);
+  virtual void remove_section(const std::string& name, bool clear = false) override;
 
   //! @brief Add a section to the binary and return the section added.
   Section& add_section(
diff --git a/include/LIEF/PE/Section.hpp b/include/LIEF/PE/Section.hpp
@@ -73,6 +73,7 @@ class LIEF_API Section : public LIEF::Section {
     const std::set<PE_SECTION_TYPES>& types(void) const;
     bool                              has_characteristic(SECTION_CHARACTERISTICS c) const;
     std::set<SECTION_CHARACTERISTICS> characteristics_list(void) const;
+    void clear(uint8_t c);
 
 
     virtual void name(const std::string& name) override;
diff --git a/src/MachO/Binary.cpp b/src/MachO/Binary.cpp
@@ -1005,6 +1005,49 @@ bool Binary::extend_segment(const SegmentCommand& segment, size_t size) {
   return true;
 }
 
+void Binary::remove_section(const std::string& name, bool clear) {
+  if (not this->has_section(name)) {
+    LOG(WARNING) << "Section '" << name << "' not found!";
+    return;
+  }
+
+  Section& sec_to_delete = this->get_section(name);
+  SegmentCommand& segment = sec_to_delete.segment();
+
+  if (clear) {
+    sec_to_delete.clear(0);
+  }
+
+
+  segment.numberof_sections(segment.numberof_sections() - 1);
+  auto&& it_section = std::find_if(
+      std::begin(segment.sections_),
+      std::end(segment.sections_),
+      [&sec_to_delete] (const Section* s) {
+        return *s == sec_to_delete;
+      });
+  CHECK_NE(it_section, std::end(segment.sections_));
+
+  const size_t lc_offset = segment.command_offset();
+  const size_t section_struct_size = this->is64_ ? sizeof(section_64) : sizeof(section_32);
+  segment.size_ -= section_struct_size;
+
+  this->header().sizeof_cmds(this->header().sizeof_cmds() - section_struct_size);
+
+  for (LoadCommand* lc : this->commands_) {
+    if (lc->command_offset() > lc_offset) {
+      lc->command_offset(lc->command_offset() - section_struct_size);
+    }
+  }
+
+  this->available_command_space_ += section_struct_size;
+
+
+
+  delete *it_section;
+  segment.sections_.erase(it_section);
+}
+
 Section* Binary::add_section(const Section& section) {
   SegmentCommand* __TEXT_segment = this->get_segment("__TEXT");
   if (__TEXT_segment == nullptr) {
diff --git a/src/MachO/Builder.tcc b/src/MachO/Builder.tcc
@@ -54,6 +54,7 @@ void Builder::build_segments(void) {
     //this->raw_.seekp(segment.file_offset());
     //this->raw_.write(content);
 
+    //const size_t original_size = segment.originalData_.size();
     segment.originalData_.clear();
 
     std::move(
@@ -104,7 +105,6 @@ void Builder::build_segments(void) {
         std::back_inserter(segment.originalData_));
 
     }
-
   }
 
 } // build_segment
diff --git a/src/MachO/Section.cpp b/src/MachO/Section.cpp
@@ -341,6 +341,28 @@ Section& Section::operator-=(MACHO_SECTION_FLAGS flag) {
 }
 
 
+void Section::clear(uint8_t v) {
+  Section::content_t clear(this->size(), v);
+  this->content(std::move(clear));
+}
+
+
+bool Section::has_segment(void) const {
+  return this->segment_ != nullptr;
+}
+
+SegmentCommand& Section::segment(void) {
+  return const_cast<SegmentCommand&>(static_cast<const Section*>(this)->segment());
+}
+
+const SegmentCommand& Section::segment(void) const {
+  if (not this->has_segment()) {
+    throw not_found("No segment associated with this section");
+  }
+  return *this->segment_;
+}
+
+
 void Section::accept(Visitor& visitor) const {
   visitor.visit(*this);
 }
diff --git a/src/PE/Binary.cpp b/src/PE/Binary.cpp
@@ -498,23 +498,29 @@ uint32_t Binary::sizeof_headers(void) const {
 
 }
 
-void Binary::delete_section(const std::string& name) {
-  Section& section_to_delete = this->get_section(name);
+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()));
 
-  this->sections_.erase(
-      std::remove_if(
-        std::begin(this->sections_),
-        std::end(this->sections_),
-        [&section_to_delete](const Section* section)
-        {
-          return section->name() == section_to_delete.name();
-        }),
-      std::end(this->sections_));
+  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;
+  }
+
+  Section* to_remove = *it_section;
+  if (clear) {
+    to_remove->clear(0);
+  }
+  delete to_remove;
+  this->sections_.erase(it_section);
 }
 
 void Binary::make_space_for_new_section(void) {
diff --git a/src/PE/Section.cpp b/src/PE/Section.cpp
@@ -233,6 +233,14 @@ void Section::accept(LIEF::Visitor& visitor) const {
   visitor.visit(*this);
 }
 
+
+void Section::clear(uint8_t c) {
+  std::fill(
+      std::begin(this->content_),
+      std::end(this->content_),
+      c);
+}
+
 bool Section::operator==(const Section& rhs) const {
   size_t hash_lhs = Hash::hash(*this);
   size_t hash_rhs = Hash::hash(rhs);
diff --git a/tests/elf/CMakeLists.txt b/tests/elf/CMakeLists.txt
@@ -228,6 +228,10 @@ if (PYTHON_TESTS_ENABLED)
     ${PYTHON_EXECUTABLE}
     "${CMAKE_CURRENT_SOURCE_DIR}/modify_relocations.py")
 
+  ADD_PYTHON_TEST(ELF_PYTHON_remove_section
+    ${PYTHON_EXECUTABLE}
+    "${CMAKE_CURRENT_SOURCE_DIR}/remove_section.py")
+
   # Examples
   # --------
   ADD_PYTHON_TEST(EXAMPLE_PYTHON_elf_reader_ls
diff --git a/tests/elf/remove_section.py b/tests/elf/remove_section.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+import unittest
+import logging
+import os
+import sys
+import stat
+import re
+import subprocess
+import tempfile
+import shutil
+from subprocess import Popen
+
+import lief
+from lief.ELF import Section
+
+from unittest import TestCase
+from utils import get_sample
+
+CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
+
+class TestRemoveSection(TestCase):
+    def setUp(self):
+        self.logger = logging.getLogger(__name__)
+        self.tmp_dir = tempfile.mkdtemp(suffix='_lief_test_section')
+        self.logger.debug("temp dir: {}".format(self.tmp_dir))
+
+
+    @unittest.skipUnless(sys.platform.startswith("linux"), "requires Linux")
+    def test_simple(self):
+        sample_path = get_sample('ELF/ELF64_x86-64_binary_ls.bin')
+        output      = os.path.join(self.tmp_dir, "ls.section")
+
+        ls = lief.parse(sample_path)
+        ls.remove_section(".text", clear=False)
+        ls.write(output)
+
+        st = os.stat(output)
+        os.chmod(output, st.st_mode | stat.S_IEXEC)
+
+        p = Popen([output, "--help"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        stdout, _ = p.communicate()
+        self.logger.debug(stdout.decode("utf8"))
+        self.assertIsNotNone(re.search(r'GNU coreutils', stdout.decode("utf8")))
+
+
+    def tearDown(self):
+        # Delete it
+        if os.path.isdir(self.tmp_dir):
+            shutil.rmtree(self.tmp_dir)
+
+if __name__ == '__main__':
+
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.DEBUG)
+
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.DEBUG)
+    root_logger.addHandler(ch)
+
+    unittest.main(verbosity=2)
+
diff --git a/tests/macho/macho_builder.py b/tests/macho/macho_builder.py
@@ -344,6 +344,23 @@ def test_ssh(self):
             self.assertIsNotNone(re.search(r'Hello World!', stdout))
 
 
+class TestRemoveSection(TestCase):
+    def setUp(self):
+        self.logger = logging.getLogger(__name__)
+
+    def test_simple(self):
+        original = lief.parse(get_sample('MachO/MachO64_x86-64_binary_section_to_remove.bin'))
+        _, output = tempfile.mkstemp(prefix="lief_sec_remove_")
+
+        original.remove_section("__to_remove")
+        original.write(output)
+
+        if sys.platform.startswith("darwin"):
+            stdout = run_program(output)
+            self.logger.debug(stdout)
+            self.assertIsNotNone(re.search(r'Hello World', stdout))
+
+
 if __name__ == '__main__':
 
     root_logger = logging.getLogger()
