Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
Fix alignment issue when removing a PE section
  • Loading branch information
romainthomas committed Aug 10, 2018
1 parent 1f875db commit 04dddd3
Show file tree
Hide file tree
Showing 4 changed files with 86 additions and 4 deletions.
5 changes: 5 additions & 0 deletions api/python/PE/objects/pyBinary.cpp
Expand Up @@ -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) "",
Expand Down
3 changes: 3 additions & 0 deletions include/LIEF/PE/Binary.hpp
Expand Up @@ -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,
Expand Down
40 changes: 36 additions & 4 deletions src/PE/Binary.cpp
Expand Up @@ -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) {
Expand Down
42 changes: 42 additions & 0 deletions tests/pe/test_pe.py
Expand Up @@ -14,6 +14,7 @@
import ctypes
import zipfile
import json
import re

from subprocess import Popen

Expand All @@ -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')
Expand Down Expand Up @@ -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)








Expand Down

0 comments on commit 04dddd3

Please sign in to comment.