Skip to content

Commit 04dddd3

Browse files
committed
Fix alignment issue when removing a PE section
1 parent 1f875db commit 04dddd3

File tree

4 files changed

+86
-4
lines changed

4 files changed

+86
-4
lines changed

api/python/PE/objects/pyBinary.cpp

+5
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,11 @@ void create<Binary>(py::module& m) {
194194

195195
.def("remove_all_relocations", &Binary::remove_all_relocations)
196196

197+
.def("remove",
198+
static_cast<void(Binary::*)(const Section&, bool)>(&Binary::remove),
199+
"Remove the " RST_CLASS_REF(lief.PE.Section) " given in first parameter",
200+
"section"_a, "clear"_a = false)
201+
197202
.def_property_readonly("data_directories",
198203
static_cast<no_const_getter<it_data_directories>>(&Binary::data_directories),
199204
"Return an iterator on the " RST_CLASS_REF(lief.PE.DataDirectory) "",

include/LIEF/PE/Binary.hpp

+3
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,9 @@ class LIEF_API Binary : public LIEF::Binary {
202202
//! @param[in] name Name of section to delete
203203
virtual void remove_section(const std::string& name, bool clear = false) override;
204204

205+
//! Remove the given section
206+
void remove(const Section& section, bool clear = false);
207+
205208
//! @brief Add a section to the binary and return the section added.
206209
Section& add_section(
207210
const Section& section,

src/PE/Binary.cpp

+36-4
Original file line numberDiff line numberDiff line change
@@ -500,27 +500,59 @@ uint32_t Binary::sizeof_headers(void) const {
500500

501501
void Binary::remove_section(const std::string& name, bool clear) {
502502

503-
this->header().numberof_sections(this->header().numberof_sections() - 1);
504-
505-
this->optional_header().sizeof_headers(this->sizeof_headers());
506-
this->optional_header().sizeof_image(static_cast<uint32_t>(this->virtual_size()));
507503

508504
auto&& it_section = std::find_if(
509505
std::begin(this->sections_),
510506
std::end(this->sections_),
511507
[&name] (const Section* section) {
512508
return section->name() == name;
513509
});
510+
514511
if (it_section == std::end(this->sections_)) {
515512
LOG(ERROR) << "Unable to find section: '" << name << "'" << std::endl;
513+
return;
514+
}
515+
516+
return this->remove(**it_section, clear);
517+
}
518+
519+
520+
void Binary::remove(const Section& section, bool clear) {
521+
auto&& it_section = std::find_if(
522+
std::begin(this->sections_),
523+
std::end(this->sections_),
524+
[&section] (const Section* s) {
525+
return *s == section;
526+
});
527+
if (it_section == std::end(this->sections_)) {
528+
LOG(ERROR) << "Unable to find section: '" << section.name() << "'" << std::endl;
529+
return;
516530
}
517531

518532
Section* to_remove = *it_section;
533+
const size_t section_index = std::distance(std::begin(this->sections_), it_section);
534+
535+
if (section_index < (this->sections_.size() - 1) and section_index > 0) {
536+
Section* previous = this->sections_[section_index - 1];
537+
const size_t raw_size_gap = (to_remove->offset() + to_remove->size()) - (previous->offset() + previous->size());
538+
previous->size(previous->size() + raw_size_gap);
539+
540+
const size_t vsize_size_gap = (to_remove->virtual_address() + to_remove->virtual_size()) - (previous->virtual_address() + previous->virtual_size());
541+
previous->virtual_size(previous->virtual_size() + vsize_size_gap);
542+
}
543+
544+
519545
if (clear) {
520546
to_remove->clear(0);
521547
}
548+
522549
delete to_remove;
523550
this->sections_.erase(it_section);
551+
552+
this->header().numberof_sections(this->header().numberof_sections() - 1);
553+
554+
this->optional_header().sizeof_headers(this->sizeof_headers());
555+
this->optional_header().sizeof_image(static_cast<uint32_t>(this->virtual_size()));
524556
}
525557

526558
void Binary::make_space_for_new_section(void) {

tests/pe/test_pe.py

+42
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import ctypes
1515
import zipfile
1616
import json
17+
import re
1718

1819
from subprocess import Popen
1920

@@ -25,6 +26,14 @@ def setUp(self):
2526
self.logger = logging.getLogger(__name__)
2627
self.maxDiff = None
2728

29+
self.tmp_dir = tempfile.mkdtemp(suffix='_lief_tests')
30+
self.logger.debug("temp dir: {}".format(self.tmp_dir))
31+
32+
33+
if sys.platform.startswith("win"):
34+
SEM_NOGPFAULTERRORBOX = 0x0002 # From MSDN
35+
ctypes.windll.kernel32.SetErrorMode(SEM_NOGPFAULTERRORBOX);
36+
2837

2938
def test_code_view_pdb(self):
3039
path = get_sample('PE/PE64_x86-64_binary_ConsoleApplication1.exe')
@@ -61,6 +70,39 @@ def test_code_view_pdb(self):
6170
'type': 'CODEVIEW'
6271
})
6372

73+
def test_remove_section(self):
74+
path = get_sample('PE/PE64_x86-64_remove_section.exe')
75+
sample = lief.parse(path)
76+
77+
output = os.path.join(self.tmp_dir, "section_removed.exe")
78+
79+
sample.remove_section("lief")
80+
sample.write(output)
81+
82+
st = os.stat(output)
83+
os.chmod(output, st.st_mode | stat.S_IEXEC)
84+
85+
if sys.platform.startswith("win"):
86+
subprocess_flags = 0x8000000 # win32con.CREATE_NO_WINDOW?
87+
p = Popen([output], shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, creationflags=subprocess_flags)
88+
89+
stdout, _ = p.communicate()
90+
self.logger.debug(stdout.decode("utf8"))
91+
self.assertIn("Hello World", stdout)
92+
93+
94+
def tearDown(self):
95+
# Delete it
96+
try:
97+
if os.path.isdir(self.tmp_dir):
98+
shutil.rmtree(self.tmp_dir)
99+
except Exception as e:
100+
self.logger.error(e)
101+
102+
103+
104+
105+
64106

65107

66108

0 commit comments

Comments
 (0)