ELFFile: allow filtering of sections by type in iter_sections (#345)

As stated in the corresponding issue, we can already filter
the output of Dynamic.iter_tags() by the type of the tag
we're looking for.

Let's adapt the iteration over the sections of the ELF file
so that it only yields sections of a certain type if the
optional type parameter is passed to iter_sections().

By doing this we can also simplify two call sites inside
the ELFFile class.

Fixes: #344

diff --git a/elftools/elf/elffile.py b/elftools/elf/elffile.py
index 0b44b295f577b98973e8c0ea1acc4e501aae4d09..779ae3b1475bddf9204c3b36d0a7ab8f035c6cd7 100644 (file)
--- a/elftools/elf/elffile.py
+++ b/elftools/elf/elffile.py
@@ -135,11 +135,17 @@ class ELFFile(object):
             self._make_section_name_map()
         return self._section_name_map.get(section_name, None)
 
-    def iter_sections(self):
-        """ Yield all the sections in the file
+    def iter_sections(self, type=None):
+        """ Yield all the sections in the file. If the optional |type|
+            parameter is passed, this method will only yield sections of the
+            given type. The parameter value must be a string containing the
+            name of the type as defined in the ELF specification, e.g.
+            'SHT_SYMTAB'.
         """
         for i in range(self.num_sections()):
-            yield self.get_section(i)
+            section = self.get_section(i)
+            if type is None or section['sh_type'] == type:
+                yield section
 
     def num_segments(self):
         """ Number of segments in the file
@@ -254,7 +260,7 @@ class ELFFile(object):
     def has_ehabi_info(self):
         """ Check whether this file appears to have arm exception handler index table.
         """
-        return any(s['sh_type'] == 'SHT_ARM_EXIDX' for s in self.iter_sections())
+        return any(self.iter_sections(type='SHT_ARM_EXIDX'))
 
     def get_ehabi_infos(self):
         """ Generally, shared library and executable contain 1 .ARM.exidx section.
@@ -265,9 +271,8 @@ class ELFFile(object):
         if self['e_type'] == 'ET_REL':
             # TODO: support relocatable file
             assert False, "Current version of pyelftools doesn't support relocatable file."
-        for section in self.iter_sections():
-            if section['sh_type'] == 'SHT_ARM_EXIDX':
-                _ret.append(EHABIInfo(section, self.little_endian))
+        for section in self.iter_sections(type='SHT_ARM_EXIDX'):
+            _ret.append(EHABIInfo(section, self.little_endian))
         return _ret if len(_ret) > 0 else None
 
     def get_machine_arch(self):

diff --git a/test/test_elffile.py b/test/test_elffile.py
index 654b2021f76196f47312e5da90ba16c2ba955f22..1edba94344c710bd689458e1088dc1e487b23870 100644 (file)
--- a/test/test_elffile.py
+++ b/test/test_elffile.py
@@ -5,6 +5,7 @@
 # This code is in the public domain
 #-------------------------------------------------------------------------------
 import unittest
+import os
 
 from elftools.elf.elffile import ELFFile
 
@@ -43,6 +44,16 @@ class TestMap(unittest.TestCase):
         self.assertEqual(tuple(elf.address_offsets(0x103FE, 4)), ())
         self.assertEqual(tuple(elf.address_offsets(0x10400, 4)), ())
 
+class TestSectionFilter(unittest.TestCase):
+
+    def test_section_filter(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'arm_exidx_test.so'), 'rb') as f:
+            elf = ELFFile(f)
+            self.assertEqual(len(list(elf.iter_sections())), 26)
+            self.assertEqual(len(list(elf.iter_sections('SHT_REL'))), 2)
+            self.assertEqual(len(list(elf.iter_sections('SHT_ARM_EXIDX'))), 1)
+            self.assertTrue(elf.has_ehabi_info())
 
 if __name__ == '__main__':
     unittest.main()