import struct
import collections

class ElfFile:
   def __init__(self, f):
      self.f = f
      self.parse()
   def parse(self):
      self.parseHeader()
   @property
   def ElfClass(self):
      return classes[self.ei_class]
   def parseHeader(self):
      # Parse identifier
      e_ident = self.f.read(16)
      ELFMAGIC = bytes([0x7f]) + b'ELF'
      assert e_ident[0:4] == ELFMAGIC
      self.ei_class = e_ident[4]
      # Change the class of this object in some strange way:
      self.__class__ = classes[self.ei_class] 
      print('ELF class: ', self.ElfClass)
      assert self.ei_class in classes.keys()
      self.ei_data = e_ident[5]
      print('ei_data:', self.ei_data)
      ei_version = e_ident[6]
      assert ei_version == 1

      # rest of the ELF header:
      e_type = self.readHalf()
      e_machine = self.readHalf()
      e_version = self.readWord()
      e_entry = self.readAddr()
      self.e_phoff = self.readOff()
      self.e_shoff = self.readOff()
      self.e_flags = self.readWord()
      e_ehsize = self.readHalf()
      self.e_phentsize = self.readHalf()
      self.e_phnum = self.readHalf()
      self.e_shentsize = self.readHalf()
      self.e_shnum = self.readHalf()
      self.e_shstrndx = self.readHalf()
      
      assert e_ehsize == self.HEADERSIZE
      print('program header size', self.e_phentsize, self.e_phnum)
      print('section header size', self.e_shentsize, self.e_shnum)
      print('string index:', self.e_shstrndx)
   def readFmt(self, fmt):
      size = struct.calcsize(fmt)
      return struct.unpack(fmt, self.f.read(size))
   def readHalf(self):
      return self.readFmt('<H')[0]
   def readWord(self):
      return self.readFmt('<I')[0]
   def parsePrograms(self):
      # Skip to start of program headers:
      self.f.seek(self.e_phoff)
      pheaders = []
      for i in range(self.e_phnum):
         pheaders.append(self.parseProgramHeader())
      return pheaders

ProgramHeader = collections.namedtuple('ProgramHeader', ['p_type', 
   'p_flags', 'p_offset', 'p_vaddr', 'p_paddr', 'p_filesz', 'p_memsz'])
      
class Elf32File(ElfFile):
   HEADERSIZE=52
   def readAddr(self):
      return self.readFmt('<I')[0]
   readOff = readAddr
   def parseProgramHeader(self):
      p_type = self.readWord()
      p_offset = self.readOff()
      p_vaddr = self.readAddr()
      p_paddr = self.readAddr()
      p_filesz = self.readWord()
      p_memsz = self.readWord()
      p_flags = self.readWord()
      p_align = self.readWord()
      return ProgramHeader(p_type, p_flags, p_offset, p_vaddr, p_paddr, p_filesz, p_memsz)
      
class Elf64File(ElfFile):
   HEADERSIZE=64
   def readAddr(self):
      return self.readFmt('<Q')[0]
   readOff = readAddr
   readXword = readAddr
   def parseProgramHeader(self):
      p_type = self.readWord()
      p_flags = self.readWord()
      p_offset = self.readOff()
      p_vaddr = self.readAddr()
      p_paddr = self.readAddr()
      p_filesz = self.readXword()
      p_memsz = self.readXword()
      p_align = self.readXword()
      return ProgramHeader(p_type, p_flags, p_offset, p_vaddr, p_paddr)

classes = {1:Elf32File, 2:Elf64File}