From 83fab3c220ced8eb64155e0c64bd62a4d69bf054 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Fri, 18 Oct 2019 15:59:12 +0200
Subject: [PATCH] trezorctl: Add set wipe-code command.

---
 python/docs/OPTIONS.rst              |  1 +
 python/src/trezorlib/cli/settings.py | 13 +++++++++++++
 python/src/trezorlib/device.py       |  7 +++++++
 python/src/trezorlib/ui.py           |  8 +++++++-
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/python/docs/OPTIONS.rst b/python/docs/OPTIONS.rst
index 67e44f44c5..85f3a7dd8a 100644
--- a/python/docs/OPTIONS.rst
+++ b/python/docs/OPTIONS.rst
@@ -362,6 +362,7 @@ Device settings.
     label             Set new device label.
     passphrase        Enable, disable or configure passphrase protection.
     pin               Set, change or remove PIN.
+    wipe-code         Set or remove the wipe code.
 
 Stellar commands.
 ~~~~~~~~~~~~~~~~~
diff --git a/python/src/trezorlib/cli/settings.py b/python/src/trezorlib/cli/settings.py
index bdc5d598bd..2bf72cc389 100644
--- a/python/src/trezorlib/cli/settings.py
+++ b/python/src/trezorlib/cli/settings.py
@@ -41,6 +41,19 @@ def pin(connect, remove):
     return device.change_pin(connect(), remove)
 
 
+@cli.command()
+@click.option("-r", "--remove", is_flag=True)
+@click.pass_obj
+def wipe_code(connect, remove):
+    """Set or remove the wipe code.
+
+    The wipe code functions as a "self-destruct PIN". If the wipe code is ever
+    entered into any PIN entry dialog, then all private data will be immediately
+    removed and the device will be reset to factory defaults.
+    """
+    return device.change_wipe_code(connect(), remove)
+
+
 @cli.command()
 # keep the deprecated -l/--label option, make it do nothing
 @click.option("-l", "--label", "_ignore", is_flag=True, hidden=True, expose_value=False)
diff --git a/python/src/trezorlib/device.py b/python/src/trezorlib/device.py
index 51decd4e3d..f95836137f 100644
--- a/python/src/trezorlib/device.py
+++ b/python/src/trezorlib/device.py
@@ -90,6 +90,13 @@ def change_pin(client, remove=False):
     return ret
 
 
+@expect(messages.Success, field="message")
+def change_wipe_code(client, remove=False):
+    ret = client.call(messages.ChangeWipeCode(remove=remove))
+    client.init_device()  # Re-read features
+    return ret
+
+
 @expect(messages.Success, field="message")
 def sd_protect(client, operation):
     ret = client.call(messages.SdProtect(operation=operation))
diff --git a/python/src/trezorlib/ui.py b/python/src/trezorlib/ui.py
index e95eb50e39..fc0f6e20f4 100644
--- a/python/src/trezorlib/ui.py
+++ b/python/src/trezorlib/ui.py
@@ -45,6 +45,8 @@ PIN_GENERIC = None
 PIN_CURRENT = PinMatrixRequestType.Current
 PIN_NEW = PinMatrixRequestType.NewFirst
 PIN_CONFIRM = PinMatrixRequestType.NewSecond
+WIPE_CODE_NEW = PinMatrixRequestType.WipeCodeFirst
+WIPE_CODE_CONFIRM = PinMatrixRequestType.WipeCodeSecond
 
 
 def echo(*args, **kwargs):
@@ -74,6 +76,10 @@ class ClickUI:
             desc = "new PIN"
         elif code == PIN_CONFIRM:
             desc = "new PIN again"
+        elif code == WIPE_CODE_NEW:
+            desc = "new wipe code"
+        elif code == WIPE_CODE_CONFIRM:
+            desc = "new wipe code again"
         else:
             desc = "PIN"
 
@@ -88,7 +94,7 @@ class ClickUI:
             except click.Abort:
                 raise Cancelled from None
             if not pin.isdigit():
-                echo("Non-numerical PIN provided, please try again")
+                echo("Non-numerical value provided, please try again")
             else:
                 return pin