imgtool: Add support for encrypting image with raw AES key

The change adds --aes-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,7 +513,7 @@ 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_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):
@@ -605,7 +605,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_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 +620,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_key is None else len(aes_key) * 8
+        self.add_header(protected_tlv_size, compression_flags, aes_key_bits)
 
         prot_tlv = TLV(self.endian, TLV_PROT_INFO_MAGIC)
 
@@ -742,12 +740,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 payload stays clear text.
+        if aes_key and not clear:
+            nonce = bytes([0] * 16)
+            cipher = Cipher(algorithms.AES(aes_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 +766,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_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_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_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 +800,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

@@ -20,6 +20,7 @@
 import base64
 import getpass
 import lzma
+import os
 import re
 import struct
 import sys
@@ -322,6 +323,14 @@ def create_lzma2_header(dictsize, pb, lc, lp):
     header.append( ( pb * 5 + lp) * 9 + lc)
     return header
 
+def match_sig_enc_key(skey, ekey):
+    ok = ((isinstance(skey, keys.ECDSA256P1) and isinstance(ekey, keys.ECDSA256P1Public)) or
+          (isinstance(skey, keys.ECDSA384P1) and isinstance(ekey, keys.ECDSA384P1Public)) or
+          (isinstance(skey, keys.RSA) and isinstance(ekey, keys.RSAPublic))
+    )
+
+    return ok
+
 class BasedIntParamType(click.ParamType):
     name = 'integer'
 
@@ -450,13 +459,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-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_key):
 
     if confirm or test:
         # Confirmed but non-padded images don't make much sense, because
@@ -472,17 +485,24 @@ 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_key:
+        enckey = load_key(encrypt) if encrypt else None
+        if enckey and not match_sig_enc_key(key, enckey):
+            # 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,9 +547,26 @@ 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_key:
+        # Converting the command line provided raw AES key to byte array;
+        # this aray will be truncated to desired len.
+        plainkey = bytes.fromhex(aes_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)
         compressed_img = image.Image(version=decode_version(version),
@@ -575,13 +612,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)