From 85bbda02217835cd4aad805078d8bac189512ea2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sosth=C3=A8ne=20Gu=C3=A9don?= <sosthene@nitrokey.com>
Date: Fri, 1 Sep 2023 11:41:11 +0200
Subject: [PATCH] Add support for factory reset command

---
 pynitrokey/cli/nk3/__init__.py |  7 +++++++
 pynitrokey/nk3/admin_app.py    | 38 ++++++++++++++++++++++++++++++++++
 pynitrokey/nk3/device.py       |  3 +++
 3 files changed, 48 insertions(+)

diff --git a/pynitrokey/cli/nk3/__init__.py b/pynitrokey/cli/nk3/__init__.py
index 67120599..884e6637 100644
--- a/pynitrokey/cli/nk3/__init__.py
+++ b/pynitrokey/cli/nk3/__init__.py
@@ -503,6 +503,13 @@ def version(ctx: Context) -> None:
         local_print(version)
 
 
+@nk3.command()
+@click.pass_obj
+def factory_reset(ctx: Context) -> None:
+    with ctx.connect_device() as device:
+        device.factory_reset()
+
+
 @nk3.command()
 @click.pass_obj
 def wink(ctx: Context) -> None:
diff --git a/pynitrokey/nk3/admin_app.py b/pynitrokey/nk3/admin_app.py
index d2265978..2a1a87e4 100644
--- a/pynitrokey/nk3/admin_app.py
+++ b/pynitrokey/nk3/admin_app.py
@@ -18,6 +18,7 @@ class AdminCommand(Enum):
     TEST_SE050 = 0x81
     GET_CONFIG = 0x82
     SET_CONFIG = 0x83
+    FACTORY_RESET = 0x84
 
 
 @enum.unique
@@ -57,6 +58,31 @@ class Status:
     variant: Optional[Variant] = None
 
 
+@enum.unique
+class FactoryResetStatus(Enum):
+    SUCCESS = 0
+    NOT_CONFIRMED = 0x01
+    APP_NOT_ALLOWED = 0x02
+    APP_FAILED_PARSE = 0x03
+
+    @classmethod
+    def from_int(cls, i: int) -> Optional["FactoryResetStatus"]:
+        for status in FactoryResetStatus:
+            if status.value == i:
+                return status
+        return None
+
+    @classmethod
+    def check(cls, i: int, msg: str) -> None:
+        status = FactoryResetStatus.from_int(i)
+        if status != FactoryResetStatus.SUCCESS:
+            if status:
+                error = str(status)
+            else:
+                error = f"unknown error {i:x}"
+            raise Exception(f"{msg}: {error}")
+
+
 @enum.unique
 class ConfigStatus(Enum):
     SUCCESS = 0
@@ -148,3 +174,15 @@ def set_config(self, key: str, value: str) -> None:
         reply = self._call(AdminCommand.SET_CONFIG, data=request, response_len=1)
         assert reply
         ConfigStatus.check(reply[0], "Failed to set config value")
+
+    def factory_reset(self) -> None:
+        try:
+            reply = self._call(AdminCommand.FACTORY_RESET, response_len=1)
+            assert reply
+        except OSError as e:
+            if e.errno == 5:
+                self.device.logger.debug("ignoring OSError after reboot", exc_info=e)
+                return
+            else:
+                raise e
+        FactoryResetStatus.check(reply[0], "Failed to factory reset the device")
diff --git a/pynitrokey/nk3/device.py b/pynitrokey/nk3/device.py
index 99886e9e..a7defe12 100644
--- a/pynitrokey/nk3/device.py
+++ b/pynitrokey/nk3/device.py
@@ -114,6 +114,9 @@ def uuid(self) -> Optional[Uuid]:
     def version(self) -> Version:
         return self.admin.version()
 
+    def factory_reset(self) -> None:
+        self.admin.factory_reset()
+
     def wink(self) -> None:
         self.device.wink()