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

Author: romainthomas

api/python/PE/objects/pyOptionalHeader.cpp

@@ -206,14 +206,24 @@ void init_PE_OptionalHeader_class(py::module& m) {
         static_cast<setter_t<uint32_t>>(&OptionalHeader::dll_characteristics),
         "The " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " characteristics")
 
+    .def("add",
+        static_cast<void (OptionalHeader::*)(DLL_CHARACTERISTICS)>(&OptionalHeader::add),
+        "Add the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) "",
+        "characteristic"_a)
+
+    .def("remove",
+        static_cast<void (OptionalHeader::*)(DLL_CHARACTERISTICS)>(&OptionalHeader::remove),
+        "Remove the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) "",
+        "characteristic"_a)
 
     .def_property_readonly("dll_characteristics_lists",
         &OptionalHeader::dll_characteristics_list,
         "" RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " as a list")
 
-    .def("has_dll_characteristics",
-        &OptionalHeader::has_dll_characteristics,
-        "``True`` if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is in the :attr:`~lief.PE.OptionalHeader.dll_characteristics`",
+    .def("has",
+        static_cast<bool (OptionalHeader::*)(DLL_CHARACTERISTICS) const>(&OptionalHeader::has),
+        "``True`` if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is in the "
+        ":attr:`~lief.PE.OptionalHeader.dll_characteristics`",
         "characteristics"_a)
 
     .def_property("sizeof_stack_reserve",
@@ -265,6 +275,13 @@ void init_PE_OptionalHeader_class(py::module& m) {
           return LIEF::Hash::hash(optional_header);
         })
 
+    .def(py::self += DLL_CHARACTERISTICS())
+    .def(py::self -= DLL_CHARACTERISTICS())
+
+    .def("__contains__",
+        static_cast<bool (OptionalHeader::*)(DLL_CHARACTERISTICS) const>(&OptionalHeader::has),
+        "Check if the given " RST_CLASS_REF(lief.PE.DLL_CHARACTERISTICS) " is present")
+
     .def("__str__", [] (const OptionalHeader& header)
         {
           std::ostringstream stream;

include/LIEF/PE/OptionalHeader.hpp

@@ -65,9 +65,12 @@ class DLL_PUBLIC OptionalHeader : public Visitable {
     uint64_t  sizeof_heap_commit(void) const;
     uint32_t  loader_flags(void) const;
     uint32_t  numberof_rva_and_size(void) const;
-    bool      has_dll_characteristics(DLL_CHARACTERISTICS c) const;
+    bool      has(DLL_CHARACTERISTICS c) const;
     std::set<DLL_CHARACTERISTICS> dll_characteristics_list(void) const;
 
+    void add(DLL_CHARACTERISTICS c);
+    void remove(DLL_CHARACTERISTICS c);
+
     void magic(PE_TYPE magic);
     void major_linker_version(uint8_t majorLinkerVersion);
     void minor_linker_version(uint8_t minorLinkerVersion);
@@ -101,6 +104,9 @@ class DLL_PUBLIC OptionalHeader : public Visitable {
 
     virtual void accept(Visitor& visitor) const override;
 
+    OptionalHeader& operator+=(DLL_CHARACTERISTICS c);
+    OptionalHeader& operator-=(DLL_CHARACTERISTICS c);
+
     bool operator==(const OptionalHeader& rhs) const;
     bool operator!=(const OptionalHeader& rhs) const;
 

src/PE/OptionalHeader.cpp

@@ -293,17 +293,26 @@ uint32_t OptionalHeader::numberof_rva_and_size(void) const {
   return this->numberOfRvaAndSize_;
 }
 
-bool OptionalHeader::has_dll_characteristics(DLL_CHARACTERISTICS c) const {
-  return (this->DLLCharacteristics_ & static_cast<uint32_t>(c)) > 0;
+bool OptionalHeader::has(DLL_CHARACTERISTICS c) const {
+  return (this->dll_characteristics() & static_cast<uint32_t>(c)) > 0;
 }
 
+void OptionalHeader::add(DLL_CHARACTERISTICS c) {
+  this->dll_characteristics(this->dll_characteristics() | static_cast<uint32_t>(c));
+}
+
+void OptionalHeader::remove(DLL_CHARACTERISTICS c) {
+  this->dll_characteristics(this->dll_characteristics() & (~ static_cast<uint32_t>(c)));
+}
+
+
 std::set<DLL_CHARACTERISTICS> OptionalHeader::dll_characteristics_list(void) const {
   std::set<DLL_CHARACTERISTICS> dll_charac;
   std::copy_if(
       std::begin(dll_characteristics_array),
       std::end(dll_characteristics_array),
       std::inserter(dll_charac, std::begin(dll_charac)),
-      std::bind(&OptionalHeader::has_dll_characteristics, this, std::placeholders::_1));
+      std::bind(&OptionalHeader::has, this, std::placeholders::_1));
 
   return dll_charac;
 }
@@ -499,6 +508,17 @@ void OptionalHeader::accept(LIEF::Visitor& visitor) const {
   visitor.visit(this->numberof_rva_and_size());
 }
 
+
+OptionalHeader& OptionalHeader::operator+=(DLL_CHARACTERISTICS c) {
+  this->add(c);
+  return *this;
+}
+
+OptionalHeader& OptionalHeader::operator-=(DLL_CHARACTERISTICS c) {
+  this->remove(c);
+  return *this;
+}
+
 bool OptionalHeader::operator==(const OptionalHeader& rhs) const {
   size_t hash_lhs = Hash::hash(*this);
   size_t hash_rhs = Hash::hash(rhs);