From 75b5fde294e320f239f9df85728ee4e6f2935098 Mon Sep 17 00:00:00 2001
From: eliben <devnull@localhost>
Date: Thu, 8 Sep 2011 16:24:49 +0300
Subject: [PATCH] Some code reorg: * created common/ dir and moved exceptions
 there (along with utils) * moved struct_parse from elffile to utils * moved
 string table functionality to its own section class

---
 elftools/common/__init__.py         |  0
 elftools/{ => common}/exceptions.py |  2 +-
 elftools/common/utils.py            | 25 +++++++++++
 elftools/elf/elffile.py             | 68 ++++++++++++-----------------
 elftools/elf/sections.py            | 17 ++++++++
 z.py                                |  3 ++
 6 files changed, 74 insertions(+), 41 deletions(-)
 create mode 100644 elftools/common/__init__.py
 rename elftools/{ => common}/exceptions.py (88%)
 create mode 100644 elftools/common/utils.py

diff --git a/elftools/common/__init__.py b/elftools/common/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/elftools/exceptions.py b/elftools/common/exceptions.py
similarity index 88%
rename from elftools/exceptions.py
rename to elftools/common/exceptions.py
index b61b463..05fa602 100644
--- a/elftools/exceptions.py
+++ b/elftools/common/exceptions.py
@@ -1,5 +1,5 @@
 #-------------------------------------------------------------------------------
-# elftools: exceptions.py
+# elftools: common/exceptions.py
 #
 # Exception classes for elftools
 #
diff --git a/elftools/common/utils.py b/elftools/common/utils.py
new file mode 100644
index 0000000..716bebd
--- /dev/null
+++ b/elftools/common/utils.py
@@ -0,0 +1,25 @@
+#-------------------------------------------------------------------------------
+# elftools: common/utils.py
+#
+# Miscellaneous utilities for elftools
+#
+# Eli Bendersky (eliben@gmail.com)
+# This code is in the public domain
+#-------------------------------------------------------------------------------
+from .exceptions import ELFParseError
+
+
+def struct_parse(struct, stream, stream_pos=None):
+    """ Convenience function for using the given struct to parse a stream (at
+        its current location).
+        If stream_pos is provided, the stream is seeked to this position before
+        the parsing is done.
+        Wraps the error thrown by construct with our own error.
+    """
+    try:
+        if stream_pos is not None:
+            stream.seek(stream_pos)
+        return struct.parse_stream(stream)
+    except ConstructError as e:
+        raise ELFParseError(e.message)
+    
diff --git a/elftools/elf/elffile.py b/elftools/elf/elffile.py
index 01be519..7d134b9 100644
--- a/elftools/elf/elffile.py
+++ b/elftools/elf/elffile.py
@@ -6,13 +6,11 @@
 # Eli Bendersky (eliben@gmail.com)
 # This code is in the public domain
 #-------------------------------------------------------------------------------
-
-from cStringIO import StringIO
-
-from ..exceptions import ELFError, ELFParseError
-from ..construct import ConstructError, CString
+from ..common.exceptions import ELFError
+from ..common.utils import struct_parse
+from ..construct import ConstructError
 from .structs import ELFStructs
-from .sections import Section
+from .sections import Section, StringTableSection
 from .segments import Segment
 
 
@@ -35,7 +33,8 @@ class ELFFile(object):
             little_endian=self.little_endian,
             elfclass=self.elfclass)
         self.header = self._parse_elf_header()
-        self._stringtable = self._get_stringtable()
+        
+        self._file_stringtable_section = self._get_file_stringtable()
     
     def num_sections(self):
         """ Number of sections in the file
@@ -119,51 +118,40 @@ class ELFFile(object):
     def _get_section_header(self, n):
         """ Find the header of section #n, parse it and return the struct 
         """
-        self.stream.seek(self._section_offset(n))
-        return self._struct_parse(self.structs.Elf_Shdr)
+        return struct_parse(
+            self.structs.Elf_Shdr,
+            self.stream,
+            stream_pos=self._section_offset(n))
     
+    def _get_section_name(self, section_header):
+        """ Given a section header, find this section's name in the file's
+            string table
+        """
+        name_offset = section_header['sh_name']
+        return self._file_stringtable_section.get_string(name_offset)
+
     def _get_segment_header(self, n):
         """ Find the header of segment #n, parse it and return the struct
         """
-        self.stream.seek(self._segment_offset(n))
-        return self._struct_parse(self.structs.Elf_Phdr)
+        return struct_parse(
+            self.structs.Elf_Phdr,
+            self.stream,
+            stream_pos=self._segment_offset(n))
     
-    def _get_section_name(self, section_header):
-        """ Given a section header, find this section's name in the file's
-            string table, and return it as a normal Python string.
-        """
-        offset = section_header['sh_name']
-        self._stringtable.seek(offset)
-        return CString('').parse_stream(self._stringtable)
-    
-    def _get_stringtable(self):
-        """ Find the file's string table section, read it and return the string
-            table as a StringIO object pointing to the section's contents.
+    def _get_file_stringtable(self):
+        """ Find the file's string table section
         """
-        # Find the section header for the stringtable header, and read the 
-        # section's contents from it
-        #
         stringtable_section_num = self['e_shstrndx']
-        stringtable_header = self._get_section_header(stringtable_section_num)
-        self.stream.seek(stringtable_header['sh_offset'])
-        return StringIO(self.stream.read(stringtable_header['sh_size']))
+        return StringTableSection(
+                header=self._get_section_header(stringtable_section_num),
+                name='',
+                stream=self.stream)
     
     def _parse_elf_header(self):
         """ Parses the ELF file header and assigns the result to attributes
             of this object.
         """
-        self.stream.seek(0)
-        return self._struct_parse(self.structs.Elf_Ehdr)
-    
-    def _struct_parse(self, struct):
-        """ Convenience method for parsing at the current stream location with
-            the given struct. Also wraps the error thrown by construct with our
-            own error.
-        """
-        try:
-            return struct.parse_stream(self.stream)
-        except ConstructError as e:
-            raise ELFParseError(e.message)
+        return struct_parse(self.structs.Elf_Ehdr, self.stream, stream_pos=0)
     
     def _assert(self, cond, msg=''):
         """ Assert that cond is True, otherwise raise ELFError(msg)
diff --git a/elftools/elf/sections.py b/elftools/elf/sections.py
index 4fb4249..502b69f 100644
--- a/elftools/elf/sections.py
+++ b/elftools/elf/sections.py
@@ -6,6 +6,8 @@
 # Eli Bendersky (eliben@gmail.com)
 # This code is in the public domain
 #-------------------------------------------------------------------------------
+from ..construct import CString
+
 
 class Section(object):
     def __init__(self, header, name, stream):
@@ -24,3 +26,18 @@ class Section(object):
         """
         return self.header[name]
 
+
+class StringTableSection(Section):
+    def __init__(self, header, name, stream):
+        super(StringTableSection, self).__init__(header, name, stream)
+        
+    def get_string(self, offset):
+        """ Get the string stored at the given offset in this string table.
+        """
+        table_offset = self['sh_offset']
+        self.stream.seek(table_offset + offset)
+        return CString('').parse_stream(self.stream)
+
+
+    
+    
diff --git a/z.py b/z.py
index 9f208d3..adcfd6c 100644
--- a/z.py
+++ b/z.py
@@ -9,6 +9,9 @@ stream = open('binfiles/z.elf', 'rb')
 
 efile = ELFFile(stream)
 
+print '===> %s sections!' % efile.num_sections() 
+print '===> %s segments!' % efile.num_segments()
+
 for sec in efile.iter_sections():
     print sec.name
 
-- 
2.30.2