Add support for cmd line parameters

From this commit we can create a binary working with mynewt's bootloader
without using any of the tools from mynewt.

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

1 files changed, 54 insertions, 19 deletions

diff --git a/patch_mynewt_img.py b/patch_mynewt_img.py
index a7bd521..634399d 100755
--- a/patch_mynewt_img.py
+++ b/patch_mynewt_img.py
@@ -1,4 +1,5 @@
 #!/usr/bin/python2
+import getopt
 import hashlib
 import mmap
 import os
@@ -9,9 +10,9 @@ from ctypes import *
 
 DEBUG = 1
 
-def create_header(img):
+def create_header(binary_file):
     # Start by getting the correct size for the image
-    image_size = os.path.getsize(img + ".elf.bin") # TODO, handle zephyr.bin
+    image_size = os.path.getsize(binary_file) # TODO, handle zephyr.bin
     print("Image size (bin): %d" % image_size)
 
     hdr = bytearray(struct.pack('I', IMAGE_MAGIC) +
@@ -28,38 +29,39 @@ def create_header(img):
                     struct.pack('I', 0) + # Build number
                     struct.pack('I', 0)) # PAD3
     if DEBUG:
-        with open(img + ".hdr", "w+b") as f:
+        with open(binary_file + ".hdr", "w+b") as f:
             f.write(hdr)
             f.close()
 
     return hdr
 
-def write_partial_img(image_name, hdr):
+def write_partial_img(binary_file, image_file, hdr):
     try:
-        with open(image_name + ".elf.bin", "rb") as f:
+        with open(binary_file, "rb") as f:
             image = f.read()
             f.close()
         print("image len: %d" % len(image))
 
-        print("[*] Writing : {0}".format(image_name))
-        with open(image_name + ".img", "w+b") as f:
+        print("[*] Writing : {0}".format(image_file))
+        with open(image_file, "w+b") as f:
             f.write(hdr)
             f.write('\xFF' * (OFFSET_VECTOR_TABLE - IMAGE_HEADER_SIZE))
             f.write(image)
             f.close()
 
     except IOError:
-        print("[*] Cannot open {0} (!) ".format(image_name))
+        print("[*] Cannot open {0} (!) ".format(image_file))
 
-def calculate_hash(image_name):
+
+def calculate_hash(image_file):
     sha256 = hashlib.sha256()
-    with open(image_name + ".img", "rb") as f:
+    with open(image_file, "rb") as f:
             sha256.update(f.read())
     return sha256.hexdigest()
 
 
-def append_hash(img, digest):
-    with open(img + ".img", "ab") as f:
+def append_hash(image_file, digest):
+    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
@@ -77,19 +79,52 @@ def append_hash(img, digest):
         f.write(digest.decode('hex'))
         f.close()
 
-def main(img):
-    # backup(img) # Backup original image
-    hdr = create_header(img)
-    write_partial_img(img, hdr)
-    digest = calculate_hash(img)
+
+def usage():
+    print(sys.argv[0] + " -b <zephyr.bin> -i <zephyr.img> -o <offset>")
+
+
+def main(argv):
+    binary_file = "zephyr.bin"
+    image_file = "zephyr.img"
+    offset = 0x80
+    try:
+        opts, args = getopt.getopt(argv, 'b:i:o:h', ['bin=', 'img=', 'offset=', "help"])
+    except getopt.GetoptError:
+        usage()
+        sys.exit(2)
+
+    for opt, arg in opts:
+        if opt in ('-h', '--help'):
+            usage()
+            sys.exit(2)
+        elif opt in ('-b', '--bin'):
+            binary_file = arg
+            print("bin called")
+        elif opt in ('-i', '--img'):
+            image_file = arg
+        elif opt in ('-o', '--offset'):
+            offset = arg
+        else:
+            usage()
+            sys.exit(2)
+
+    if DEBUG:
+        print("Input (bin): %s" % binary_file)
+        print("Output (img): %s" % image_file)
+        print("Offset to reset vector: %s" % offset)
+
+    hdr = create_header(binary_file)
+    write_partial_img(binary_file, image_file, hdr)
+    digest = calculate_hash(image_file)
     print("Hash of intermediate image: %s" % digest)
 
-    append_hash(img, digest)
+    append_hash(image_file, digest)
 
 if __name__ == "__main__":
     if sys.argv[1]:
         try:
-            main(sys.argv[1])
+            main(sys.argv[1:])
         except KeyboardInterrupt:
             sys.exit(0)
     else: