ELFFile: allow filtering by segment type in iter_segments() (#375)

This is very similar to the filtering implemented for
sections in commit d71faebcd58e.

diff --git a/elftools/elf/elffile.py b/elftools/elf/elffile.py
index acddc8952bfb212ad7d25e3c67eabdff20ae75c5..f248268815089422838fa2d5950818719e12db99 100644 (file)
--- a/elftools/elf/elffile.py
+++ b/elftools/elf/elffile.py
@@ -168,11 +168,17 @@ class ELFFile(object):
         segment_header = self._get_segment_header(n)
         return self._make_segment(segment_header)
 
-    def iter_segments(self):
-        """ Yield all the segments in the file
+    def iter_segments(self, type=None):
+        """ Yield all the segments in the file. If the optional |type|
+            parameter is passed, this method will only yield segments 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.
+            'PT_LOAD'.
         """
         for i in range(self.num_segments()):
-            yield self.get_segment(i)
+            segment = self.get_segment(i)
+            if type is None or segment['p_type'] == type:
+                yield segment
 
     def address_offsets(self, start, size=1):
         """ Yield a file offset for each ELF segment containing a memory region.
@@ -181,10 +187,8 @@ class ELFFile(object):
             offset of the region is yielded.
         """
         end = start + size
-        for seg in self.iter_segments():
-            # consider LOAD only to prevent same address being yielded twice
-            if seg['p_type'] != 'PT_LOAD':
-                continue
+        # consider LOAD only to prevent same address being yielded twice
+        for seg in self.iter_segments(type='PT_LOAD'):
             if (start >= seg['p_vaddr'] and
                 end <= seg['p_vaddr'] + seg['p_filesz']):
                 yield start - seg['p_vaddr'] + seg['p_offset']

diff --git a/test/test_elffile.py b/test/test_elffile.py
index 1edba94344c710bd689458e1088dc1e487b23870..cc929b9f5beff4c7a139ae0d3fccea645d391855 100644 (file)
--- a/test/test_elffile.py
+++ b/test/test_elffile.py
@@ -14,12 +14,18 @@ class TestMap(unittest.TestCase):
     def test_address_offsets(self):
         class MockELF(ELFFile):
             __init__ = object.__init__
-            def iter_segments(self):
-                return iter((
-                    dict(p_type='PT_PHDR', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
-                    dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
-                    dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
-                ))
+            def iter_segments(self, type=None):
+                if type == 'PT_LOAD':
+                    return iter((
+                        dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                    ))
+                else:
+                    return iter((
+                        dict(p_type='PT_PHDR', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10200, p_filesz=0x200, p_offset=0x100),
+                        dict(p_type='PT_LOAD', p_vaddr=0x10100, p_filesz=0x100, p_offset=0x400),
+                    ))
 
         elf = MockELF()