imgtool: Add support for encrypting image with raw AES key

The change adds --aes-raw-key option that allows to pass a key
via command line. The key is used to encrypt the image and there
is not key exchange TLV added to the image.
The options is provided for encrypting images for devices that store
AES key on them so they do not expect it to be passed with image,
in encrypted form.

Signed-off-by: Dominik Ermel <dominik.ermel@nordicsemi.no>

Author: de-nordic

scripts/imgtool/image.py

@@ -513,12 +513,13 @@ def ecies_hkdf(self, enckey, plainkey, hmac_sha_alg):
 
     def create(self, key, public_key_format, enckey, dependencies=None,
                sw_type=None, custom_tlvs=None, compression_tlvs=None,
-               compression_type=None, encrypt_keylen=128, clear=False,
+               compression_type=None, aes_raw_key=None, clear=False,
                fixed_sig=None, pub_key=None, vector_to_sign=None,
-               user_sha='auto', hmac_sha='auto', is_pure=False, keep_comp_size=False,
-               dont_encrypt=False):
+               user_sha='auto', hmac_sha='auto', is_pure=False, keep_comp_size=False):
         self.enckey = enckey
 
+        dont_encrypt = bool(not aes_raw_key)
+
         # key decides on sha, then pub_key; of both are none default is used
         check_key = key if key is not None else pub_key
         hash_algorithm, hash_tlv = key_and_user_sha_to_alg_and_tlv(check_key, user_sha, is_pure)
@@ -605,7 +606,7 @@ def create(self, key, public_key_format, enckey, dependencies=None,
         #
         # This adds the padding if image is not aligned to the 16 Bytes
         # in encrypted mode
-        if self.enckey is not None and dont_encrypt is False:
+        if aes_raw_key is not None and dont_encrypt is False:
             pad_len = len(self.payload) % 16
             if pad_len > 0:
                 pad = bytes(16 - pad_len)
@@ -620,10 +621,8 @@ def create(self, key, public_key_format, enckey, dependencies=None,
             if compression_type == "lzma2armthumb":
                 compression_flags |= IMAGE_F['COMPRESSED_ARM_THUMB']
         # This adds the header to the payload as well
-        if encrypt_keylen == 256:
-            self.add_header(enckey, protected_tlv_size, compression_flags, 256)
-        else:
-            self.add_header(enckey, protected_tlv_size, compression_flags)
+        aes_key_bits = 0 if aes_raw_key is None else len(aes_raw_key) * 8
+        self.add_header(protected_tlv_size, compression_flags, aes_key_bits)
 
         prot_tlv = TLV(self.endian, TLV_PROT_INFO_MAGIC)
 
@@ -742,12 +741,18 @@ def create(self, key, public_key_format, enckey, dependencies=None,
         if protected_tlv_off is not None:
             self.payload = self.payload[:protected_tlv_off]
 
-        if enckey is not None and dont_encrypt is False:
-            if encrypt_keylen == 256:
-                plainkey = os.urandom(32)
-            else:
-                plainkey = os.urandom(16)
+        # When passed AES key and clear flag, then do not encrypt, because the key
+        # is only passed to be stored in encryption key TLV, the paylad stays clear text.
+        if aes_raw_key and not clear:
+            nonce = bytes([0] * 16)
+            cipher = Cipher(algorithms.AES(aes_raw_key), modes.CTR(nonce),
+                            backend=default_backend())
+            encryptor = cipher.encryptor()
+            img = bytes(self.payload[self.header_size:])
+            self.payload[self.header_size:] = \
+                encryptor.update(img) + encryptor.finalize()
 
+        if enckey is not None and dont_encrypt is False:
             if not isinstance(enckey, rsa.RSAPublic):
                 if hmac_sha == 'auto' or hmac_sha == '256':
                     hmac_sha = '256'
@@ -762,35 +767,26 @@ def create(self, key, public_key_format, enckey, dependencies=None,
 
             if isinstance(enckey, rsa.RSAPublic):
                 cipherkey = enckey._get_public().encrypt(
-                    plainkey, padding.OAEP(
+                    aes_raw_key, padding.OAEP(
                         mgf=padding.MGF1(algorithm=hashes.SHA256()),
                         algorithm=hashes.SHA256(),
                         label=None))
                 self.enctlv_len = len(cipherkey)
                 tlv.add('ENCRSA2048', cipherkey)
             elif isinstance(enckey, ecdsa.ECDSA256P1Public):
-                cipherkey, mac, pubk = self.ecies_hkdf(enckey, plainkey, hmac_sha_alg)
+                cipherkey, mac, pubk = self.ecies_hkdf(enckey, aes_raw_key, hmac_sha_alg)
                 enctlv = pubk + mac + cipherkey
                 self.enctlv_len = len(enctlv)
                 tlv.add('ENCEC256', enctlv)
             elif isinstance(enckey, x25519.X25519Public):
-                cipherkey, mac, pubk = self.ecies_hkdf(enckey, plainkey, hmac_sha_alg)
+                cipherkey, mac, pubk = self.ecies_hkdf(enckey, aes_raw_key, hmac_sha_alg)
                 enctlv = pubk + mac + cipherkey
                 self.enctlv_len = len(enctlv)
                 if (hmac_sha == '256'):
                     tlv.add('ENCX25519', enctlv)
                 else:
                     tlv.add('ENCX25519_SHA512', enctlv)
 
-            if not clear:
-                nonce = bytes([0] * 16)
-                cipher = Cipher(algorithms.AES(plainkey), modes.CTR(nonce),
-                                backend=default_backend())
-                encryptor = cipher.encryptor()
-                img = bytes(self.payload[self.header_size:])
-                self.payload[self.header_size:] = \
-                    encryptor.update(img) + encryptor.finalize()
-
         self.payload += prot_tlv.get()
         self.payload += tlv.get()
 
@@ -805,11 +801,11 @@ def get_signature(self):
     def get_infile_data(self):
         return self.infile_data
 
-    def add_header(self, enckey, protected_tlv_size, compression_flags, aes_length=128):
+    def add_header(self, protected_tlv_size, compression_flags, aes_length=0):
         """Install the image header."""
 
         flags = 0
-        if enckey is not None:
+        if aes_length != 0:
             if aes_length == 128:
                 flags |= IMAGE_F['ENCRYPTED_AES128']
             else:

scripts/imgtool/main.py

@@ -23,6 +23,7 @@
 import re
 import struct
 import sys
+import os
 
 import click
 
@@ -450,13 +451,17 @@ def convert(self, value, param, ctx):
               help='Unique vendor identifier, format: (<raw_uuid>|<domain_name)>')
 @click.option('--cid', default=None, required=False,
               help='Unique image class identifier, format: (<raw_uuid>|<image_class_name>)')
-def sign(key, public_key_format, align, version, pad_sig, header_size,
+@click.option('--aes-raw-key', default=None, required=False,
+              help='String representing raw AES key, format: hex byte string of 32 or 64'
+                   'hexadecimal characters')
+@click.pass_context
+def sign(ctx, key, public_key_format, align, version, pad_sig, header_size,
          pad_header, slot_size, pad, confirm, test, max_sectors, overwrite_only,
          endian, encrypt_keylen, encrypt, compression, infile, outfile,
          dependencies, load_addr, hex_addr, erased_val, save_enctlv,
          security_counter, boot_record, custom_tlv, rom_fixed, max_align,
          clear, fix_sig, fix_sig_pubkey, sig_out, user_sha, hmac_sha, is_pure,
-         vector_to_sign, non_bootable, vid, cid):
+         vector_to_sign, non_bootable, vid, cid, aes_raw_key):
 
     if confirm or test:
         # Confirmed but non-padded images don't make much sense, because
@@ -472,17 +477,30 @@ def sign(key, public_key_format, align, version, pad_sig, header_size,
                       non_bootable=non_bootable, vid=vid, cid=cid)
     compression_tlvs = {}
     img.load(infile)
+
     key = load_key(key) if key else None
-    enckey = load_key(encrypt) if encrypt else None
-    if enckey and key and ((isinstance(key, keys.ECDSA256P1) and
-         not isinstance(enckey, keys.ECDSA256P1Public))
-       or (isinstance(key, keys.ECDSA384P1) and
-           not isinstance(enckey, keys.ECDSA384P1Public))
-            or (isinstance(key, keys.RSA) and
-                not isinstance(enckey, keys.RSAPublic))):
-        # FIXME
-        raise click.UsageError("Signing and encryption must use the same "
-                               "type of key")
+    enckey = None
+    if not aes_raw_key:
+        enckey = load_key(encrypt) if encrypt else None
+        if enckey and key:
+            if ((isinstance(key, keys.ECDSA256P1) and
+              not isinstance(enckey, keys.ECDSA256P1Public))
+            or (isinstance(key, keys.ECDSA384P1) and
+                not isinstance(enckey, keys.ECDSA384P1Public))
+                 or (isinstance(key, keys.RSA) and
+                     not isinstance(enckey, keys.RSAPublic))):
+                # FIXME
+                raise click.UsageError("Signing and encryption must use the same "
+                                        "type of key")
+    else:
+       if encrypt:
+           encrypt = None
+           print('Raw AES key overrides --key, there will be no encrypted key added to the image')
+       if clear:
+           clear = False
+           print('Raw AES key overrides --clear, image will be encrypted')
+       if ctx.get_parameter_source('encrypt_keylen') != click.core.ParameterSource.DEFAULT:
+           print('Raw AES key len overrides --encrypt-keylen')
 
     if pad_sig and hasattr(key, 'pad_sig'):
         key.pad_sig = True
@@ -527,11 +545,28 @@ def sign(key, public_key_format, align, version, pad_sig, header_size,
             'Pure signatures, currently, enforces preferred hash algorithm, '
             'and forbids sha selection by user.')
 
+    plainkey = None
+    if aes_raw_key:
+        # Converting the command line provided raw AES key to byte array;
+        # this aray will be truncated to desired len.
+        plainkey = bytes.fromhex(aes_raw_key)
+        plainkey_len = len(plainkey)
+        if plainkey_len not in (16, 32):
+            raise click.UsageError("Provided keylen, {int(plainkey_len)} in bytes, not supported")
+    elif enckey:
+        if encrypt_keylen == 256:
+            encrypt_keylen_bytes = 32
+        else:
+            encrypt_keylen_bytes = 16
+
+        # No AES plain key and there is request to encrypt, generate random AES key
+        plainkey = os.urandom(encrypt_keylen_bytes)
+
     if compression in ["lzma2", "lzma2armthumb"]:
         img.create(key, public_key_format, enckey, dependencies, boot_record,
-               custom_tlvs, compression_tlvs, None, int(encrypt_keylen), clear,
+               custom_tlvs, compression_tlvs, None, None, clear,
                baked_signature, pub_key, vector_to_sign, user_sha=user_sha,
-               hmac_sha=hmac_sha, is_pure=is_pure, keep_comp_size=False, dont_encrypt=True)
+               hmac_sha=hmac_sha, is_pure=is_pure, keep_comp_size=False)
         compressed_img = image.Image(version=decode_version(version),
                   header_size=header_size, pad_header=pad_header,
                   pad=pad, confirm=confirm, align=int(align),
@@ -575,13 +610,13 @@ def sign(key, public_key_format, align, version, pad_sig, header_size,
                 keep_comp_size = True
             compressed_img.create(key, public_key_format, enckey,
                dependencies, boot_record, custom_tlvs, compression_tlvs,
-               compression, int(encrypt_keylen), clear, baked_signature,
+               compression, plainkey, clear, baked_signature,
                pub_key, vector_to_sign, user_sha=user_sha, hmac_sha=hmac_sha,
                is_pure=is_pure, keep_comp_size=keep_comp_size)
             img = compressed_img
     else:
         img.create(key, public_key_format, enckey, dependencies, boot_record,
-               custom_tlvs, compression_tlvs, None, int(encrypt_keylen), clear,
+               custom_tlvs, compression_tlvs, None, plainkey, clear,
                baked_signature, pub_key, vector_to_sign, user_sha=user_sha,
                hmac_sha=hmac_sha, is_pure=is_pure)
     img.save(outfile, hex_addr)