Improve readability and stability

Clean up debug prints and add exception handling where things
potentially could go wrong, so that error messages are a bit more clear
to read.

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

1 files changed, 121 insertions, 70 deletions

diff --git a/zep2newt.py b/zep2newt.py
index 3139065..087f2a4 100755
--- a/zep2newt.py
+++ b/zep2newt.py
@@ -8,8 +8,9 @@ from argparse import ArgumentParser
 from newtimg import *
 from ctypes import *
 
-DEBUG = 1
+DEBUG = False
 
+################################################################################
 def get_args():
     parser = ArgumentParser(description='Script to create images on a format \
                             that Mynewts bootloader expects')
@@ -52,7 +53,7 @@ def get_args():
 
     return parser.parse_args()
 
-
+################################################################################
 def create_header(binary_file, sig_type, vtable_offs):
     """
     Create a header on a format that Mynewt's bootloader expects. Based on
@@ -75,10 +76,15 @@ def create_header(binary_file, sig_type, vtable_offs):
         tlv_size = tlv_size + 4 + ECDSA_SIZE
         flags = IMAGE_F_ECDSA224_SHA256 | IMAGE_F_SHA256
 
-    # Get the correct size for the image
-    image_size = os.path.getsize(binary_file)
-    if DEBUG:
-        print("[*] Binary size %d (0x%x) of %s" % (image_size, image_size, binary_file))
+    image_size = 0
+    try:
+        # Get the correct size for the image
+        image_size = os.path.getsize(binary_file)
+        if DEBUG:
+            print("[*] Binary size %d (0x%x) of %s" % (image_size, image_size, binary_file))
+    except (OSError, IOError):
+        print("[ERROR]: Cannot open %s" % binary_file)
+        sys.exit(1)
 
     hdr = bytearray(struct.pack('I', IMAGE_MAGIC) +
                     struct.pack('H', tlv_size) +
@@ -94,12 +100,16 @@ def create_header(binary_file, sig_type, vtable_offs):
                     struct.pack('I', 0) + # Build number
                     struct.pack('I', 0)) # PAD3
     if DEBUG:
-        with open(binary_file + ".hdr", "w+b") as f:
-            f.write(hdr)
-            f.close()
+        try:
+            with open(binary_file + ".hdr", "w+b") as f:
+                f.write(hdr)
+                f.close()
+        except (OSError, IOError):
+            print("[ERROR]: Cannot write to %s (!) " % (binary_file + ".hdr"))
+            sys.exit(1)
     return hdr
 
-
+################################################################################
 def write_partial_img(binary_file, image_file, hdr, vtable_offs):
     try:
         with open(binary_file, "rb") as f:
@@ -108,8 +118,9 @@ def write_partial_img(binary_file, image_file, hdr, vtable_offs):
         if DEBUG:
             print("[*] Read %d bytes from %s" % (len(image), binary_file))
 
-    except IOError:
-        print("[*] Cannot open %s (!) ", image_file)
+    except (OSError, IOError):
+        print("[ERROR]: Cannot open %s" % (binary_file))
+        sys.exit(1)
 
     try:
         with open(image_file, "w+b") as f:
@@ -123,36 +134,54 @@ def write_partial_img(binary_file, image_file, hdr, vtable_offs):
             sz = os.path.getsize(image_file)
             print("[*] Wrote %d (0x%x) bytes to %s" % (sz, sz, image_file))
 
-    except IOError:
-        print("[*] Cannot open %s (!) ", image_file)
+    except (OSError, IOError):
+        print("[ERROR]: Cannot write to %s" % (image_file))
+        sys.exit(1)
 
 
+################################################################################
 def calculate_hash(image_file):
     sha256 = hashlib.sha256()
-    with open(image_file, "rb") as f:
+    try:
+        with open(image_file, "rb") as f:
             sha256.update(f.read())
-    return sha256.hexdigest()
+            f.close()
+
+    except IOError:
+        print("[ERROR]: Cannot open %s" % (image_file))
+        sys.exit(1)
 
+    digest = sha256.hexdigest()
+    if DEBUG:
+        print("[*] Hash of intermediate image: %s" % digest)
+    return digest
 
+################################################################################
 def append_hash(image_file, digest):
-    with open(image_file, "ab") as f:
+    try:
+        with open(image_file, "ab") as f:
+
+            # Start by settings the TLV type
+            # https://github.com/apache/incubator-mynewt-newt/blob/master/newt/image/image.go#L109-L116
+            tlv_type = struct.pack('b', IMAGE_TLV_SHA256)
+
+            # Next 1 byte padding
+            tlv_pad = '\x00'
 
-        # Start by settings the TLV type
-        # https://github.com/apache/incubator-mynewt-newt/blob/master/newt/image/image.go#L109-L116
-        tlv_type = struct.pack('b', IMAGE_TLV_SHA256)
+            # Finally the size of the TLV, for SHA256 that is 32 bytes
+            tlv_len = struct.pack('h', SHA256_DIGEST_SIZE)
 
-        # Next 1 byte padding
-        tlv_pad = '\x00'
+            f.write(tlv_type)
+            f.write(tlv_pad)
+            f.write(tlv_len)
+            f.write(digest.decode('hex'))
+            f.close()
 
-        # Finally the size of the TLV, for SHA256 that is 32 bytes
-        tlv_len = struct.pack('h', SHA256_DIGEST_SIZE)
-        
-        f.write(tlv_type)
-        f.write(tlv_pad)
-        f.write(tlv_len)
-        f.write(digest.decode('hex'))
-        f.close()
+    except IOError:
+        print("[ERROR]: Cannot open/append to %s" % (image_file))
+        sys.exit(1)
 
+################################################################################
 def create_jlink_script(image_file, offset, erase):
     """
     Creates a jlink script to flash the created binary.
@@ -160,55 +189,77 @@ def create_jlink_script(image_file, offset, erase):
     @erase: whether the script first shall erase or not when flashing.
     @offset: where in flash to store the image.
     """
-    with open("flash_zephyr.jlink", "w+") as f:
-        f.write("device nrf52\n")
-        f.write("power on\n")
-        f.write("sleep 10\n")
-        f.write("si 1\n")
-        f.write("speed auto\n")
-        if erase:
-            f.write("erase\n")
-        f.write("loadfile %s %s\n" % (image_file, hex(int(offset, 16))))
-        f.write("q\n")
-        f.close()
-        print("\nTo flash the Image for nrf52 run:")
-        print(" $ JLinkExe -CommanderScript flash_zephyr.jlink")
+    jlink_file = "flash_zephyr.jlink"
+    try:
+        with open(jlink_file, "w+") as f:
+            f.write("device nrf52\n")
+            f.write("power on\n")
+            f.write("sleep 10\n")
+            f.write("si 1\n")
+            f.write("speed auto\n")
+            if erase:
+                f.write("erase\n")
+            f.write("loadfile %s %s\n" % (image_file, hex(int(offset, 16))))
+            f.write("q\n")
+            f.close()
+        if DEBUG:
+            print("\n[*] To flash the Image for nrf52 run:")
+            print("     JLinkExe -CommanderScript %s" % jlink_file)
 
+    except IOError:
+        print("[ERROR]: Cannot create to %s" % (jlink_file))
+        sys.exit(1)
+
+################################################################################
 def create_jlink_bit_script(bit_file, bitoffset="0x7bff8"):
-    with open("flash_bit.jlink", "w+") as f:
-        f.write("device nrf52\n")
-        f.write("power on\n")
-        f.write("sleep 10\n")
-        f.write("si 1\n")
-        f.write("speed auto\n")
-        f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
-        f.write("q\n")
-        f.close()
-        print("\nTo flash Boot Image Trailer for nrf52 run:")
-        print(" $ JLinkExe -CommanderScript flash_bit.jlink")
+    jlink_file = "flash_bit.jlink"
+    try:
+        with open(jlink_file, "w+") as f:
+            f.write("device nrf52\n")
+            f.write("power on\n")
+            f.write("sleep 10\n")
+            f.write("si 1\n")
+            f.write("speed auto\n")
+            f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
+            f.write("q\n")
+            f.close()
+        if DEBUG:
+            print("\n[*] To flash Boot Image Trailer for nrf52 run:")
+            print("     JLinkExe -CommanderScript %s" % jlink_file)
+
+    except IOError:
+        print("[ERROR]: Cannot create to %s" % (jlink_file))
+        sys.exit(1)
 
+################################################################################
 def create_jlink_clear_bit_script(bit_file, bitoffset="0x7bff8"):
-    with open("flash_clear_bit.jlink", "w+") as f:
-        f.write("device nrf52\n")
-        f.write("power on\n")
-        f.write("sleep 10\n")
-        f.write("si 1\n")
-        f.write("speed auto\n")
-        f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
-        f.write("q\n")
-        f.close()
-        print("\nTo flash Boot Image Trailer for nrf52 run:")
-        print(" $ JLinkExe -CommanderScript flash_clear_bit.jlink")
+    jlink_file = "flash_clear_bit.jlink"
+    try:
+        with open(jlink_file, "w+") as f:
+            f.write("device nrf52\n")
+            f.write("power on\n")
+            f.write("sleep 10\n")
+            f.write("si 1\n")
+            f.write("speed auto\n")
+            f.write("loadfile %s %s\n" % (bit_file, hex(int(bitoffset, 16))))
+            f.write("q\n")
+            f.close()
+        if DEBUG:
+            print("\n[*] To clear Boot Image Trailer for nrf52 run:")
+            print("     JLinkExe -CommanderScript %s" % jlink_file)
 
+    except IOError:
+        print("[ERROR]: Cannot create to %s" % (jlink_file))
+        sys.exit(1)
 
+################################################################################
 def main(argv):
     args = get_args()
     erase = False
 
     if args.verbose:
-        print("Arguments: ")
-        print(args)
-        print("\n")
+        for a in vars(args):
+            print("Arg -> %s: %s" % (a, getattr(args, a)))
         global DEBUG
         DEBUG = True
 
@@ -224,9 +275,10 @@ def main(argv):
     # Now we have a header and the binary itself and we should get the hash of
     # those concatenated.
     digest = calculate_hash(args.image_file)
-    print("Hash of intermediate image: %s" % digest)
     append_hash(args.image_file, digest)
+    print("[*] Successfully created: %s" % args.image_file)
 
+    # Misc function related to flashing
     create_jlink_script(args.image_file, args.flash_offs_addr, erase)
     curdir = os.path.dirname(sys.argv[0])
     create_jlink_bit_script(curdir + "/boot_image_trailer.bin")
@@ -236,6 +288,5 @@ def main(argv):
     if args.f:
         os.system("JLinkExe -CommanderScript flash_zephyr.jlink")
 
-
 if __name__ == "__main__":
     main(sys.argv)