1 files changed, 48 insertions, 0 deletions
diff --git a/zep2newt.py b/zep2newt.py
index 087f2a4..70a35b2 100755
--- a/zep2newt.py
+++ b/zep2newt.py
@@ -7,6 +7,9 @@ import sys
 from argparse import ArgumentParser
 from newtimg import *
 from ctypes import *
+from Crypto.Signature import PKCS1_v1_5
+from Crypto.Hash import SHA256
+from Crypto.PublicKey import RSA
 
 DEBUG = False
 
@@ -182,6 +185,44 @@ def append_hash(image_file, digest):
         sys.exit(1)
 
 ################################################################################
+def append_rsa_signature(image_file, key_file, bin_digest):
+    signature = None
+    try:
+        with open(key_file, "rb") as f:
+            rsa_key = RSA.importKey(f.read())
+            f.close()
+            rsa = PKCS1_v1_5.new(rsa_key)
+            digest = SHA256.new()
+            digest.update(bin_digest)
+            signature = rsa.sign(digest)
+    except (OSError, IOError):
+        print("[ERROR]: Cannot open %s" % (key_file))
+        sys.exit(1)
+
+    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_RSA2048)
+
+            # Next 1 byte padding
+            tlv_pad = '\x00'
+
+            # Finally the size of the TLV, for SHA256 that is 32 bytes
+            tlv_len = struct.pack('h', RSA_SIZE)
+
+            f.write(tlv_type)
+            f.write(tlv_pad)
+            f.write(tlv_len)
+            f.write(signature)
+            f.close()
+
+    except (OSError, 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.
@@ -276,6 +317,13 @@ def main(argv):
     # those concatenated.
     digest = calculate_hash(args.image_file)
     append_hash(args.image_file, digest)
+
+    if args.sig_type == "RSA":
+        append_rsa_signature(args.image_file, args.key_file, digest)
+    elif args.sig_tye == "EC":
+        print("[ERROR]: ECDSA not implemented")
+        sys.exit(1)
+
     print("[*] Successfully created: %s" % args.image_file)
 
     # Misc function related to flashing