From 729b1d5f70418c6cf3941e42809f7ba39315bbcd Mon Sep 17 00:00:00 2001
From: Robin Krahl <robin@nitrokey.com>
Date: Thu, 3 Feb 2022 11:18:36 +0100
Subject: [PATCH] nk3: Improve error message for missing confirmation

This patch improves the error message shown if the user does not confirm
a reboot operation with a touch button press.

Fixes https://github.com/Nitrokey/pynitrokey/issues/173
---
 pynitrokey/cli/nk3/__init__.py  | 17 +++++++++++++++--
 pynitrokey/nk3/device.py        | 12 +++++++++++-
 pynitrokey/nk3/exceptions.py    | 17 +++++++++++++++++
 pynitrokey/stubs/fido2/ctap.pyi | 11 ++++++++++-
 4 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 pynitrokey/nk3/exceptions.py

diff --git a/pynitrokey/cli/nk3/__init__.py b/pynitrokey/cli/nk3/__init__.py
index 17d235ae..a0afedd9 100644
--- a/pynitrokey/cli/nk3/__init__.py
+++ b/pynitrokey/cli/nk3/__init__.py
@@ -29,6 +29,7 @@
     check_firmware_image,
 )
 from pynitrokey.nk3.device import BootMode, Nitrokey3Device
+from pynitrokey.nk3.exceptions import TimeoutException
 from pynitrokey.nk3.updates import get_latest_update, get_update
 from pynitrokey.nk3.utils import Version
 
@@ -117,7 +118,13 @@ def reboot(ctx: Context, bootloader: bool) -> None:
                 local_print(
                     "Please press the touch button to reboot the device into bootloader mode ..."
                 )
-                device.reboot(BootMode.BOOTROM)
+                try:
+                    device.reboot(BootMode.BOOTROM)
+                except TimeoutException:
+                    local_critical(
+                        "The reboot was not confirmed with the touch button.",
+                        support_hint=False,
+                    )
             else:
                 local_critical(
                     "A Nitrokey 3 device in bootloader mode can only reboot into firmware mode."
@@ -354,7 +361,13 @@ def update(ctx: Context, image: Optional[str], experimental: bool) -> None:
             local_print(
                 "Please press the touch button to reboot the device into bootloader mode ..."
             )
-            device.reboot(BootMode.BOOTROM)
+            try:
+                device.reboot(BootMode.BOOTROM)
+            except TimeoutException:
+                local_critical(
+                    "The reboot was not confirmed with the touch button.",
+                    support_hint=False,
+                )
 
             local_print("")
 
diff --git a/pynitrokey/nk3/device.py b/pynitrokey/nk3/device.py
index 8eebfed3..20e6abf1 100644
--- a/pynitrokey/nk3/device.py
+++ b/pynitrokey/nk3/device.py
@@ -13,11 +13,13 @@
 from enum import Enum
 from typing import List, Optional
 
+from fido2.ctap import CtapError
 from fido2.hid import CtapHidDevice, open_device
 
 from pynitrokey.fido2 import device_path_to_str
 
 from .base import Nitrokey3Base
+from .exceptions import TimeoutException
 from .utils import Version
 
 RNG_LEN = 57
@@ -77,7 +79,15 @@ def reboot(self, mode: BootMode = BootMode.FIRMWARE) -> None:
             if mode == BootMode.FIRMWARE:
                 self._call(Command.REBOOT)
             elif mode == BootMode.BOOTROM:
-                self._call(Command.UPDATE)
+                try:
+                    self._call(Command.UPDATE)
+                except CtapError as e:
+                    # The admin app returns an Invalid Length error if the user confirmation
+                    # request times out
+                    if e.code == CtapError.ERR.INVALID_LENGTH:
+                        raise TimeoutException()
+                    else:
+                        raise e
         except OSError as e:
             # OS error is expected as the device does not respond during the reboot
             self.logger.debug("ignoring OSError after reboot", exc_info=e)
diff --git a/pynitrokey/nk3/exceptions.py b/pynitrokey/nk3/exceptions.py
new file mode 100644
index 00000000..c3abbe55
--- /dev/null
+++ b/pynitrokey/nk3/exceptions.py
@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2022 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.
+
+
+class Nitrokey3Exception(Exception):
+    pass
+
+
+class TimeoutException(Nitrokey3Exception):
+    def __init__(self) -> None:
+        super().__init__("The user confirmation request timed out")
diff --git a/pynitrokey/stubs/fido2/ctap.pyi b/pynitrokey/stubs/fido2/ctap.pyi
index 04739c9c..51cf2cec 100644
--- a/pynitrokey/stubs/fido2/ctap.pyi
+++ b/pynitrokey/stubs/fido2/ctap.pyi
@@ -7,5 +7,14 @@
 # http://opensource.org/licenses/MIT>, at your option. This file may not be
 # copied, modified, or distributed except according to those terms.
 
+from enum import IntEnum, unique
+from typing import Union
+
 class CtapDevice: ...
-class CtapError(Exception): ...
+
+class CtapError(Exception):
+    class UNKNOWN_ERR(int): ...
+    @unique
+    class ERR(IntEnum):
+        INVALID_LENGTH: int
+    code: Union[UNKNOWN_ERR, ERR]