From 46ae4bd3a5f7b0e355913b3497f0e96ac2778ef5 Mon Sep 17 00:00:00 2001
From: Matthew Fernandez <matthew.fernandez@gmail.com>
Date: Mon, 24 Nov 2014 14:26:41 +1100
Subject: [PATCH] ELF: Add `get_symbol_by_name` functionality.

This commit implements the equivalent of `get_section_by_name` for symbols for
ELF files.
---
 elftools/elf/sections.py        | 15 +++++++++++
 test/test_get_symbol_by_name.py | 45 +++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 test/test_get_symbol_by_name.py

diff --git a/elftools/elf/sections.py b/elftools/elf/sections.py
index 1380d6b..83b5781 100644
--- a/elftools/elf/sections.py
+++ b/elftools/elf/sections.py
@@ -81,6 +81,7 @@ class SymbolTableSection(Section):
                 'Expected entry size of section %r to be > 0' % name)
         elf_assert(self['sh_size'] % self['sh_entsize'] == 0,
                 'Expected section size to be a multiple of entry size in section %r' % name)
+        self._symbol_name_map = None
 
     def num_symbols(self):
         """ Number of symbols in the table
@@ -100,6 +101,20 @@ class SymbolTableSection(Section):
         name = self.stringtable.get_string(entry['st_name'])
         return Symbol(entry, name)
 
+    def get_symbol_by_name(self, name):
+        """ Get a symbol by its name. Return None if no symbol by the given
+            name exists.
+        """
+        # The first time this method is called, construct a name to number
+        # mapping
+        #
+        if self._symbol_name_map is None:
+            self._symbol_name_map = {}
+            for i, sym in enumerate(self.iter_symbols()):
+                self._symbol_name_map[sym.name] = i
+        symnum = self._symbol_name_map.get(name)
+        return None if symnum is None else self.get_symbol(symnum)
+
     def iter_symbols(self):
         """ Yield all the symbols in the table
         """
diff --git a/test/test_get_symbol_by_name.py b/test/test_get_symbol_by_name.py
new file mode 100644
index 0000000..cad104f
--- /dev/null
+++ b/test/test_get_symbol_by_name.py
@@ -0,0 +1,45 @@
+# Tests the functionality of the ELF file function `get_symbol_by_name`.
+
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest
+import os
+
+from utils import setup_syspath; setup_syspath()
+from elftools.elf.elffile import ELFFile
+
+class TestGetSymbolByName(unittest.TestCase):
+    def test_existing_symbol(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'simple_gcc.elf.arm'), 'rb') as f:
+            elf = ELFFile(f)
+
+            # Find the symbol table.
+            symtab = elf.get_section_by_name(b'.symtab')
+            self.assertIsNotNone(symtab)
+
+            # Test we can find a symbol by its name.
+            main = symtab.get_symbol_by_name(b'main')
+            self.assertIsNotNone(main)
+
+            # Test it is actually the symbol we expect.
+            self.assertEqual(main.name, b'main')
+            self.assertEqual(main['st_value'], 0x8068)
+            self.assertEqual(main['st_size'], 0x28)
+
+    def test_missing_symbol(self):
+        with open(os.path.join('test', 'testfiles_for_unittests',
+                               'simple_gcc.elf.arm'), 'rb') as f:
+            elf = ELFFile(f)
+
+            # Find the symbol table.
+            symtab = elf.get_section_by_name(b'.symtab')
+            self.assertIsNotNone(symtab)
+
+            # Test we get None when we look up a symbol that doesn't exist.
+            undef = symtab.get_symbol_by_name(b'non-existent symbol')
+            self.assertIsNone(undef)
+
+if __name__ == '__main__':
+    unittest.main()
-- 
2.30.2