Merge pull request Nitrokey#106 from robin-nitrokey/nk3

Add basic support for managing Nitrokey 3 devices

Makefile

@@ -9,7 +9,7 @@ ISORT_FLAGS=--py 35 --extend-skip pynitrokey/nethsm/client
 MYPY_FLAGS=--config-file mypy.ini
 
 # whitelist of directories for flake8
-FLAKE8_DIRS=pynitrokey/nethsm
+FLAKE8_DIRS=pynitrokey/nethsm pynitrokey/cli/nk3.py pynitrokey/nk3
 
 # setup development environment
 init: update-venv

mypy.ini

@@ -1,9 +1,13 @@
 [mypy]
 
+# enable strict checks for new code
+[mypy-pynitrokey.cli.nk3,pynitrokey.nk3.*]
+disallow_untyped_defs = True
+
 # pynitrokey.nethsm.client is auto-generated
 [mypy-pynitrokey.nethsm.client.*]
 ignore_errors = True
 
 # libraries without annotations
-[mypy-cbor.*,cffi.*,ecdsa.*,fido2.*,intelhex.*,nacl.*,nkdfu.*,serial.*,urllib3.*,usb.*,usb1.*]
+[mypy-cbor.*,cffi.*,click.*,ecdsa.*,fido2.*,intelhex.*,nacl.*,nkdfu.*,serial.*,urllib3.*,usb.*,usb1.*]
 ignore_missing_imports = True

pynitrokey/cli/__init__.py

@@ -17,6 +17,7 @@
 import pynitrokey.fido2.operations
 from pynitrokey.cli.fido2 import fido2
 from pynitrokey.cli.nethsm import nethsm
+from pynitrokey.cli.nk3 import nk3
 from pynitrokey.cli.pro import pro
 from pynitrokey.cli.start import start
 from pynitrokey.cli.storage import storage
@@ -47,12 +48,16 @@ def nitropy():
     handler = logging.FileHandler(filename=LOG_FN, delay=True)
     logging.basicConfig(format=LOG_FORMAT, level=logging.DEBUG, handlers=[handler])
 
-    print("Nitrokey tool for Nitrokey FIDO2, Nitrokey Start & NetHSM", file=sys.stderr)
+    print(
+        "Nitrokey tool for Nitrokey FIDO2, Nitrokey Start, Nitrokey 3 & NetHSM",
+        file=sys.stderr,
+    )
     check_root()
 
 
 nitropy.add_command(fido2)
 nitropy.add_command(nethsm)
+nitropy.add_command(nk3)
 nitropy.add_command(start)
 nitropy.add_command(storage)
 nitropy.add_command(pro)
@@ -73,6 +78,7 @@ def ls():
 
     fido2.commands["list"].callback()
     start.commands["list"].callback()
+    nk3.commands["list"].callback()
 
 
 nitropy.add_command(ls)

pynitrokey/cli/nk3.py

@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021 Nitrokey Developers
+#
+# Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
+# http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
+# http://opensource.org/licenses/MIT>, at your option. This file may not be
+# copied, modified, or distributed except according to those terms.
+
+from typing import List, Optional, TypeVar
+
+import click
+
+from pynitrokey.helpers import local_print
+from pynitrokey.nk3 import list as list_nk3
+from pynitrokey.nk3.base import Nitrokey3Base
+from pynitrokey.nk3.device import Nitrokey3Device
+
+T = TypeVar("T", bound="Nitrokey3Base")
+
+
+class Context:
+    def __init__(self, path: Optional[str]) -> None:
+        self.path = path
+
+    def list(self) -> List[Nitrokey3Base]:
+        devices = []
+        for device in list_nk3():
+            if not self.path or self.path == device.path:
+                devices.append(device)
+        return devices
+
+    def _select_unique(self, name: str, devices: List[T]) -> T:
+        if len(devices) == 0:
+            msg = f"No {name} device found"
+            if self.path:
+                msg += f" at path {self.path}"
+            raise click.ClickException(msg)
+
+        if len(devices) > 1:
+            raise click.ClickException(
+                f"Multiple {name} devices found -- use the --path option to select one"
+            )
+
+        return devices[0]
+
+    def connect(self) -> Nitrokey3Base:
+        return self._select_unique("Nitrokey 3", self.list())
+
+    def connect_device(self) -> Nitrokey3Device:
+        devices = [
+            device for device in self.list() if isinstance(device, Nitrokey3Device)
+        ]
+        return self._select_unique("Nitrokey 3", devices)
+
+
+@click.group()
+@click.option("-p", "--path", "path", help="The path of the Nitrokey 3 device")
+@click.pass_context
+def nk3(ctx: click.Context, path: Optional[str]) -> None:
+    """Interact with Nitrokey 3, see subcommands."""
+    ctx.obj = Context(path)
+
+
+@nk3.command()
+def list() -> None:
+    """List all Nitrokey 3 devices."""
+    local_print(":: 'Nitrokey 3' keys")
+    for device in list_nk3():
+        with device as device:
+            uuid = device.uuid()
+            if uuid:
+                local_print(f"{device.path}: {device.name} {device.uuid():X}")
+            else:
+                local_print(f"{device.path}: {device.name}")
+
+
+@nk3.command()
+@click.pass_obj
+def reboot(ctx: Context) -> None:
+    """Reboot the key."""
+    with ctx.connect() as device:
+        device.reboot()
+
+
+@nk3.command()
+@click.option(
+    "-l",
+    "--length",
+    "length",
+    default=57,
+    help="The length of the generated data (default: 57)",
+)
+@click.pass_obj
+def rng(ctx: Context, length: int) -> None:
+    """Generate random data on the key."""
+    with ctx.connect_device() as device:
+        while length > 0:
+            rng = device.rng()
+            local_print(rng[:length].hex())
+            length -= len(rng)
+
+
+@nk3.command()
+@click.pass_obj
+def version(ctx: Context) -> None:
+    """Query the firmware version of the key."""
+    with ctx.connect_device() as device:
+        version = device.version()
+        local_print(version)
+
+
+@nk3.command()
+@click.pass_obj
+def wink(ctx: Context) -> None:
+    """Send wink command to the key (blinks LED a few times)."""
+    with ctx.connect_device() as device:
+        device.wink()

pynitrokey/nk3/__init__.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021 Nitrokey Developers
+#
+# Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
+# http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
+# http://opensource.org/licenses/MIT>, at your option. This file may not be
+# copied, modified, or distributed except according to those terms.
+
+from typing import List
+
+from .base import Nitrokey3Base
+
+VID_NITROKEY = 0x20A0
+PID_NITROKEY3_DEVICE = 0x42B2
+PID_NITROKEY3_BOOTLOADER = 0x42DD
+
+
+def list() -> List[Nitrokey3Base]:
+    from .device import Nitrokey3Device
+
+    return [device for device in Nitrokey3Device.list()]

pynitrokey/nk3/base.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021 Nitrokey Developers
+#
+# Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
+# http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
+# http://opensource.org/licenses/MIT>, at your option. This file may not be
+# copied, modified, or distributed except according to those terms.
+
+from abc import ABC, abstractmethod
+from typing import Optional, TypeVar
+
+T = TypeVar("T", bound="Nitrokey3Base")
+
+
+class Nitrokey3Base(ABC):
+    """Base class for Nitrokey 3 devices, running the firmware or the bootloader."""
+
+    def __enter__(self: T) -> T:
+        return self
+
+    def __exit__(self, exc_type: None, exc_val: None, exc_tb: None) -> None:
+        self.close()
+
+    @property
+    @abstractmethod
+    def path(self) -> str:
+        ...
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        ...
+
+    @abstractmethod
+    def close(self) -> None:
+        ...
+
+    @abstractmethod
+    def reboot(self) -> None:
+        ...
+
+    @abstractmethod
+    def uuid(self) -> Optional[int]:
+        ...

pynitrokey/nk3/device.py

@@ -0,0 +1,133 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021 Nitrokey Developers
+#
+# Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
+# http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
+# http://opensource.org/licenses/MIT>, at your option. This file may not be
+# copied, modified, or distributed except according to those terms.
+
+import enum
+import errno
+from enum import Enum
+from typing import List, Optional
+
+from fido2.hid import CtapHidDevice
+
+from . import PID_NITROKEY3_DEVICE, VID_NITROKEY
+from .base import Nitrokey3Base
+
+RNG_LEN = 57
+UUID_LEN = 16
+VERSION_LEN = 4
+
+
+@enum.unique
+class Command(Enum):
+    """Vendor-specific CTAPHID commands for the Nitrokey 3."""
+
+    UPDATE = 0x51
+    REBOOT = 0x53
+    RNG = 0x60
+    VERSION = 0x61
+    UUID = 0x62
+
+
+@enum.unique
+class BootMode(Enum):
+    FIRMWARE = enum.auto()
+    BOOTROM = enum.auto()
+
+
+class Version:
+    def __init__(self, major: int, minor: int, patch: int) -> None:
+        self.major = major
+        self.minor = minor
+        self.patch = patch
+
+    def __repr__(self) -> str:
+        return f"Version(major={self.major}, minor={self.minor}, patch={self.patch}"
+
+    def __str__(self) -> str:
+        return f"v{self.major}.{self.minor}.{self.patch}"
+
+
+class Nitrokey3Device(Nitrokey3Base):
+    """A Nitrokey 3 device running the firmware."""
+
+    def __init__(self, device: CtapHidDevice) -> None:
+        (vid, pid) = (device.descriptor.vid, device.descriptor.pid)
+        if (vid, pid) != (VID_NITROKEY, PID_NITROKEY3_DEVICE):
+            raise ValueError(
+                "Not a Nitrokey 3 device: expected VID:PID "
+                f"{VID_NITROKEY:x}:{PID_NITROKEY3_DEVICE:x}, got {vid:x}:{pid:x}"
+            )
+
+        self.device = device
+
+    @property
+    def path(self) -> str:
+        return self.device.descriptor.path
+
+    @property
+    def name(self) -> str:
+        return "Nitrokey 3"
+
+    def close(self) -> None:
+        self.device.close()
+
+    def reboot(self, mode: BootMode = BootMode.FIRMWARE) -> None:
+        try:
+            if mode == BootMode.FIRMWARE:
+                self._call(Command.REBOOT)
+            elif mode == BootMode.BOOTROM:
+                self._call(Command.UPDATE)
+        except OSError as e:
+            if e.errno == errno.EIO:
+                # IO error is expected as the device does not respond during the reboot
+                pass
+            else:
+                raise e
+
+    def uuid(self) -> Optional[int]:
+        uuid = self._call(Command.UUID)
+        if len(uuid) == 0:
+            # Firmware version 1.0.0 does not support querying the UUID
+            return None
+        if len(uuid) != UUID_LEN:
+            raise ValueError(f"UUID response has invalid length {len(uuid)}")
+        return int.from_bytes(uuid, "big")
+
+    def version(self) -> Version:
+        version_bytes = self._call(Command.VERSION, response_len=VERSION_LEN)
+        version = int.from_bytes(version_bytes, "big")
+        major = version >> 22
+        minor = (version >> 6) & ((1 << 16) - 1)
+        patch = version & ((1 << 6) - 1)
+        return Version(major=major, minor=minor, patch=patch)
+
+    def wink(self) -> None:
+        self.device.wink()
+
+    def rng(self) -> bytes:
+        return self._call(Command.RNG, response_len=RNG_LEN)
+
+    def _call(self, command: Command, response_len: Optional[int] = None) -> bytes:
+        response = self.device.call(command.value)
+        if response_len is not None and response_len != len(response):
+            raise ValueError(
+                f"The response for the CTAPHID {command.name} command has an unexpected length "
+                f"(expected: {response_len}, actual: {len(response)})"
+            )
+        return response
+
+    @staticmethod
+    def list() -> List["Nitrokey3Device"]:
+        devices = []
+        for device in CtapHidDevice.list_devices():
+            try:
+                devices.append(Nitrokey3Device(device))
+            except ValueError:
+                # not a Nitrokey 3 device, skip
+                pass
+        return devices