From a74b8abf22c6a7ceb804d0cd036c86c63da767a9 Mon Sep 17 00:00:00 2001 From: Andreas Ziegler Date: Mon, 25 Oct 2021 16:31:14 +0200 Subject: [PATCH] ELFFile: allow filtering by segment type in iter_segments() (#375) This is very similar to the filtering implemented for sections in commit d71faebcd58e. --- elftools/elf/elffile.py | 18 +++++++++++------- test/test_elffile.py | 18 ++++++++++++------ 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/elftools/elf/elffile.py b/elftools/elf/elffile.py index acddc89..f248268 100644 --- 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 1edba94..cc929b9 100644 --- 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() -- 2.30.2