write_partial_img pads correct to vtable

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

1 files changed, 16 insertions, 7 deletions

diff --git a/zep2newt.py b/zep2newt.py
index ff5c14c..3139065 100755
--- a/zep2newt.py
+++ b/zep2newt.py
@@ -78,7 +78,7 @@ def create_header(binary_file, sig_type, vtable_offs):
     # Get the correct size for the image
     image_size = os.path.getsize(binary_file)
     if DEBUG:
-        print("Binary size %d (0x%x)" % (image_size, image_size))
+        print("[*] Binary size %d (0x%x) of %s" % (image_size, image_size, binary_file))
 
     hdr = bytearray(struct.pack('I', IMAGE_MAGIC) +
                     struct.pack('H', tlv_size) +
@@ -100,22 +100,31 @@ def create_header(binary_file, sig_type, vtable_offs):
     return hdr
 
 
-def write_partial_img(binary_file, image_file, hdr):
+def write_partial_img(binary_file, image_file, hdr, vtable_offs):
     try:
         with open(binary_file, "rb") as f:
             image = f.read()
             f.close()
-        print("image len: %d" % len(image))
+        if DEBUG:
+            print("[*] Read %d bytes from %s" % (len(image), binary_file))
 
-        print("[*] Writing : {0}".format(image_file))
+    except IOError:
+        print("[*] Cannot open %s (!) ", image_file)
+
+    try:
         with open(image_file, "w+b") as f:
             f.write(hdr)
-            f.write('\xFF' * (OFFSET_VECTOR_TABLE - IMAGE_HEADER_SIZE))
+            # Calculate how much to pad before the actual image with the vector
+            # table starts.
+            f.write('\xFF' * (vtable_offs - IMAGE_HEADER_SIZE))
             f.write(image)
             f.close()
+        if DEBUG:
+            sz = os.path.getsize(image_file)
+            print("[*] Wrote %d (0x%x) bytes to %s" % (sz, sz, image_file))
 
     except IOError:
-        print("[*] Cannot open {0} (!) ".format(image_file))
+        print("[*] Cannot open %s (!) ", image_file)
 
 
 def calculate_hash(image_file):
@@ -210,7 +219,7 @@ def main(argv):
     hdr = create_header(args.binary_file, args.sig_type, vtable_offs)
 
     # Write the image itself
-    write_partial_img(args.binary_file, args.image_file, hdr)
+    write_partial_img(args.binary_file, args.image_file, hdr, vtable_offs)
 
     # Now we have a header and the binary itself and we should get the hash of
     # those concatenated.