added a method for returning the index of a section by name (#331)

* added an method for returning the index of a section by name

Signed-off-by: Jonathan <yonbruchim@gmail.com>
* changed naming of init mapping function

Signed-off-by: Jonathan <yonbruchim@gmail.com>
* Fixed CR

Added a test file containing 3 tests
1. test index of existing section
2. test index of missing section
3. test index of existing section when section_map_name is None

Signed-off-by: Jonathan Bruchim <yonbruchim@gmail.com>

diff --git a/elftools/elf/elffile.py b/elftools/elf/elffile.py
index 6a00e0a03b2ecebecb3b876a8c173dad34c44c83..0b44b295f577b98973e8c0ea1acc4e501aae4d09 100644 (file)
--- a/elftools/elf/elffile.py
+++ b/elftools/elf/elffile.py
@@ -120,12 +120,21 @@ class ELFFile(object):
         # mapping
         #
         if self._section_name_map is None:
-            self._section_name_map = {}
-            for i, sec in enumerate(self.iter_sections()):
-                self._section_name_map[sec.name] = i
+            self._make_section_name_map()
         secnum = self._section_name_map.get(name, None)
         return None if secnum is None else self.get_section(secnum)
 
+    def get_section_index(self, section_name):
+        """ Gets the index of the section by name. Return None if no such
+            section name exists.
+        """
+        # The first time this method is called, construct a name to number
+        # mapping
+        #
+        if self._section_name_map is None:
+            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
         """
@@ -574,6 +583,11 @@ class ELFFile(object):
         else:
             return Section(section_header, name, self)
 
+    def _make_section_name_map(self):
+        self._section_name_map = {}
+        for i, sec in enumerate(self.iter_sections()):
+            self._section_name_map[sec.name] = i
+
     def _make_symbol_table_section(self, section_header, name):
         """ Create a SymbolTableSection
         """

diff --git a/test/test_get_section_index.py b/test/test_get_section_index.py
new file mode 100644 (file)
index 0000000..0935b7d
--- /dev/null
+++ b/test/test_get_section_index.py
@@ -0,0 +1,60 @@
+#-------------------------------------------------------------------------------
+# Tests the functionality of get_section_index
+#
+# Jonathan Bruchim (YonBruchim@gmail.com)
+# This code is in the public domain
+#-------------------------------------------------------------------------------
+import unittest
+import os
+
+from elftools.elf.elffile import ELFFile
+
+
+class TestGetSectionIndex(unittest.TestCase):
+    def test_existing_section(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'simple_gcc.elf.arm'), 'rb') as f:
+            elf = ELFFile(f)
+
+            # Find the symbol table.
+            data_section_index = elf.get_section_index('.data')
+            self.assertIsNotNone(data_section_index)
+
+            # Test we can find a symbol by its name.
+            data_section = elf.get_section(data_section_index)
+            self.assertIsNotNone(data_section)
+
+            # Test it is actually the symbol we expect.
+            self.assertEqual(data_section.name, '.data')
+
+    def test_missing_section(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'simple_gcc.elf.arm'), 'rb') as f:
+            elf = ELFFile(f)
+
+            # try getting a missing section index
+            missing_section_index = elf.get_section_index('non-existent section')
+            self.assertIsNone(missing_section_index)
+
+    def test_uninitialized_section_name_map(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'simple_gcc.elf.arm'), 'rb') as f:
+            elf = ELFFile(f)
+
+            elf._section_name_map = None
+
+            # Find the symbol table.
+            data_section_index = elf.get_section_index('.data')
+            self.assertIsNotNone(data_section_index)
+
+            # Test we can find a symbol by its name.
+            data_section = elf.get_section(data_section_index)
+            self.assertIsNotNone(data_section)
+
+            # Test it is actually the symbol we expect.
+            self.assertEqual(data_section.name, '.data')
+
+
+if __name__ == '__main__':
+    unittest.main()
+