Refactor imgtool.py to handle PKCS mcu-tools#11 URIs

ref: mcu-tools#599
which is still not mainlined

ref: grandcentrix/mcuboot@d7c58b62
part 3 of a 3 commit series

Author: amoskopp

scripts/imgtool/keys/__init__.py

@@ -50,22 +50,21 @@ class PasswordRequired(Exception):
 
 
 def load(path, passwd=None):
-    """Try loading a key from the given path.
-      Returns None if the password wasn't specified."""
-    if path == "pkcs11":
+    if path.startswith("pkcs11:"):
         try:
-            return PKCS11()
+            return PKCS11(path)  # assume a PKCS #11 URI according to RFC7512
         except pkcs11.exceptions.PinIncorrect:
             print('ERROR: WRONG PIN')
             sys.exit(1)
         except pkcs11.exceptions.PinLocked:
             print('ERROR: WRONG PIN, MAX ATTEMPTS REACHED. CONTACT YOUR SECURITY OFFICER.')
             sys.exit(1)
         except pkcs11.exceptions.DataLenRange:
-            print('ERROR: \'DataLenRange\' (maybe the PIN is too short?)')
+            print('ERROR: PIN IS TOO SHORT OR TOO LONG')
             sys.exit(1)
 
-    """Try loading a key from the given path.  Returns None if the password wasn't specified."""
+    """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()
     try:

scripts/imgtool/keys/imgtool_keys_pkcs11.py

@@ -8,23 +8,83 @@
 import pkcs11
 import pkcs11.util.ec
 
+from pathlib import Path
+from urllib.parse import unquote, urlparse
+
 from .general import KeyClass
 
 
+def find_library(filename):
+    search_directory = '/usr'
+    for root, directorios, filenames in os.walk(search_directory):
+        if filename in filenames:
+            return os.path.join(
+                root,
+                filename,
+            )
+
+def unquote_to_bytes(urlencoded_string):
+    """Replace %xx escapes by their single-character equivalent,
+using the “iso-8859-1” encoding to decode all 8-bit values.
+    """
+    return bytes(
+        unquote(urlencoded_string, encoding='iso-8859-1'),
+        encoding='iso-8859-1'
+    )
+
+def get_pkcs11_uri_params(uri):
+    uri_tokens = urlparse(uri)
+    assert(uri_tokens.scheme == 'pkcs11')
+    assert(uri_tokens.query == '')
+    assert(uri_tokens.fragment == '')
+    return {
+        unquote_to_bytes(key): unquote_to_bytes(value)
+        for key, value
+        in [
+            line.split('=')
+            for line
+            in uri_tokens.path.split(';')
+        ]
+    }
+
 class PKCS11UsageError(Exception):
     pass
 
 
 class PKCS11(KeyClass):
-    def __init__(self, env=os.environ):
-        lib = pkcs11.lib(env['PKCS11_MODULE'])
-        self.token = lib.get_token(token_label=env['PKCS11_TOKEN_LABEL'])
-        self.key_id = bytes.fromhex(env['PKCS11_KEY_ID'])
-        self.session = self.token.open(user_pin=env['PKCS11_PIN'])
-
-    def __del__(self):
-        if hasattr(self, 'session'):
-            self.session.close()
+    def __init__(self, uri, env=os.environ):
+        if not 'PKCS11_PIN' in env.keys():
+            raise RuntimeError("Environment variable PKCS11_PIN not set. Set it to the user PIN.")
+        params = get_pkcs11_uri_params(uri)
+        assert(b'serial' in params.keys())
+        self.user_pin = env['PKCS11_PIN']
+
+        # Try to find PKCS #11 module path like this:
+        # Use the environment variable $PKCS11_MODULE,
+        # then fall back to OpenSC (for NitroKey HSM).
+        pkcs11_module_path = ''
+        try:
+            pkcs11_module_path = env['PKCS11_MODULE']
+            if not Path(pkcs11_module_path).is_file():
+                raise RuntimeError('$PKCS11_MODULE is not a file: %s' % pkcs11_module_path)
+        except KeyError:
+            pkcs11_module_path = find_library('opensc-pkcs11.so')
+            if pkcs11_module_path is None:
+                raise RuntimeError('$PKCS11_MODULE not set and OpenSC not found.')
+
+        lib = ''
+        try:
+            lib = pkcs11.lib(pkcs11_module_path)
+        except RuntimeError:
+            pass  # happens if lib does not exist or is corrupt
+        if '' == lib:
+            raise RuntimeError('PKCS #11 module not loaded.')
+
+        self.token = lib.get_token(token_serial=params[b'serial'])
+        # try to open a session to see if the PIN is valid
+        with self.token.open(user_pin=self.user_pin) as session:
+            pass
+        self.key_id = params[b'id']
 
     def shortname(self):
         return "ecdsa"
@@ -33,11 +93,16 @@ def _unsupported(self, name):
         raise PKCS11UsageError("Operation {} requires private key".format(name))
 
     def get_public_bytes(self):
-        pub = self.session.get_key(
-            id=self.key_id, key_type=pkcs11.KeyType.EC, object_class=pkcs11.ObjectClass.PUBLIC_KEY)
-        key = oscrypto.asymmetric.load_public_key(
-            pkcs11.util.ec.encode_ec_public_key(pub))
-        key = pkcs11.util.ec.encode_ec_public_key(pub)
+        with self.token.open(user_pin=self.user_pin) as session:
+            pub = session.get_key(
+                id=self.key_id,
+                key_type=pkcs11.KeyType.EC,
+                object_class=pkcs11.ObjectClass.PUBLIC_KEY
+            )
+            key = oscrypto.asymmetric.load_public_key(
+                pkcs11.util.ec.encode_ec_public_key(pub)
+            )
+            key = pkcs11.util.ec.encode_ec_public_key(pub)
         return key
 
     def get_private_bytes(self, minimal):
@@ -48,11 +113,16 @@ def export_private(self, path, passwd=None):
 
     def export_public(self, path):
         """Write the public key to the given file."""
-        pub = self.session.get_key(
-            id=self.key_id, key_type=pkcs11.KeyType.EC, object_class=pkcs11.ObjectClass.PUBLIC_KEY)
-        key = oscrypto.asymmetric.load_public_key(
-            pkcs11.util.ec.encode_ec_public_key(pub))
-        pem = oscrypto.asymmetric.dump_public_key(key, encoding='pem')
+        with self.token.open(user_pin=self.user_pin) as session:
+            pub = session.get_key(
+                id=self.key_id,
+                key_type=pkcs11.KeyType.EC,
+                object_class=pkcs11.ObjectClass.PUBLIC_KEY
+            )
+            key = oscrypto.asymmetric.load_public_key(
+                pkcs11.util.ec.encode_ec_public_key(pub)
+            )
+            pem = oscrypto.asymmetric.dump_public_key(key, encoding='pem')
 
         with open(path, 'wb') as f:
             f.write(pem)
@@ -74,10 +144,16 @@ def sig_len(self):
 
     def raw_sign(self, payload):
         """Return the actual signature"""
-        priv = self.session.get_key(
-            id=self.key_id, key_type=pkcs11.KeyType.EC, object_class=pkcs11.ObjectClass.PRIVATE_KEY)
-        sig = priv.sign(hashlib.sha256(payload).digest(),
-                        mechanism=pkcs11.Mechanism.ECDSA)
+        with self.token.open(user_pin=self.user_pin) as session:
+            priv = session.get_key(
+                id=self.key_id,
+                key_type=pkcs11.KeyType.EC,
+                object_class=pkcs11.ObjectClass.PRIVATE_KEY
+            )
+            sig = priv.sign(
+                hashlib.sha256(payload).digest(),
+                mechanism=pkcs11.Mechanism.ECDSA
+            )
         return pkcs11.util.ec.encode_ecdsa_signature(sig)
 
     def sign(self, payload):