Fix issue with double SHA256 hashes at the endHEADmaster

Signed-off-by: Joakim Bech <joakim.bech@linaro.org>

1 files changed, 26 insertions, 10 deletions

diff --git a/patch_mynewt_img.py b/patch_mynewt_img.py
index 3707d31..a6b48dc 100755
--- a/patch_mynewt_img.py
+++ b/patch_mynewt_img.py
@@ -7,6 +7,9 @@ import sys
 from ctypes import *
 
 def patch_header(img):
+    """
+    Patch the image with the new header size and with the new image size.
+    """
     image_size = os.path.getsize(img + ".elf.bin")
     print("Image size (bin): %d" % image_size)
 
@@ -19,22 +22,34 @@ def patch_header(img):
             f.write('\x80\x00')
 
             # Patch the image size, start by convert int to hex string.
-            image_size = struct.pack('i', image_size)
+            image_size_hex = struct.pack('i', image_size)
             f.seek(12)
-            f.write(image_size)
+            f.write(image_size_hex)
 
-            # 
             f.close()
     except IOError:
         print("[*] Cannot open {0} (!) ".format(img))
 
+    return image_size
 
-def extend_image(img):
+
+def extend_image(img, img_size):
     try:
-        with open(img + ".img", "r+b") as f:
-            print("[*] Loading : {0}".format(img))
-            mm=mmap.mmap(f.fileno(),0)
-            f.write(mm[0:32] + ('\xFF' * 0x60) + mm[32:])
+        # Read the header and image locally
+        with open(img + ".img", "rb") as f:
+            print("[*] Loading %s (img_size: %d)" % (img, img_size))
+            hdr = f.read(0x20)
+            image = f.read(img_size)
+            f.close()
+
+        # Clear the file (why doesn't truncate work?)
+        open(img + ".img").close()
+
+        # Store the header, pad it and then store the image itself
+        with open(img + ".img", "w+b") as f:
+            f.write(hdr)
+            f.write(('\xFF' * 0x60))
+            f.write(image)
             f.close()
     except IOError:
         print("[*] Cannot open {0} (!) ".format(img))
@@ -68,8 +83,9 @@ def append_hash(img, digest):
         f.close()
 
 def main(img):
-    patch_header(img)
-    extend_image(img)
+    # backup(img) # Backup original image
+    size = patch_header(img)
+    extend_image(img, size)
 
     digest = calculate_hash(img)
     print("Hash of intermediate image: %s" % digest)