Pull irom0.text from the appropriate location depending on the header…

… type. Works for e9 headers, still testing e4. For issue #3

esp_rom.py

@@ -10,28 +10,34 @@
 class EspRom(object):
     def __init__(self, rom_name, rom_bytes_stream, flash_layout):
         self.name = rom_name
+        self.sections = []
         self.contents = rom_bytes_stream.read()
         rom_bytes_stream.seek(0)
-        HeaderClass = EspRomHeader.get_header_type(self.contents[0])
-        self.header = HeaderClass(rom_bytes_stream.read(HeaderClass.ROM_HEADER_SIZE))
-
-        self.sections = []
-        for i in range(0, self.header.sect_count):
-            section = EspRomSection(rom_bytes_stream)
-            self.sections.append(section)
 
-        # read the irom0.text section from flash, non-OTA case.
-        irom_section = flash_layout['.irom0.text']
-        irom_address = 0x40200000
-        irom_size = irom_section_size * 1024
-
-        rom_bytes_stream.seek(irom_section.offset)
-        irom_text_contents = rom_bytes_stream.read(irom_size)
+        self.header = EspRomHeader.get_header(rom_bytes_stream)
+        if self.header.is_new():
+            # the new header format includes .irom0.text directly after,
+            # followed by an e9 header
+            irom_address = 0x40200000
+            irom_size = self.header.length
+            irom_text_contents = rom_bytes_stream.read(irom_size)
+            self.header = EspRomE9Header(rom_bytes_stream)
+        else:
+            # read the irom0.text section from flash, non-OTA case.
+            irom_section = flash_layout['.irom0.text']
+            irom_address = 0x40200000
+            irom_size = irom_section.size * 1024
+            limit = irom_section.offset + irom_size
+            irom_text_contents = self.contents[irom_section.offset:limit]
 
         # add .irom0.text section
         section = EspRomSection(StringIO(irom_text_contents), irom_address, irom_size)
         self.sections.append(section)
 
+        for i in range(0, self.header.sect_count):
+            section = EspRomSection(rom_bytes_stream)
+            self.sections.append(section)
+
     def __str__(self):
         rep = "EspRom("
         rep += "name: %s, " % (self.name)
@@ -44,15 +50,18 @@ def __str__(self):
 
 class EspRomHeader(object):
     @staticmethod
-    def get_header_type(header_byte):
-        if header_byte == '\xe9':
-            return EspRomE9Header
-        elif header_byte == '\xe4':
-            return EspRomE4Header
+    def get_header(rom_bytes_stream):
+        header_type = rom_bytes_stream.read(1)
+        rom_bytes_stream.seek(-1, 1) # relative position
+
+        if header_type == '\xe9':
+            return EspRomE9Header(rom_bytes_stream)
+        elif header_type == '\xe4':
+            return EspRomE4Header(rom_bytes_stream)
         else:
             raise RomParseException(
-                "EspRomHeader.get_header_type: unrecognized magic_number %s"
-                    % (header_byte))
+                "EspRomHeader.get_header: unrecognized magic_number %s"
+                    % (header_type))
 
     def __init__(self):
         pass
@@ -61,7 +70,7 @@ def __init__(self):
 class EspRomE9Header(EspRomHeader):
     ROM_HEADER_SIZE = 8
 
-    def __init__(self, rom_header_bytes):
+    def __init__(self, rom_bytes_stream):
         # typedef struct {
         #     uint8 magic;
         #     uint8 sect_count;
@@ -70,6 +79,8 @@ def __init__(self, rom_header_bytes):
         #     uint32 entry_addr;
         # } rom_header;
 
+        rom_header_bytes = rom_bytes_stream.read(EspRomE9Header.ROM_HEADER_SIZE)
+
         if len(rom_header_bytes) != EspRomE9Header.ROM_HEADER_SIZE:
             raise RomParseException(
                 "EspRomE9Header.init(): len(rom_header_bytes) is %d bytes != 8 bytes."
@@ -88,6 +99,9 @@ def __init__(self, rom_header_bytes):
 
         super(EspRomE9Header, self).__init__()
 
+    def is_new(self):
+        return False
+
     def __str__(self):
         rep = "EspRomE9Header("
         rep += "magic: 0x%02x, " % (self.magic)
@@ -102,7 +116,7 @@ def __str__(self):
 class EspRomE4Header(EspRomHeader):
     ROM_HEADER_SIZE = 16
 
-    def __init__(self, rom_header_bytes):
+    def __init__(self, rom_bytes_stream):
         # typedef struct {
         #     uint8 magic1;
         #     uint8 magic2;
@@ -112,6 +126,8 @@ def __init__(self, rom_header_bytes):
         #     uint32 length;
         # } rom_header;
 
+        rom_header_bytes = rom_bytes_stream.read(EspRomE4Header.ROM_HEADER_SIZE)
+
         if len(rom_header_bytes) != EspRomE4Header.ROM_HEADER_SIZE:
             raise RomParseException(
                 "EspRomE4Header.init(): len(rom_header_bytes) is %d bytes != 16 bytes."
@@ -136,6 +152,9 @@ def __init__(self, rom_header_bytes):
 
         super(EspRomE4Header, self).__init__()
 
+    def is_new(self):
+        return True
+
     def __str__(self):
         rep = "EspRomE4Header("
         rep += "magic1: 0x%02x, " % (self.magic1)
@@ -171,7 +190,7 @@ def __init__(self, rom_bytes_stream, address=None, length=None):
         else:
             # support specified length and address for non-OTA case and new header
             self.address = address
-            self.length = address
+            self.length = length
 
         self.contents = rom_bytes_stream.read(self.length)
 
@@ -188,5 +207,5 @@ def __str__(self):
         return rep
 
 
-class RomParseException(object):
+class RomParseException(Exception):
     pass