imgtool: Improve key type checks and invalid pass-phrase handling

Key type checking moved to separate function and added support for
ed25519, enckey is enforced to be a public key

Invalid pass-phrase was not reachable in methods as in that case an
exception raised instead of returning "None".
load_key function improved by adding handling for invalid pass-phrase
and FileNotFound exception.

Signed-off-by: Rustam Ismayilov <[email protected]>
Change-Id: I0caa0a6100fc95c8339403e991d6de0337c7a725

Author: rustammendel

scripts/imgtool/keys/__init__.py

@@ -1,5 +1,5 @@
 # Copyright 2017 Linaro Limited
-# Copyright 2023 Arm Limited
+# Copyright 2023-2024 Arm Limited
 #
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -58,14 +58,24 @@ def load(path, passwd=None):
     except TypeError as e:
         msg = str(e)
         if "private key is encrypted" in msg:
+            print(msg)
             return None
         raise e
-    except ValueError:
+    except ValueError as e:
+        msg1 = str(e)
         # 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(
+        try:
+            pk = serialization.load_pem_public_key(
                 raw_pem,
                 backend=default_backend())
+        except ValueError as e:
+            # If loading as public key also fails, that indicates wrong
+            # passphrase input
+            msg2 = str(e)
+            if ("password may be incorrect" in msg1 and
+                    "Are you sure this is a public key" in msg2):
+                raise Exception("Invalid passphrase")
 
     if isinstance(pk, RSAPrivateKey):
         if pk.key_size not in RSA_KEY_SIZES:

scripts/imgtool/keys/general.py

@@ -2,15 +2,25 @@
 
 # SPDX-License-Identifier: Apache-2.0
 
-import binascii
-import io
 import os
 import sys
+
+from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes, PublicKeyTypes
 from cryptography.hazmat.primitives.hashes import Hash, SHA256
 
+from imgtool import keys
+
 AUTOGEN_MESSAGE = "/* Autogenerated by imgtool.py, do not edit. */"
 
 
+def key_types_matching(key: PrivateKeyTypes, enckey: PublicKeyTypes):
+    type_dict = {keys.ECDSA256P1: keys.ECDSA256P1Public,
+                 keys.ECDSA384P1: keys.ECDSA384P1Public,
+                 keys.Ed25519: keys.X25519Public,
+                 keys.RSA: keys.RSAPublic}
+    return type_dict[type(key)] == type(enckey)
+
+
 class FileHandler(object):
     def __init__(self, file, *args, **kwargs):
         self.file_in = file
@@ -34,7 +44,7 @@ def _emit(self, header, trailer, encoded_bytes, indent, file=sys.stdout,
               len_format=None):
         with FileHandler(file, 'w') as file:
             self._emit_to_output(header, trailer, encoded_bytes, indent,
-                                     file, len_format)
+                                 file, len_format)
 
     def _emit_to_output(self, header, trailer, encoded_bytes, indent, file,
                         len_format):
@@ -62,27 +72,27 @@ def _emit_raw(self, encoded_bytes, file):
 
     def emit_c_public(self, file=sys.stdout):
         self._emit(
-                header="const unsigned char {}_pub_key[] = {{"
-                       .format(self.shortname()),
-                trailer="};",
-                encoded_bytes=self.get_public_bytes(),
-                indent="    ",
-                len_format="const unsigned int {}_pub_key_len = {{}};"
-                           .format(self.shortname()),
-                file=file)
+            header="const unsigned char {}_pub_key[] = {{"
+            .format(self.shortname()),
+            trailer="};",
+            encoded_bytes=self.get_public_bytes(),
+            indent="    ",
+            len_format="const unsigned int {}_pub_key_len = {{}};"
+            .format(self.shortname()),
+            file=file)
 
     def emit_c_public_hash(self, file=sys.stdout):
         digest = Hash(SHA256())
         digest.update(self.get_public_bytes())
         self._emit(
-                header="const unsigned char {}_pub_key_hash[] = {{"
-                       .format(self.shortname()),
-                trailer="};",
-                encoded_bytes=digest.finalize(),
-                indent="    ",
-                len_format="const unsigned int {}_pub_key_hash_len = {{}};"
-                           .format(self.shortname()),
-                file=file)
+            header="const unsigned char {}_pub_key_hash[] = {{"
+            .format(self.shortname()),
+            trailer="};",
+            encoded_bytes=digest.finalize(),
+            indent="    ",
+            len_format="const unsigned int {}_pub_key_hash_len = {{}};"
+            .format(self.shortname()),
+            file=file)
 
     def emit_raw_public(self, file=sys.stdout):
         self._emit_raw(self.get_public_bytes(), file=file)
@@ -94,22 +104,22 @@ def emit_raw_public_hash(self, file=sys.stdout):
 
     def emit_rust_public(self, file=sys.stdout):
         self._emit(
-                header="static {}_PUB_KEY: &[u8] = &["
-                       .format(self.shortname().upper()),
-                trailer="];",
-                encoded_bytes=self.get_public_bytes(),
-                indent="    ",
-                file=file)
+            header="static {}_PUB_KEY: &[u8] = &["
+            .format(self.shortname().upper()),
+            trailer="];",
+            encoded_bytes=self.get_public_bytes(),
+            indent="    ",
+            file=file)
 
     def emit_public_pem(self, file=sys.stdout):
         with FileHandler(file, 'w') as file:
             print(str(self.get_public_pem(), 'utf-8'), file=file, end='')
 
     def emit_private(self, minimal, format, file=sys.stdout):
         self._emit(
-                header="const unsigned char enc_priv_key[] = {",
-                trailer="};",
-                encoded_bytes=self.get_private_bytes(minimal, format),
-                indent="    ",
-                len_format="const unsigned int enc_priv_key_len = {};",
-                file=file)
+            header="const unsigned char enc_priv_key[] = {",
+            trailer="};",
+            encoded_bytes=self.get_private_bytes(minimal, format),
+            indent="    ",
+            len_format="const unsigned int enc_priv_key_len = {};",
+            file=file)

scripts/imgtool/main.py

@@ -1,7 +1,7 @@
 #! /usr/bin/env python3
 #
 # Copyright 2017-2020 Linaro Limited
-# Copyright 2019-2023 Arm Limited
+# Copyright 2019-2024 Arm Limited
 #
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -24,6 +24,7 @@
 import sys
 import base64
 from imgtool import image, imgtool_version
+from imgtool.keys.general import key_types_matching
 from imgtool.version import decode_version
 from imgtool.dumpinfo import dump_imginfo
 from .keys import (
@@ -87,12 +88,29 @@ def save_signature(sigfile, sig):
 
 
 def load_key(keyfile):
-    # TODO: better handling of invalid pass-phrase
-    key = keys.load(keyfile)
+    try:
+        key = keys.load(keyfile)
+    except FileNotFoundError as e:
+        print("Key File Not Found in the path: " + keyfile)
+        raise e
     if key is not None:
         return key
+
+    # Key is password protected
     passwd = getpass.getpass("Enter key passphrase: ").encode('utf-8')
-    return keys.load(keyfile, passwd)
+    try:
+        key = keys.load(keyfile, passwd)
+    except Exception as e:
+        msg = str(e)
+        if "Invalid passphrase" in msg:
+            print(msg)
+            exit(1)
+        else:
+            raise e
+    if key is None:
+        print("Could not load key for unknown error")
+        exit(1)
+    return key
 
 
 def get_password():
@@ -145,9 +163,8 @@ def getpub(key, encoding, lang, output):
 
     if not output:
         output = sys.stdout
-    if key is None:
-        print("Invalid passphrase")
-    elif lang == 'c' or encoding == 'lang-c':
+
+    if lang == 'c' or encoding == 'lang-c':
         key.emit_c_public(file=output)
     elif lang == 'rust' or encoding == 'lang-rust':
         key.emit_rust_public(file=output)
@@ -177,9 +194,8 @@ def getpubhash(key, output, encoding):
 
     if not output:
         output = sys.stdout
-    if key is None:
-        print("Invalid passphrase")
-    elif encoding == 'lang-c':
+
+    if encoding == 'lang-c':
         key.emit_c_public_hash(file=output)
     elif encoding == 'raw':
         key.emit_raw_public_hash(file=output)
@@ -200,8 +216,6 @@ def getpubhash(key, output, encoding):
 @click.command(help='Dump private key from keypair')
 def getpriv(key, minimal, format):
     key = load_key(key)
-    if key is None:
-        print("Invalid passphrase")
     try:
         key.emit_private(minimal, format)
     except (RSAUsageError, ECDSAUsageError, Ed25519UsageError,
@@ -430,16 +444,9 @@ def sign(key, public_key_format, align, version, pad_sig, header_size,
     img.load(infile)
     key = load_key(key) if key else None
     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")
+    if enckey and key and not key_types_matching(key, enckey):
+        raise click.UsageError("Encryption must use the public pair of the "
+                               "same type of key used for signing")
 
     if pad_sig and hasattr(key, 'pad_sig'):
         key.pad_sig = True