imgtool: add AES‑KW encryption support and tests

docs/encrypted_images.md

@@ -176,7 +176,8 @@ To extract the public key in source file form, use
 format, use `imgtool getpub -k <input.pem> -e pem`.
 
 If using AES-KW, follow the steps in the next section to generate the
-required keys.
+required keys. The base64-encoded KEK can be passed to `imgtool sign`
+via the `--encrypt` option.
 
 ## [Creating your keys with Unix tooling](#creating-your-keys-with-unix-tooling)
 
@@ -187,6 +188,6 @@ required keys.
 * 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 (`newt` only), the `kek` can be generated with a
+* If using AES-KW, 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`

docs/imgtool.md

@@ -134,6 +134,8 @@ primary slot and adds a header and trailer that the bootloader is expecting:
       -E, --encrypt filename          Encrypt image using the provided public key.
                                       (Not supported in direct-xip or ram-load
                                       mode.)
+                                      For AES-KW, pass a base64-encoded KEK file
+                                      (16 bytes for 128-bit, 32 bytes for 256-bit).
       --save-enctlv                   When upgrading, save encrypted key TLVs
                                       instead of plain keys. Enable when
                                       BOOT_SWAP_SAVE_ENCTLV config option was set.

scripts/imgtool/image.py

@@ -36,6 +36,7 @@
 from cryptography.hazmat.primitives.asymmetric import ec, padding
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.keywrap import aes_key_wrap
 from cryptography.hazmat.primitives.kdf.hkdf import HKDF
 from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
 from intelhex import IntelHex
@@ -749,7 +750,7 @@ def create(self, key, public_key_format, enckey, dependencies=None,
             else:
                 plainkey = os.urandom(16)
 
-            if not isinstance(enckey, rsa.RSAPublic):
+            if not isinstance(enckey, rsa.RSAPublic) and not isinstance(enckey, keys.AESKWKey):
                 if hmac_sha == 'auto' or hmac_sha == '256':
                     hmac_sha = '256'
                     hmac_sha_alg = hashes.SHA256()
@@ -769,6 +770,15 @@ def create(self, key, public_key_format, enckey, dependencies=None,
                         label=None))
                 self.enctlv_len = len(cipherkey)
                 tlv.add('ENCRSA2048', cipherkey)
+            elif isinstance(enckey, keys.AESKWKey):
+                if not ((encrypt_keylen == 128 and len(enckey.kek) == 16) or
+                        (encrypt_keylen == 256 and len(enckey.kek) == 32)):
+                    raise click.UsageError(
+                        "AES-KW KEK size must match --encrypt-keylen"
+                    )
+                cipherkey = aes_key_wrap(enckey.kek, plainkey)
+                self.enctlv_len = len(cipherkey)
+                tlv.add('ENCKW', cipherkey)
             elif isinstance(enckey, ecdsa.ECDSA256P1Public):
                 cipherkey, mac, pubk = self.ecies_hkdf(enckey, plainkey, hmac_sha_alg)
                 enctlv = pubk + mac + cipherkey

scripts/imgtool/keys/__init__.py

@@ -19,6 +19,8 @@
 Cryptographic key management for imgtool.
 """
 
+import base64
+
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.ec import (
@@ -50,6 +52,7 @@
     "X25519",
     "X25519Public",
     "X25519UsageError",
+    "AESKWKey",
 ]
 
 
@@ -58,29 +61,63 @@ class PasswordRequired(Exception):
     password was not specified."""
 
 
-def load(path, passwd=None):
+class AESKWKey:
+    def __init__(self, kek):
+        self.kek = kek
+
+
+def load(path, passwd=None, allow_aes=False):
     """Try loading a key from the given path.
       Returns None if the password wasn't specified."""
     with open(path, 'rb') as f:
         raw_pem = f.read()
+    load_error = None
     try:
         pk = serialization.load_pem_private_key(
                 raw_pem,
                 password=passwd,
                 backend=default_backend())
     # Unfortunately, the crypto library raises unhelpful exceptions,
     # so we have to look at the text.
-    except TypeError as e:
-        msg = str(e)
-        if "private key is encrypted" in msg:
+    except Exception as e:
+        load_error = e
+        pk = None
+
+    if isinstance(load_error, TypeError):
+        msg = str(load_error)
+        if "private key is encrypted" in msg and passwd is None:
             return None
-        raise e
-    except ValueError:
-        # This seems to happen if the key is a public key, let's try
-        # loading it as a public key.
-        pk = serialization.load_pem_public_key(
-                raw_pem,
-                backend=default_backend())
+
+    # This seems to happen if the key is a public key, let's try
+    # loading it as a public key.
+    if isinstance(load_error, ValueError):
+        try:
+            pk = serialization.load_pem_public_key(
+                    raw_pem,
+                    backend=default_backend())
+        except Exception as e:
+            load_error = e
+            pk = None
+        else:
+            load_error = None
+
+    # Try base64 on malformed PEM
+    kek = None
+    if isinstance(load_error, ValueError) and allow_aes:
+        try:
+            kek = base64.b64decode(raw_pem.strip(), validate=True)
+        except Exception:
+            raise load_error
+
+    if kek is not None:
+        if len(kek) not in (16, 32):
+            raise Exception(
+                f"Invalid AES key length: {len(kek)} bytes. Expected 16 or 32 bytes after base64 decode."
+            )
+        return AESKWKey(kek)
+
+    if load_error is not None:
+        raise load_error
 
     if isinstance(pk, RSAPrivateKey):
         if pk.key_size not in RSA_KEY_SIZES:

scripts/imgtool/main.py

@@ -98,13 +98,13 @@ def save_signature(sigfile, sig):
         f.write(signature)
 
 
-def load_key(keyfile):
+def load_key(keyfile, allow_aes=False):
     # TODO: better handling of invalid pass-phrase
-    key = keys.load(keyfile)
+    key = keys.load(keyfile, allow_aes=allow_aes)
     if key is not None:
         return key
     passwd = getpass.getpass("Enter key passphrase: ").encode('utf-8')
-    return keys.load(keyfile, passwd)
+    return keys.load(keyfile, passwd, allow_aes=allow_aes)
 
 
 def get_password():
@@ -473,16 +473,18 @@ def sign(key, public_key_format, align, version, pad_sig, header_size,
     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))
+    enckey = load_key(encrypt, allow_aes=True) if encrypt else None
+    if enckey and key and not isinstance(enckey, keys.AESKWKey):
+        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")
+                not isinstance(enckey, keys.RSAPublic))
+        ):
+            # FIXME
+            raise click.UsageError("Signing and encryption must use the same "
+                                   "type of key")
 
     if pad_sig and hasattr(key, 'pad_sig'):
         key.pad_sig = True

scripts/tests/test_encryption_kw.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import struct
+from pathlib import Path
+
+import pytest
+from click.testing import CliRunner
+from imgtool import image
+from imgtool.main import imgtool
+
+VERSION = '1.2.3'
+HEADER_SIZE = 32
+SLOT_SIZE = 4096
+
+
+@pytest.fixture
+def key_file():
+    return Path(__file__).parents[2] / 'root-ec-p256.pem'
+
+
+def parse_tlvs(data, endian="<"):
+    header = data[:image.IMAGE_HEADER_SIZE]
+    _, _, hdr_sz, _, img_sz, flags = struct.unpack(endian + "IIHHII", header[:20])
+    tlv_off = hdr_sz + img_sz
+
+    magic, tlv_tot = struct.unpack(endian + "HH", data[tlv_off:tlv_off + image.TLV_INFO_SIZE])
+    if magic == image.TLV_PROT_INFO_MAGIC:
+        tlv_off += tlv_tot
+        magic, tlv_tot = struct.unpack(endian + "HH", data[tlv_off:tlv_off + image.TLV_INFO_SIZE])
+
+    assert magic == image.TLV_INFO_MAGIC
+
+    tlv_end = tlv_off + tlv_tot
+    tlv_off += image.TLV_INFO_SIZE
+    tlvs = []
+    while tlv_off < tlv_end:
+        tlv_type, _, tlv_len = struct.unpack(endian + "BBH", data[tlv_off:tlv_off + image.TLV_SIZE])
+        tlv_off += image.TLV_SIZE
+        payload = data[tlv_off:tlv_off + tlv_len]
+        tlvs.append((tlv_type, payload))
+        tlv_off += tlv_len
+
+    return flags, tlvs
+
+
+@pytest.mark.parametrize(
+    "keylen, kek_bytes, expected_len, enc_flag",
+    [
+        (128, bytes(range(16)), 24, image.IMAGE_F["ENCRYPTED_AES128"]),
+        (256, bytes(range(32)), 40, image.IMAGE_F["ENCRYPTED_AES256"]),
+    ],
+)
+def test_encrypt_aes_kw(keylen, kek_bytes, expected_len, enc_flag, tmpdir):
+    runner = CliRunner()
+
+    infile = tmpdir / "in.bin"
+    outfile = tmpdir / "out.bin"
+    kekfile = tmpdir / "kek.b64"
+
+    with infile.open("wb") as f:
+        f.write(bytes([0x11]) * 256)
+    with kekfile.open("wb") as f:
+        f.write(base64.b64encode(kek_bytes))
+
+    result = runner.invoke(
+        imgtool,
+        [
+            "sign",
+            f"--header-size={HEADER_SIZE}",
+            f"--slot-size={SLOT_SIZE}",
+            f"--version={VERSION}",
+            "--pad-header",
+            f"--encrypt={kekfile}",
+            f"--encrypt-keylen={keylen}",
+            str(infile),
+            str(outfile),
+        ],
+    )
+    assert result.exit_code == 0
+    assert outfile.exists()
+
+    with outfile.open("rb") as f:
+        data = f.read()
+    flags, tlvs = parse_tlvs(data)
+    assert flags & enc_flag
+
+    enckw_tlvs = [payload for tlv_type, payload in tlvs if tlv_type == image.TLV_VALUES["ENCKW"]]
+    assert len(enckw_tlvs) == 1
+    assert len(enckw_tlvs[0]) == expected_len
+
+
+@pytest.mark.parametrize(
+    "keylen, kek_bytes",
+    [
+        (128, bytes(range(16))),
+        (256, bytes(range(32))),
+    ],
+)
+def test_encrypt_aes_kw_with_signing_key(keylen, kek_bytes, tmpdir, key_file):
+    runner = CliRunner()
+
+    infile = tmpdir / "in.bin"
+    outfile = tmpdir / "out.bin"
+    kekfile = tmpdir / "kek.b64"
+
+    with infile.open("wb") as f:
+        f.write(bytes([0x22]) * 256)
+    with kekfile.open("wb") as f:
+        f.write(base64.b64encode(kek_bytes))
+
+    result = runner.invoke(
+        imgtool,
+        [
+            "sign",
+            f"--header-size={HEADER_SIZE}",
+            f"--slot-size={SLOT_SIZE}",
+            f"--version={VERSION}",
+            "--pad-header",
+            f"--encrypt={kekfile}",
+            f"--encrypt-keylen={keylen}",
+            f"--key={key_file}",
+            str(infile),
+            str(outfile),
+        ],
+    )
+    assert result.exit_code == 0
+    assert outfile.exists()