From: Andreas Ziegler <andreas.ziegler@fau.de>
Date: Tue, 12 Jan 2021 00:27:24 +0000 (+0100)
Subject: ELFFile: allow filtering of sections by type in iter_sections (#345)
X-Git-Tag: v0.28~23
X-Git-Url: https://git.libre-soc.org/?a=commitdiff_plain;h=d71faebcd58ef26ce973044f2a822a494c42a3c1;p=pyelftools.git

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 0b44b29..779ae3b 100644
--- 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 654b202..1edba94 100644
--- 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()