Put correct TLV size in the header

Depending on verification type, SHA256, RSA or EC, calculate and store
the correct tlv_size.

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

2 files changed, 17 insertions, 5 deletions

diff --git a/newtimg.py b/newtimg.py
index ca55307..1c5c7a9 100644
--- a/newtimg.py
+++ b/newtimg.py
@@ -13,6 +13,10 @@ IMAGE_HEADER_SIZE =		32
 # Digest sizes
 SHA256_DIGEST_SIZE =		32
 
+# Signature sizes
+RSA_SIZE =                      256
+ECDSA_SIZE =                    68
+
 # 
 # Image trailer TLV types.
 # 
diff --git a/zep2newt.py b/zep2newt.py
index dc75068..a0305e6 100755
--- a/zep2newt.py
+++ b/zep2newt.py
@@ -24,8 +24,9 @@ def get_args():
                         default='zephyr.img.bin', \
                         help='Name of *.img file (output)')
 
-    parser.add_argument('--sig', required=False, dest='sig_type', default=None, \
-                        help='Type of signature <RSA|EC> (not implemented)')
+    parser.add_argument('--sig', required=False, dest='sig_type', \
+                        default='SHA256', \
+                        help='Type of signature <SHA256|RSA|EC>')
 
     parser.add_argument('--off', required=False, dest='flash_offs_addr', \
                         default='0x8000', \
@@ -51,13 +52,20 @@ def get_args():
     return parser.parse_args()
 
 
-def create_header(binary_file):
+def create_header(binary_file, sig_type):
     # Start by getting the correct size for the image
     image_size = os.path.getsize(binary_file) # TODO, handle zephyr.bin
     print("Image size (bin): %d" % image_size)
 
+    # The SHA256 hash is always used
+    tlv_size = SHA256_DIGEST_SIZE + 4
+    if sig_type == "RSA":
+        tlv_size = tlv_size + 4 + RSA_SIZE
+    elif sig_type == "EC":
+        tlv_size = tlv_size + 4 + ECDSA_SIZE
+
     hdr = bytearray(struct.pack('I', IMAGE_MAGIC) +
-                    struct.pack('H', SHA256_DIGEST_SIZE + 4) + # TLV size
+                    struct.pack('H', tlv_size) + # TLV size
                     struct.pack('B', 0) + # Key ID
                     struct.pack('B', 0) + # PAD 1
                     struct.pack('H', OFFSET_VECTOR_TABLE) + # New HDR SIZE
@@ -183,7 +191,7 @@ def main(argv):
         global DEBUG
         DEBUG = True
 
-    hdr = create_header(args.binary_file)
+    hdr = create_header(args.binary_file, args.sig_type)
     write_partial_img(args.binary_file, args.image_file, hdr)
 
     digest = calculate_hash(args.image_file)