Skip to content

Commit

Permalink
Fix #20
Browse files Browse the repository at this point in the history
  • Loading branch information
romainthomas committed Apr 21, 2017
1 parent eb1c0d3 commit e0634c1
Show file tree
Hide file tree
Showing 9 changed files with 123 additions and 30 deletions.
5 changes: 5 additions & 0 deletions api/python/PE/objects/pyBinary.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,11 @@ void init_PE_Binary_class(py::module& m) {
.def_property_readonly("resources_manager",
static_cast<no_const_getter<ResourcesManager>>(&Binary::get_resources_manager))

.def_property_readonly("overlay",
static_cast<no_const_getter<std::vector<uint8_t>&>>(&Binary::overlay),
"Return the overlay content",
py::return_value_policy::reference)

.def("add_import_function",
&Binary::add_import_function,
py::return_value_policy::reference)
Expand Down
4 changes: 4 additions & 0 deletions api/python/PE/objects/pyBuilder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,10 @@ void init_PE_Builder_class(py::module& m) {
&Builder::build_imports,
py::return_value_policy::reference)

.def("build_overlay",
static_cast<Builder& (Builder::*)(bool)>(&Builder::build_overlay),
py::return_value_policy::reference)

.def("write",
&Builder::write)

Expand Down
62 changes: 36 additions & 26 deletions include/LIEF/PE/Binary.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,15 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
Debug& get_debug(void);
const Debug& get_debug(void) const;

// =======
// Overlay
// =======

//! @brief Return the overlay content
const std::vector<uint8_t>& overlay(void) const;
std::vector<uint8_t>& overlay(void);


// =========================
// Methods to manage Imports
// =========================
Expand Down Expand Up @@ -293,32 +302,33 @@ class DLL_PUBLIC Binary : public LIEF::Binary {
void update_lookup_address_table_offset(void);
void update_iat(void);

PE_TYPE type_;
DosHeader dos_header_;
Header header_;
OptionalHeader optional_header_;

bool has_tls_;
bool has_imports_;
bool has_signature_;
bool has_exports_;
bool has_resources_;
bool has_exceptions_;
bool has_relocations_;
bool has_debug_;
bool has_configuration_;

Signature signature_;
TLS tls_;
sections_t sections_;
data_directories_t data_directories_;
symbols_t symbols_;
strings_table_t strings_table_;
relocations_t relocations_;
ResourceNode* resources_;
imports_t imports_;
Export export_;
Debug debug_;
PE_TYPE type_;
DosHeader dos_header_;
Header header_;
OptionalHeader optional_header_;

bool has_tls_;
bool has_imports_;
bool has_signature_;
bool has_exports_;
bool has_resources_;
bool has_exceptions_;
bool has_relocations_;
bool has_debug_;
bool has_configuration_;

Signature signature_;
TLS tls_;
sections_t sections_;
data_directories_t data_directories_;
symbols_t symbols_;
strings_table_t strings_table_;
relocations_t relocations_;
ResourceNode* resources_;
imports_t imports_;
Export export_;
Debug debug_;
std::vector<uint8_t> overlay_;

};

Expand Down
3 changes: 3 additions & 0 deletions include/LIEF/PE/Builder.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class DLL_PUBLIC Builder
Builder& build_relocations(bool flag = true);
Builder& build_tls(bool flag = true);
Builder& build_resources(bool flag);
Builder& build_overlay(bool flag);

const std::vector<uint8_t>& get_build(void);
void write(const std::string& filename) const;
Expand Down Expand Up @@ -83,6 +84,7 @@ class DLL_PUBLIC Builder
void build_string_table(void);
void build_relocation(void);
void build_resources(void);
void build_overlay(void);

void compute_resources_size(
ResourceNode *node,
Expand All @@ -108,6 +110,7 @@ class DLL_PUBLIC Builder
bool build_relocations_;
bool build_tls_;
bool build_resources_;
bool build_overlay_;

};

Expand Down
1 change: 1 addition & 0 deletions include/LIEF/PE/Parser.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ class DLL_PUBLIC Parser : public LIEF::Parser {
void build_string_table(void);
void build_symbols(void);
void build_signature(void);
void build_overlay(void);

ResourceNode* build_resource_node(
const pe_resource_directory_table *directoryTable, uint32_t baseOffset);
Expand Down
11 changes: 11 additions & 0 deletions src/PE/Binary.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -925,6 +925,17 @@ std::vector<uint8_t> Binary::get_content_from_virtual_address(uint64_t virtual_a

}

// Overlay
// =======

const std::vector<uint8_t>& Binary::overlay(void) const {
return this->overlay_;
}

std::vector<uint8_t>& Binary::overlay(void) {
return const_cast<std::vector<uint8_t>&>(static_cast<const Binary*>(this)->overlay());
}

// Resource manager
// ===============

Expand Down
37 changes: 33 additions & 4 deletions src/PE/Builder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <algorithm>
#include <fstream>
#include <iterator>
#include <numeric>

#include "easylogging++.h"

Expand All @@ -41,7 +42,8 @@ Builder::Builder(Binary* binary) :
patch_imports_{false},
build_relocations_{false},
build_tls_{false},
build_resources_{false}
build_resources_{false},
build_overlay_{true}
{}


Expand Down Expand Up @@ -69,7 +71,10 @@ Builder& Builder::build_resources(bool flag) {
return *this;
}


Builder& Builder::build_overlay(bool flag) {
this->build_overlay_ = flag;
return *this;
}


void Builder::write(const std::string& filename) const {
Expand All @@ -85,8 +90,6 @@ void Builder::write(const std::string& filename) const {
}




void Builder::build(void) {

LOG(DEBUG) << "Rebuilding" << std::endl;
Expand Down Expand Up @@ -151,6 +154,12 @@ void Builder::build(void) {

LOG(DEBUG) << "[+] Rebuilding string table" << std::endl;
//this->build_string_table();

if (this->binary_->overlay().size() > 0 and this->build_overlay_) {
this->build_overlay();
}


}

const std::vector<uint8_t>& Builder::get_build(void) {
Expand Down Expand Up @@ -402,6 +411,25 @@ void Builder::build_string_table(void) {
//TODO
}

void Builder::build_overlay(void) {
LOG(DEBUG) << "Building overlay";

const uint64_t last_section_offset = std::accumulate(
std::begin(this->binary_->sections_),
std::end(this->binary_->sections_), 0,
[] (uint64_t offset, const Section* section) {
return std::max<uint64_t>(section->offset() + section->size(), offset);
});

LOG(DEBUG) << "Overlay offset: 0x" << std::hex << last_section_offset;
LOG(DEBUG) << "Overlay size: " << std::dec << this->binary_->overlay().size();

const size_t saved_offset = this->ios_.tellp();
this->ios_.seekp(last_section_offset);
this->ios_.write(this->binary_->overlay());
this->ios_.seekp(saved_offset);
}

Builder& Builder::operator<<(const DosHeader& dos_header) {

pe_dos_header dosHeader;
Expand Down Expand Up @@ -519,6 +547,7 @@ std::ostream& operator<<(std::ostream& os, const Builder& b) {
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;
return os;
}

Expand Down
28 changes: 28 additions & 0 deletions src/PE/Parser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <iterator>
#include <iostream>
#include <string>
#include <numeric>

#include <mbedtls/platform.h>
#include <mbedtls/oid.h>
Expand Down Expand Up @@ -1018,6 +1019,33 @@ void Parser::build_signature(void) {
this->binary_->has_signature_ = true;
}


void Parser::build_overlay(void) {
LOG(DEBUG) << "Parsing Overlay";
const uint64_t last_section_offset = std::accumulate(
std::begin(this->binary_->sections_),
std::end(this->binary_->sections_), 0,
[] (uint64_t offset, const Section* section) {
return std::max<uint64_t>(section->offset() + section->size(), offset);
});

LOG(DEBUG) << "Overlay offset: 0x" << std::hex << last_section_offset;

const uint64_t overlay_size = this->stream_->size() - last_section_offset;

LOG(DEBUG) << "Overlay size: " << std::dec << overlay_size;

const uint8_t* ptr_to_overlay = reinterpret_cast<const uint8_t*>(this->stream_->read(
last_section_offset,
overlay_size));

this->binary_->overlay_ = {
ptr_to_overlay,
ptr_to_overlay + overlay_size
};

}

//
// Return the Binary constructed
//
Expand Down
2 changes: 2 additions & 0 deletions src/PE/Parser.tcc
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@ void Parser::build(void) {
} catch (const corrupted& e) {
LOG(WARNING) << e.what();
}

this->build_overlay();
}

template<typename PE_T>
Expand Down

0 comments on commit e0634c1

Please sign in to comment.