Change to defined constants

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

2 files changed, 68 insertions, 7 deletions

diff --git a/newtimg.py b/newtimg.py
new file mode 100644
index 0000000..8124f0b
--- /dev/null
+++ b/newtimg.py
@@ -0,0 +1,60 @@
+IMAGE_MAGIC =			0x96f3b83
+IMAGE_MAGIC_NONE =		0xffffffff
+
+#
+# Image header flags.
+#
+IMAGE_F_PIC =			0x00000001
+IMAGE_F_SHA256 =		0x00000002  # Image contains hash TLV
+IMAGE_F_PKCS15_RSA2048_SHA256 = 0x00000004  # PKCS15 w/RSA and SHA
+IMAGE_F_ECDSA224_SHA256 =	0x00000008  # ECD
+IMAGE_HEADER_SIZE =		32
+
+# Digest sizes
+SHA256_DIGEST_SIZE =		32
+
+# 
+# Image trailer TLV types.
+# 
+IMAGE_TLV_SHA256 =		1   # SHA256 of image hdr and body
+IMAGE_TLV_RSA2048 =		2   # RSA2048 of hash output
+IMAGE_TLV_ECDSA224 =		3   # ECDSA of hash output
+
+
+OFFSET_IH_MAGIC = 0
+OFFSET_IH_TLV_SIZE = 4
+OFFSET_IH_KEY_ID = 6
+OFFSET_IH_PAD1 = 7
+OFFSET_IH_HDR_SIZE = 8
+OFFSET_IH_PAD2 = 10
+OFFSET_IH_IMG_SIZE = 12
+OFFSET_IH_FLAGS = 16
+OFFSET_IH_VER_MAJOR = 20
+OFFSET_IH_VER_MINOR = 21
+OFFSET_IH_REVISION = 22
+OFFSET_IH_BUILD_NUM = 24
+OFFSET_IH_PAD3 = 28
+#struct image_version { uint8_t iv_major; uint8_t iv_minor; uint16_t iv_revision;
+#    uint32_t iv_build_num;
+#};
+#
+#/** Image header.  All fields are in little endian byte order. */
+#struct image_header {
+#    uint32_t ih_magic;
+#    uint16_t ih_tlv_size; /* Trailing TLVs */
+#    uint8_t  ih_key_id;
+#    uint8_t  _pad1;
+#    uint16_t ih_hdr_size;
+#    uint16_t _pad2;
+#    uint32_t ih_img_size; /* Does not include header. */
+#    uint32_t ih_flags;
+#    struct image_version ih_ver;
+#    uint32_t _pad3;
+#};
+#
+#/** Image trailer TLV format. All fields in little endian. */
+#struct image_tlv {
+#    uint8_t  it_type;
+#    uint8_t  _pad;
+#    uint16_t it_len;
+#};
diff --git a/patch_mynewt_img.py b/patch_mynewt_img.py
index a6b48dc..f77b4e7 100755
--- a/patch_mynewt_img.py
+++ b/patch_mynewt_img.py
@@ -4,6 +4,7 @@ import mmap
 import os
 import struct
 import sys
+from newtimg import *
 from ctypes import *
 
 def patch_header(img):
@@ -18,12 +19,12 @@ def patch_header(img):
             print("[*] Loading : {0}".format(img))
 
             # Patch the header size (Hardcoded to 0x80)
-            f.seek(8)
+            f.seek(OFFSET_IH_HDR_SIZE)
             f.write('\x80\x00')
 
             # Patch the image size, start by convert int to hex string.
             image_size_hex = struct.pack('i', image_size)
-            f.seek(12)
+            f.seek(OFFSET_IH_IMG_SIZE)
             f.write(image_size_hex)
 
             f.close()
@@ -38,7 +39,7 @@ def extend_image(img, img_size):
         # 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)
+            hdr = f.read(IMAGE_HEADER_SIZE)
             image = f.read(img_size)
             f.close()
 
@@ -48,7 +49,7 @@ def extend_image(img, img_size):
         # 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(('\xFF' * (0x80 - IMAGE_HEADER_SIZE)))
             f.write(image)
             f.close()
     except IOError:
@@ -68,13 +69,13 @@ def append_hash(img, digest):
 
         # 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', 1) # 1: SHA256, 2: RSA2048, 3: ECDSA224
+        tlv_type = struct.pack('b', IMAGE_TLV_SHA256)
 
         # Next 1 byte padding
         tlv_pad = '\x00'
 
         # Finally the size of the TLV, for SHA256 that is 32 bytes
-        tlv_len = struct.pack('h', 0x20)
+        tlv_len = struct.pack('h', SHA256_DIGEST_SIZE)
         
         f.write(tlv_type)
         f.write(tlv_pad)
@@ -100,4 +101,4 @@ if __name__ == "__main__":
             sys.exit(0)
     else:
         print("[*] Not enough arguments (!) ")
-        print("[*] Usage : parser.py bin")
+        print("[*] Usage : patch_mynewt_img.py")