boot: Add AES256 support for image encryption

Support only works when using mbedtls as the cryptographic library.

Signed-off-by: Salome Thirot <salome.thirot@arm.com>
diff --git a/NOTICE b/NOTICE
index 0b3435b..ce203c1 100644
--- a/NOTICE
+++ b/NOTICE
@@ -8,4 +8,4 @@
 Runtime Inc, copyright 2015.
 
 Portions of this software were developed at
-Arm Limited, copyright 2019-2020.
+Arm Limited, copyright 2019-2021.
diff --git a/boot/bootutil/include/bootutil/caps.h b/boot/bootutil/include/bootutil/caps.h
index 8ead3be..9bf060c 100644
--- a/boot/bootutil/include/bootutil/caps.h
+++ b/boot/bootutil/include/bootutil/caps.h
@@ -1,5 +1,6 @@
 /*
  * Copyright (c) 2017 Linaro Limited
+ * Copyright (c) 2021 Arm Limited
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -47,6 +48,7 @@
 #define BOOTUTIL_CAP_DOWNGRADE_PREVENTION   (1<<12)
 #define BOOTUTIL_CAP_ENC_X25519             (1<<13)
 #define BOOTUTIL_CAP_BOOTSTRAP              (1<<14)
+#define BOOTUTIL_CAP_AES256                 (1<<15)
 
 /*
  * Query the number of images this bootloader is configured for.  This
diff --git a/boot/bootutil/include/bootutil/crypto/aes_ctr.h b/boot/bootutil/include/bootutil/crypto/aes_ctr.h
index 8422f45..e69b037 100644
--- a/boot/bootutil/include/bootutil/crypto/aes_ctr.h
+++ b/boot/bootutil/include/bootutil/crypto/aes_ctr.h
@@ -21,11 +21,15 @@
 
 #if defined(MCUBOOT_USE_MBED_TLS)
     #include <mbedtls/aes.h>
-    #define BOOTUTIL_CRYPTO_AES_CTR_KEY_SIZE (16)
+    #include "bootutil/enc_key_public.h"
+    #define BOOTUTIL_CRYPTO_AES_CTR_KEY_SIZE BOOT_ENC_KEY_SIZE
     #define BOOTUTIL_CRYPTO_AES_CTR_BLOCK_SIZE (16)
 #endif /* MCUBOOT_USE_MBED_TLS */
 
 #if defined(MCUBOOT_USE_TINYCRYPT)
+    #if defined(MCUBOOT_AES_256)
+        #error "Cannot use AES-256 for encryption with Tinycrypt library."
+    #endif
     #include <string.h>
     #include <tinycrypt/aes.h>
     #include <tinycrypt/ctr_mode.h>
diff --git a/boot/bootutil/include/bootutil/crypto/aes_kw.h b/boot/bootutil/include/bootutil/crypto/aes_kw.h
index 925f46a..cf3194f 100644
--- a/boot/bootutil/include/bootutil/crypto/aes_kw.h
+++ b/boot/bootutil/include/bootutil/crypto/aes_kw.h
@@ -23,6 +23,9 @@
 #endif /* MCUBOOT_USE_MBED_TLS */
 
 #if defined(MCUBOOT_USE_TINYCRYPT)
+    #if defined(MCUBOOT_AES_256)
+        #error "Cannot use AES-256 for encryption with Tinycrypt library."
+    #endif
     #include <tinycrypt/aes.h>
     #include <tinycrypt/constants.h>
 #endif /* MCUBOOT_USE_TINYCRYPT */
diff --git a/boot/bootutil/include/bootutil/enc_key_public.h b/boot/bootutil/include/bootutil/enc_key_public.h
index 634f842..420d0d6 100644
--- a/boot/bootutil/include/bootutil/enc_key_public.h
+++ b/boot/bootutil/include/bootutil/enc_key_public.h
@@ -2,7 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2018-2019 JUUL Labs
- * Copyright (c) 2019 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  * Copyright (c) 2021 Nordic Semiconductor ASA
  *
  * Original license:
@@ -32,12 +32,16 @@
 extern "C" {
 #endif
 
+#ifdef MCUBOOT_AES_256
+#define BOOT_ENC_KEY_SIZE       32
+#else
 #define BOOT_ENC_KEY_SIZE       16
+#endif
 
 #define TLV_ENC_RSA_SZ    256
-#define TLV_ENC_KW_SZ     24
-#define TLV_ENC_EC256_SZ  (65 + 32 + 16)
-#define TLV_ENC_X25519_SZ (32 + 32 + 16)
+#define TLV_ENC_KW_SZ     BOOT_ENC_KEY_SIZE + 8
+#define TLV_ENC_EC256_SZ  (65 + 32 + BOOT_ENC_KEY_SIZE)
+#define TLV_ENC_X25519_SZ (32 + 32 + BOOT_ENC_KEY_SIZE)
 
 #if defined(MCUBOOT_ENCRYPT_RSA)
 #define BOOT_ENC_TLV_SIZE TLV_ENC_RSA_SZ
@@ -53,4 +57,4 @@
 }
 #endif
 
-#endif /* BOOTUTIL_ENC_KEY_PUBLIC_H */
\ No newline at end of file
+#endif /* BOOTUTIL_ENC_KEY_PUBLIC_H */
diff --git a/boot/bootutil/include/bootutil/image.h b/boot/bootutil/include/bootutil/image.h
index 38304f0..fc94915 100644
--- a/boot/bootutil/include/bootutil/image.h
+++ b/boot/bootutil/include/bootutil/image.h
@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2016-2019 Linaro LTD
  * Copyright (c) 2016-2019 JUUL Labs
- * Copyright (c) 2019-2020 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  *
  * Original license:
  *
@@ -50,7 +50,8 @@
  * Image header flags.
  */
 #define IMAGE_F_PIC                      0x00000001 /* Not supported. */
-#define IMAGE_F_ENCRYPTED                0x00000004 /* Encrypted. */
+#define IMAGE_F_ENCRYPTED_AES128         0x00000004 /* Encrypted using AES128. */
+#define IMAGE_F_ENCRYPTED_AES256         0x00000008 /* Encrypted using AES256. */
 #define IMAGE_F_NON_BOOTABLE             0x00000010 /* Split image app. */
 /*
  * Indicates that this image should be loaded into RAM instead of run
@@ -89,7 +90,7 @@
 #define IMAGE_TLV_RSA3072_PSS       0x23   /* RSA3072 of hash output */
 #define IMAGE_TLV_ED25519           0x24   /* ed25519 of hash output */
 #define IMAGE_TLV_ENC_RSA2048       0x30   /* Key encrypted with RSA-OAEP-2048 */
-#define IMAGE_TLV_ENC_KW128         0x31   /* Key encrypted with AES-KW-128 */
+#define IMAGE_TLV_ENC_KW            0x31   /* Key encrypted with AES-KW 128 or 256*/
 #define IMAGE_TLV_ENC_EC256         0x32   /* Key encrypted with ECIES-EC256 */
 #define IMAGE_TLV_ENC_X25519        0x33   /* Key encrypted with ECIES-X25519 */
 #define IMAGE_TLV_DEPENDENCY        0x40   /* Image depends on other image */
@@ -148,7 +149,8 @@
     uint16_t it_len;    /* Data length (not including TLV header). */
 };
 
-#define IS_ENCRYPTED(hdr) ((hdr)->ih_flags & IMAGE_F_ENCRYPTED)
+#define IS_ENCRYPTED(hdr) (((hdr)->ih_flags && IMAGE_F_ENCRYPTED_AES128) \
+                        || ((hdr)->ih_flags && IMAGE_F_ENCRYPTED_AES256))
 #define MUST_DECRYPT(fap, idx, hdr) \
     ((fap)->fa_id == FLASH_AREA_IMAGE_SECONDARY(idx) && IS_ENCRYPTED(hdr))
 
diff --git a/boot/bootutil/src/caps.c b/boot/bootutil/src/caps.c
index a21bc09..87bac1a 100644
--- a/boot/bootutil/src/caps.c
+++ b/boot/bootutil/src/caps.c
@@ -2,6 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2017 Linaro Limited
+ * Copyright (c) 2021 Arm Limited
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -68,6 +69,9 @@
 #if defined(MCUBOOT_BOOTSTRAP)
     res |= BOOTUTIL_CAP_BOOTSTRAP;
 #endif
+#if defined(MCUBOOT_AES_256)
+    res |= BOOTUTIL_CAP_AES256;
+#endif
 
     return res;
 }
diff --git a/boot/bootutil/src/encrypted.c b/boot/bootutil/src/encrypted.c
index d14855b..278f51d 100644
--- a/boot/bootutil/src/encrypted.c
+++ b/boot/bootutil/src/encrypted.c
@@ -2,7 +2,7 @@
  * SPDX-License-Identifier: Apache-2.0
  *
  * Copyright (c) 2018-2019 JUUL Labs
- * Copyright (c) 2019 Arm Limited
+ * Copyright (c) 2019-2021 Arm Limited
  */
 
 #include "mcuboot_config/mcuboot_config.h"
@@ -434,20 +434,20 @@
 #if defined(MCUBOOT_ENCRYPT_RSA)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_RSA2048
 #elif defined(MCUBOOT_ENCRYPT_KW)
-#    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_KW128
+#    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_KW
 #elif defined(MCUBOOT_ENCRYPT_EC256)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_EC256
 #    define EC_PUBK_INDEX       (0)
 #    define EC_TAG_INDEX        (65)
 #    define EC_CIPHERKEY_INDEX  (65 + 32)
-_Static_assert(EC_CIPHERKEY_INDEX + 16 == EXPECTED_ENC_LEN,
+_Static_assert(EC_CIPHERKEY_INDEX + BOOT_ENC_KEY_SIZE == EXPECTED_ENC_LEN,
         "Please fix ECIES-P256 component indexes");
 #elif defined(MCUBOOT_ENCRYPT_X25519)
 #    define EXPECTED_ENC_TLV    IMAGE_TLV_ENC_X25519
 #    define EC_PUBK_INDEX       (0)
 #    define EC_TAG_INDEX        (32)
 #    define EC_CIPHERKEY_INDEX  (32 + 32)
-_Static_assert(EC_CIPHERKEY_INDEX + 16 == EXPECTED_ENC_LEN,
+_Static_assert(EC_CIPHERKEY_INDEX + BOOT_ENC_KEY_SIZE == EXPECTED_ENC_LEN,
         "Please fix ECIES-X25519 component indexes");
 #endif
 
@@ -455,7 +455,7 @@
  * Decrypt an encryption key TLV.
  *
  * @param buf An encryption TLV read from flash (build time fixed length)
- * @param enckey An AES-128 key sized buffer to store to plain key.
+ * @param enckey An AES-128 or AES-256 key sized buffer to store to plain key.
  */
 int
 boot_enc_decrypt(const uint8_t *buf, uint8_t *enckey)
@@ -507,7 +507,7 @@
 
 #if defined(MCUBOOT_ENCRYPT_KW)
 
-    assert(*bootutil_enc_key.len == 16);
+    assert(*bootutil_enc_key.len == BOOT_ENC_KEY_SIZE);
     rc = key_unwrap(buf, enckey);
 
 #endif /* defined(MCUBOOT_ENCRYPT_KW) */
@@ -586,13 +586,13 @@
 
     bootutil_hmac_sha256_init(&hmac);
 
-    rc = bootutil_hmac_sha256_set_key(&hmac, &derived_key[16], 32);
+    rc = bootutil_hmac_sha256_set_key(&hmac, &derived_key[BOOT_ENC_KEY_SIZE], 32);
     if (rc != 0) {
         (void)bootutil_hmac_sha256_drop(&hmac);
         return -1;
     }
 
-    rc = bootutil_hmac_sha256_update(&hmac, &buf[EC_CIPHERKEY_INDEX], 16);
+    rc = bootutil_hmac_sha256_update(&hmac, &buf[EC_CIPHERKEY_INDEX], BOOT_ENC_KEY_SIZE);
     if (rc != 0) {
         (void)bootutil_hmac_sha256_drop(&hmac);
         return -1;
diff --git a/docs/design.md b/docs/design.md
index a50feeb..d498926 100644
--- a/docs/design.md
+++ b/docs/design.md
@@ -3,7 +3,7 @@
 
   - Copyright (c) 2017-2020 Linaro LTD
   - Copyright (c) 2017-2019 JUUL Labs
-  - Copyright (c) 2019-2020 Arm Limited
+  - Copyright (c) 2019-2021 Arm Limited
 
   - Original license:
 
@@ -96,6 +96,8 @@
  * Image header flags.
  */
 #define IMAGE_F_PIC                      0x00000001 /* Not supported. */
+#define IMAGE_F_ENCRYPTED_AES128         0x00000004 /* Encrypted using AES128. */
+#define IMAGE_F_ENCRYPTED_AES256         0x00000008 /* Encrypted using AES256. */
 #define IMAGE_F_NON_BOOTABLE             0x00000010 /* Split image app. */
 #define IMAGE_F_RAM_LOAD                 0x00000020
 
@@ -110,7 +112,8 @@
 #define IMAGE_TLV_RSA3072_PSS       0x23   /* RSA3072 of hash output */
 #define IMAGE_TLV_ED25519           0x24   /* ED25519 of hash output */
 #define IMAGE_TLV_ENC_RSA2048       0x30   /* Key encrypted with RSA-OAEP-2048 */
-#define IMAGE_TLV_ENC_KW128         0x31   /* Key encrypted with AES-KW-128 */
+#define IMAGE_TLV_ENC_KW            0x31   /* Key encrypted with AES-KW-128 or
+                                              256 */
 #define IMAGE_TLV_ENC_EC256         0x32   /* Key encrypted with ECIES-P256 */
 #define IMAGE_TLV_ENC_X25519        0x33   /* Key encrypted with ECIES-X25519 */
 #define IMAGE_TLV_DEPENDENCY        0x40   /* Image depends on other image */
@@ -1285,4 +1288,4 @@
   the shell returns, and it is possible to investigate the results. It is also
   possible to stop the test with _Ctrl+c_. The parameters to the
   `execute_test.sh` are `SKIP_SIZE`, `BUILD_TYPE`, `DAMAGE_TYPE`, `FIH_LEVEL` in
-  order.
\ No newline at end of file
+  order.
diff --git a/docs/encrypted_images.md b/docs/encrypted_images.md
index 74443db..838f493 100644
--- a/docs/encrypted_images.md
+++ b/docs/encrypted_images.md
@@ -59,9 +59,9 @@
 that the encrypted flag is set and TLV data is OK, then it decrypts each
 image block before sending the data to the hash routines.
 
-The image is encrypted using AES-CTR-128, with a counter that starts
-from zero (over the payload blocks) and increments by 1 for each 16-byte
-block. AES-CTR-128 was chosen for speed/simplicity and allowing for any
+The image is encrypted using AES-CTR-128 or AES-CTR-256, with a counter
+that starts from zero (over the payload blocks) and increments by 1 for each
+16-byte block. AES-CTR was chosen for speed/simplicity and allowing for any
 block to be encrypted/decrypted without requiring knowledge of any other
 block (allowing for simple resume operations on swap interruptions).
 
@@ -70,14 +70,15 @@
 but randomizing a 16-byte block with a TRNG should make it highly
 improbable that duplicates ever happen.
 
-To distribute this AES-CTR-128 key, new TLVs were defined. The key can be
-encrypted using either RSA-OAEP, AES-KW-128, ECIES-P256 or ECIES-X25519.
+To distribute this AES-CTR key, new TLVs were defined. The key can be
+encrypted using either RSA-OAEP, AES-KW (128 or 256 bits depending on the
+AES-CTR key length), ECIES-P256 or ECIES-X25519.
 
 For RSA-OAEP a new TLV with value `0x30` is added to the image, for
-AES-KW-128 a new TLV with value `0x31` is added to the image, for
+AES-KW a new TLV with value `0x31` is added to the image, for
 ECIES-P256 a new TLV with value `0x32` is added, and for ECIES-X25519 a
 newt TLV with value `0x33` is added. The contents of those TLVs
-are the results of applying the given operations over the AES-CTR-128 key.
+are the results of applying the given operations over the AES-CTR key.
 
 ## [ECIES encryption](#ecies-encryption)
 
@@ -94,17 +95,17 @@
 * Derive the new keys from the secret using HKDF (built on HMAC-SHA256). We
   are not using a `salt` and using an `info` of `MCUBoot_ECIES_v1`, generating
   48 bytes of key material.
-* A new random encryption key of 16 bytes is generated (for AES-128). This is
+* A new random encryption key is generated (for AES). This is
   the AES key used to encrypt the images.
-* The key is encrypted with AES-128-CTR and a `nonce` of 0 using the first
-  16 bytes of key material generated previously by the HKDF.
+* The key is encrypted with AES-128-CTR or AES-256-CTR and a `nonce` of 0 using
+  the first 16 bytes of key material generated previously by the HKDF.
 * The encrypted key now goes through a HMAC-SHA256 using the remaining 32
   bytes of key material from the HKDF.
 
 The final TLV is built from the 65 bytes for ECIES-P256 or 32 bytes for
 ECIES-X25519, which correspond to the ephemeral public key, followed by the
-32 bytes of MAC tag and the 16 bytes of the encrypted key, resulting in a TLV
-of 113 bytes for ECIES-P256 or 80 bytes for ECIES-X25519.
+32 bytes of MAC tag and the 16 or 32 bytes of the encrypted key, resulting in
+a TLV of 113 or 129 bytes for ECIES-P256 and 80 or 96 bytes for ECIES-X25519.
 
 The implemenation of ECIES-P256 is named ENC_EC256 in the source code and
 artifacts while ECIES-X25519 is named ENC_X25519.
@@ -149,7 +150,7 @@
 `imgtool getpub -k <input.pem> -l <lang>`, where lang can be one of `c` or
 `rust` (defaults to `c`).
 
-If using AES-KW-128, follow the steps in the next section to generate the
+If using AES-KW, follow the steps in the next section to generate the
 required keys.
 
 ## [Creating your keys with Unix tooling](#creating-your-keys-with-unix-tooling)
@@ -161,5 +162,6 @@
 * If using ECIES-X25519, generate a private key passing the option `-t x25519`
   to `imgtool keygen` command. To generate public key PEM file the following
   command can be used: `openssl pkey -in <generated-private-key.pem> -pubout`
-* If using AES-KW-128 (`newt` only), the `kek` can be generated with a
-  command like `dd if=/dev/urandom bs=1 count=16 | base64 > my_kek.b64`
+* If using AES-KW (`newt` only), the `kek` can be generated with a
+  command like (change count to 32 for a 256 bit key)
+  `dd if=/dev/urandom bs=1 count=16 | base64 > my_kek.b64`
diff --git a/enc-aes256kw.b64 b/enc-aes256kw.b64
new file mode 100644
index 0000000..866aa6d
--- /dev/null
+++ b/enc-aes256kw.b64
@@ -0,0 +1 @@
+5FxRRtIcgjXMGhmvofKqIMiMf0Bs2yKqarXLqvixW7Q=
diff --git a/scripts/imgtool/image.py b/scripts/imgtool/image.py
index 257c892..684c6b3 100644
--- a/scripts/imgtool/image.py
+++ b/scripts/imgtool/image.py
@@ -1,6 +1,6 @@
 # Copyright 2018 Nordic Semiconductor ASA
 # Copyright 2017-2020 Linaro Limited
-# Copyright 2019-2020 Arm Limited
+# Copyright 2019-2021 Arm Limited
 #
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -51,7 +51,8 @@
 # Image header flags.
 IMAGE_F = {
         'PIC':                   0x0000001,
-        'ENCRYPTED':             0x0000004,
+        'ENCRYPTED_AES128':      0x0000004,
+        'ENCRYPTED_AES256':      0x0000008,
         'NON_BOOTABLE':          0x0000010,
         'RAM_LOAD':              0x0000020,
         'ROM_FIXED':             0x0000100,
@@ -67,7 +68,7 @@
         'RSA3072': 0x23,
         'ED25519': 0x24,
         'ENCRSA2048': 0x30,
-        'ENCKW128': 0x31,
+        'ENCKW': 0x31,
         'ENCEC256': 0x32,
         'ENCX25519': 0x33,
         'DEPENDENCY': 0x40,
@@ -288,7 +289,7 @@
         return cipherkey, ciphermac, pubk
 
     def create(self, key, public_key_format, enckey, dependencies=None,
-               sw_type=None, custom_tlvs=None):
+               sw_type=None, custom_tlvs=None, encrypt_keylen=128):
         self.enckey = enckey
 
         # Calculate the hash of the public key
@@ -361,7 +362,10 @@
                     self.payload.extend(pad)
 
         # This adds the header to the payload as well
-        self.add_header(enckey, protected_tlv_size)
+        if encrypt_keylen == 256:
+            self.add_header(enckey, protected_tlv_size, 256)
+        else:
+            self.add_header(enckey, protected_tlv_size)
 
         prot_tlv = TLV(self.endian, TLV_PROT_INFO_MAGIC)
 
@@ -429,7 +433,10 @@
             self.payload = self.payload[:protected_tlv_off]
 
         if enckey is not None:
-            plainkey = os.urandom(16)
+            if encrypt_keylen == 256:
+                plainkey = os.urandom(32)
+            else:
+                plainkey = os.urandom(16)
 
             if isinstance(enckey, rsa.RSAPublic):
                 cipherkey = enckey._get_public().encrypt(
@@ -462,12 +469,15 @@
 
         self.check_trailer()
 
-    def add_header(self, enckey, protected_tlv_size):
+    def add_header(self, enckey, protected_tlv_size, aes_length=128):
         """Install the image header."""
 
         flags = 0
         if enckey is not None:
-            flags |= IMAGE_F['ENCRYPTED']
+            if aes_length == 128:
+                flags |= IMAGE_F['ENCRYPTED_AES128']
+            else:
+                flags |= IMAGE_F['ENCRYPTED_AES256']
         if self.load_addr != 0:
             # Indicates that this image should be loaded into RAM
             # instead of run directly from flash.
diff --git a/scripts/imgtool/main.py b/scripts/imgtool/main.py
index 6188e78..dd6c044 100755
--- a/scripts/imgtool/main.py
+++ b/scripts/imgtool/main.py
@@ -1,7 +1,7 @@
 #! /usr/bin/env python3
 #
 # Copyright 2017-2020 Linaro Limited
-# Copyright 2019-2020 Arm Limited
+# Copyright 2019-2021 Arm Limited
 #
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -250,6 +250,10 @@
 @click.option('-E', '--encrypt', metavar='filename',
               help='Encrypt image using the provided public key. '
                    '(Not supported in direct-xip or ram-load mode.)')
+@click.option('--encrypt-keylen', default='128',
+              type=click.Choice(['128','256']),
+              help='When encrypting the image using AES, select a 128 bit or '
+                   '256 bit key len.')
 @click.option('-e', '--endian', type=click.Choice(['little', 'big']),
               default='little', help="Select little or big endian")
 @click.option('--overwrite-only', default=False, is_flag=True,
@@ -295,9 +299,9 @@
                .hex extension, otherwise binary format is used''')
 def sign(key, public_key_format, align, version, pad_sig, header_size,
          pad_header, slot_size, pad, confirm, max_sectors, overwrite_only,
-         endian, encrypt, infile, outfile, dependencies, load_addr, hex_addr,
-         erased_val, save_enctlv, security_counter, boot_record, custom_tlv,
-         rom_fixed):
+         endian, encrypt_keylen, encrypt, infile, outfile, dependencies,
+         load_addr, hex_addr, erased_val, save_enctlv, security_counter,
+         boot_record, custom_tlv, rom_fixed):
 
     if confirm:
         # Confirmed but non-padded images don't make much sense, because
@@ -344,7 +348,7 @@
             custom_tlvs[tag] = value.encode('utf-8')
 
     img.create(key, public_key_format, enckey, dependencies, boot_record,
-               custom_tlvs)
+               custom_tlvs, int(encrypt_keylen))
     img.save(outfile, hex_addr)
 
 
diff --git a/sim/src/tlv.rs b/sim/src/tlv.rs
index b6c3c96..a7b2c2d 100644
--- a/sim/src/tlv.rs
+++ b/sim/src/tlv.rs
@@ -49,7 +49,7 @@
     RSA3072 = 0x23,
     ED25519 = 0x24,
     ENCRSA2048 = 0x30,
-    ENCKW128 = 0x31,
+    ENCKW = 0x31,
     ENCEC256 = 0x32,
     ENCX25519 = 0x33,
     DEPENDENCY = 0x40,
@@ -181,7 +181,7 @@
     pub fn new_enc_kw() -> TlvGen {
         TlvGen {
             flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCKW128],
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
@@ -190,7 +190,7 @@
     pub fn new_rsa_kw() -> TlvGen {
         TlvGen {
             flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::RSA2048, TlvKinds::ENCKW128],
+            kinds: vec![TlvKinds::SHA256, TlvKinds::RSA2048, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
@@ -199,7 +199,7 @@
     pub fn new_ecdsa_kw() -> TlvGen {
         TlvGen {
             flags: TlvFlags::ENCRYPTED as u32,
-            kinds: vec![TlvKinds::SHA256, TlvKinds::ECDSA256, TlvKinds::ENCKW128],
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ECDSA256, TlvKinds::ENCKW],
             ..Default::default()
         }
     }
@@ -479,7 +479,7 @@
             result.extend_from_slice(&encbuf);
         }
 
-        if self.kinds.contains(&TlvKinds::ENCKW128) {
+        if self.kinds.contains(&TlvKinds::ENCKW) {
             let key_bytes = base64::decode(
                 include_str!("../../enc-aes128kw.b64").trim()).unwrap();
 
@@ -491,7 +491,7 @@
             };
 
             assert!(encbuf.len() == 24);
-            result.write_u16::<LittleEndian>(TlvKinds::ENCKW128 as u16).unwrap();
+            result.write_u16::<LittleEndian>(TlvKinds::ENCKW as u16).unwrap();
             result.write_u16::<LittleEndian>(24).unwrap();
             result.extend_from_slice(&encbuf);
         }