Skip to content

Commit 5666351

Browse files
committed
Improve API of PE's OptionalHeader
API changes: - 'has_dll_characteristics' renamed to 'has' - 'add' to add a characteristic - added - 'remove' to remove a characteristic - added - operator+= to add a characteristic - added - operator-= to remove a characteristic - added
1 parent 629d9aa commit 5666351

File tree

3 files changed

+50
-7
lines changed

3 files changed

+50
-7
lines changed

api/python/PE/objects/pyOptionalHeader.cpp

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,14 +206,24 @@ void init_PE_OptionalHeader_class(py::module& m) {
206206
static_cast<setter_t<uint32_t>>(&OptionalHeader::dll_characteristics),
207207
"The " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " characteristics")
208208

209+
.def("add",
210+
static_cast<void (OptionalHeader::*)(DLL_CHARACTERISTICS)>(&OptionalHeader::add),
211+
"Add the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) "",
212+
"characteristic"_a)
213+
214+
.def("remove",
215+
static_cast<void (OptionalHeader::*)(DLL_CHARACTERISTICS)>(&OptionalHeader::remove),
216+
"Remove the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) "",
217+
"characteristic"_a)
209218

210219
.def_property_readonly("dll_characteristics_lists",
211220
&OptionalHeader::dll_characteristics_list,
212221
"" RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " as a list")
213222

214-
.def("has_dll_characteristics",
215-
&OptionalHeader::has_dll_characteristics,
216-
"``True`` if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is in the :attr:`~lief.PE.OptionalHeader.dll_characteristics`",
223+
.def("has",
224+
static_cast<bool (OptionalHeader::*)(DLL_CHARACTERISTICS) const>(&OptionalHeader::has),
225+
"``True`` if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is in the "
226+
":attr:`~lief.PE.OptionalHeader.dll_characteristics`",
217227
"characteristics"_a)
218228

219229
.def_property("sizeof_stack_reserve",
@@ -265,6 +275,13 @@ void init_PE_OptionalHeader_class(py::module& m) {
265275
return LIEF::Hash::hash(optional_header);
266276
})
267277

278+
.def(py::self += DLL_CHARACTERISTICS())
279+
.def(py::self -= DLL_CHARACTERISTICS())
280+
281+
.def("__contains__",
282+
static_cast<bool (OptionalHeader::*)(DLL_CHARACTERISTICS) const>(&OptionalHeader::has),
283+
"Check if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is present")
284+
268285
.def("__str__", [] (const OptionalHeader& header)
269286
{
270287
std::ostringstream stream;

include/LIEF/PE/OptionalHeader.hpp

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,12 @@ class DLL_PUBLIC OptionalHeader : public Visitable {
6565
uint64_t sizeof_heap_commit(void) const;
6666
uint32_t loader_flags(void) const;
6767
uint32_t numberof_rva_and_size(void) const;
68-
bool has_dll_characteristics(DLL_CHARACTERISTICS c) const;
68+
bool has(DLL_CHARACTERISTICS c) const;
6969
std::set<DLL_CHARACTERISTICS> dll_characteristics_list(void) const;
7070

71+
void add(DLL_CHARACTERISTICS c);
72+
void remove(DLL_CHARACTERISTICS c);
73+
7174
void magic(PE_TYPE magic);
7275
void major_linker_version(uint8_t majorLinkerVersion);
7376
void minor_linker_version(uint8_t minorLinkerVersion);
@@ -101,6 +104,9 @@ class DLL_PUBLIC OptionalHeader : public Visitable {
101104

102105
virtual void accept(Visitor& visitor) const override;
103106

107+
OptionalHeader& operator+=(DLL_CHARACTERISTICS c);
108+
OptionalHeader& operator-=(DLL_CHARACTERISTICS c);
109+
104110
bool operator==(const OptionalHeader& rhs) const;
105111
bool operator!=(const OptionalHeader& rhs) const;
106112

src/PE/OptionalHeader.cpp

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -293,17 +293,26 @@ uint32_t OptionalHeader::numberof_rva_and_size(void) const {
293293
return this->numberOfRvaAndSize_;
294294
}
295295

296-
bool OptionalHeader::has_dll_characteristics(DLL_CHARACTERISTICS c) const {
297-
return (this->DLLCharacteristics_ & static_cast<uint32_t>(c)) > 0;
296+
bool OptionalHeader::has(DLL_CHARACTERISTICS c) const {
297+
return (this->dll_characteristics() & static_cast<uint32_t>(c)) > 0;
298298
}
299299

300+
void OptionalHeader::add(DLL_CHARACTERISTICS c) {
301+
this->dll_characteristics(this->dll_characteristics() | static_cast<uint32_t>(c));
302+
}
303+
304+
void OptionalHeader::remove(DLL_CHARACTERISTICS c) {
305+
this->dll_characteristics(this->dll_characteristics() & (~ static_cast<uint32_t>(c)));
306+
}
307+
308+
300309
std::set<DLL_CHARACTERISTICS> OptionalHeader::dll_characteristics_list(void) const {
301310
std::set<DLL_CHARACTERISTICS> dll_charac;
302311
std::copy_if(
303312
std::begin(dll_characteristics_array),
304313
std::end(dll_characteristics_array),
305314
std::inserter(dll_charac, std::begin(dll_charac)),
306-
std::bind(&OptionalHeader::has_dll_characteristics, this, std::placeholders::_1));
315+
std::bind(&OptionalHeader::has, this, std::placeholders::_1));
307316

308317
return dll_charac;
309318
}
@@ -499,6 +508,17 @@ void OptionalHeader::accept(LIEF::Visitor& visitor) const {
499508
visitor.visit(this->numberof_rva_and_size());
500509
}
501510

511+
512+
OptionalHeader& OptionalHeader::operator+=(DLL_CHARACTERISTICS c) {
513+
this->add(c);
514+
return *this;
515+
}
516+
517+
OptionalHeader& OptionalHeader::operator-=(DLL_CHARACTERISTICS c) {
518+
this->remove(c);
519+
return *this;
520+
}
521+
502522
bool OptionalHeader::operator==(const OptionalHeader& rhs) const {
503523
size_t hash_lhs = Hash::hash(*this);
504524
size_t hash_rhs = Hash::hash(rhs);

0 commit comments

Comments
 (0)