Add update support for the Nitrokey Pro

Adds pro submenu
Adds update support through nkdfu
Allows to activate bootloader
Runs DFU update with Nitrokey Pro
Supports common error scenarios

Note:
	- uses always the 3.6.0 libnitrokey header

pynitrokey/cli/__init__.py

@@ -15,6 +15,7 @@
 
 import pynitrokey
 import pynitrokey.fido2.operations
+from pynitrokey.cli.pro import pro
 from pynitrokey.cli.fido2 import fido2
 from pynitrokey.cli.nethsm import nethsm
 from pynitrokey.cli.start import start
@@ -52,6 +53,7 @@ def nitropy():
 nitropy.add_command(nethsm)
 nitropy.add_command(start)
 nitropy.add_command(storage)
+nitropy.add_command(pro)
 
 
 @click.command()

pynitrokey/cli/pro.py

@@ -0,0 +1,122 @@
+import platform
+import sys
+
+import click
+import nkdfu
+import logging
+import nkdfu.dfu as dfu
+import intelhex as ih
+import usb1
+
+from pynitrokey.helpers import local_print
+from pynitrokey.libnk import NitrokeyPro, DeviceNotFound
+
+print = local_print
+vendor = "20a0:42b4"
+
+@click.group()
+def pro():
+    """Interact with Nitrokey Pro keys, see subcommands."""
+    pass
+
+
+@click.command()
+def list():
+    """list connected devices"""
+
+    local_print(":: 'Nitrokey Pro' keys:")
+    for dct in NitrokeyPro.list_devices():
+        local_print(dct)
+
+
+@click.command()
+@click.option("-p", "--password", default="12345678",
+              help="update password to be used instead of default")
+def enable_update(password):
+    """enable firmware update for NK Pro device"""
+
+    local_print("Enabling firmware update mode")
+    nks = NitrokeyPro()
+    libnk_version_current = nks.library_version()
+    libnk_version_required = (3, 6)
+    if not libnk_version_current > libnk_version_required:
+        local_print(f'You need libnitrokey {libnk_version_required} to run this command. Currently installed: {libnk_version_current}.')
+        return 1
+
+    try:
+        nks.connect()
+        if nks.enable_firmware_update(password) == 0:
+            local_print("Setting firmware update mode - success!")
+        local_print("Done")
+    except DeviceNotFound:
+        local_print(f"No {nks.friendly_name} device found")
+        local_print("If connected, perhaps already in update mode?")
+
+
+@click.command()
+@click.argument("firmware_path")
+def update(firmware_path: str):
+    """
+    Run firmware update with the provided binary.
+
+    FIRMWARE_PATH: A path to the firmware file. File name should end with .bin.
+    """
+    print = local_print
+    # TODO(szszsz): extract logic to nkdfu, leaving only end-user error handling
+    assert firmware_path.endswith('bin')
+    vendor = "20a0:42b4"
+    product = None
+    if vendor is not None:
+        if ':' in vendor:
+            vendor, product = vendor.split(':')
+            product = int(product, 16)
+        vendor = int(vendor, 16)
+    dev = None
+    bus = None
+    with usb1.USBContext() as context:
+        for device in context.getDeviceList():
+            if (vendor is not None and (vendor != device.getVendorID() or \
+                                        (product is not None and product != device.getProductID()))) \
+                    or (bus is not None and (bus != device.getBusNumber() or \
+                                             (dev is not None and dev != device.getDeviceAddress()))):
+                continue
+            break
+        else:
+            print('No Nitrokey Pro found in the update mode.')
+            print('If you have Nitrokey Pro connected please run (requires libnitrokey):')
+            print('$ nitropy pro enable-update')
+            sys.exit(1)
+        dfu_device = None
+        try:
+            dfu_device = dfu.DFU(device.open())
+        except usb1.USBErrorAccess as e:
+            print(f'Cannot connect to the device: {device} -> {e}')
+            if 'LIBUSB_ERROR_ACCESS' in str(e) and platform.system().lower() == 'linux':
+                print('Try to install UDEV rules, e.g. by executing the following:')  # TODO add command for that
+                print('$ curl https://raw.githubusercontent.com/Nitrokey/libnitrokey/master/data/41-nitrokey.rules | sudo tee /usr/lib/udev/rules.d/41-nitrokey.rules')
+                print('$ sudo udevadm control --reload-rules; sudo udevadm trigger')
+            sys.exit(1)
+        except Exception as e:
+            print(f'Cannot connect to the device: {device} -> {e}')
+            sys.exit(1)
+
+        print(f'Using firmware file {firmware_path}')
+        hex_firmware = ih.IntelHex()
+        hex_firmware.fromfile(open(firmware_path, 'rb'), "bin")
+        data = hex_firmware.tobinarray()
+        try:
+            print((dfu_device.download(data)))
+            print('Please reinsert device to the USB port to complete the process')
+        except nkdfu.dfu.DFUBadSate as e:
+            print(f'Cannot connect to the device: {device} -> {e}')
+            print(f'Reinsert device to the USB port and try again (DFU connects, but reports invalid state)')
+            sys.exit(1)
+        except Exception as e:
+            print(f'Cannot connect to the device: {device} -> {e}')
+            sys.exit(1)
+
+
+pro.add_command(update)
+pro.add_command(list)
+pro.add_command(enable_update)
+

pynitrokey/libnk.py

@@ -19,7 +19,7 @@
 
 SPDX-License-Identifier: LGPL-3.0
 """
-
+import os
 import sys
 from pathlib import Path
 from random import randint
@@ -63,11 +63,20 @@ def _get_c_library():
         + list(Path("/lib").glob("libnitrokey.so.*")) \
         + list(Path("/usr/lib/x86_64-linux-gnu").glob("libnitrokey.so.*"))
 
+    env_path = os.getenv("LIBNK_PATH")
+    if env_path:
+        print(f'Using env variable supplied path: {env_path}')
+        libs += [Path(env_path)]
+        libs += list(Path(env_path).glob("libnitrokey.so*"))
+
     load_lib = None
     load_header = None
     for lib in libs:
         for ver in avail_versions:
-            if ver in lib.as_posix():
+            # TODO load library and check its declared version instead of checking the filename
+            # Fixing on 3.6.0 for the time being
+            ver = '3.6.0'
+            if ver in lib.as_posix() or True:
                 load_lib = lib.as_posix()
                 load_header = (header_parent_path / header.format(ver)).as_posix()
 
@@ -465,6 +474,10 @@ def _connect(self):
         """only connects to NitrokeyPro devices"""
         return self.api.NK_login(b'P')
 
+    def enable_firmware_update(self, password):
+        """set nk storage device to firmware update"""
+        return self.api.NK_enable_firmware_update_pro(c_enc(password))
+
 
 class BaseSlots:
     def __init__(self, parent):

pyproject.toml

@@ -24,6 +24,7 @@ requires = [
   "urllib3",
 	"cffi",
 	"cbor",
+	"nkdfu"
 ]
 classifiers=[
     "License :: OSI Approved :: MIT License",