From f867b43251ff5ef7e6b3bdcd7b652be7db1a3d65 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Sun, 11 Aug 2019 20:52:19 +0200
Subject: [PATCH] trezorctl: Add sd-protect command.

---
 python/docs/OPTIONS.rst               |  3 ++-
 python/src/trezorlib/cli/trezorctl.py | 32 ++++++++++++++++++++++++++-
 python/src/trezorlib/device.py        |  6 +++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/python/docs/OPTIONS.rst b/python/docs/OPTIONS.rst
index 54f2a388a7..cb3a4dfd19 100644
--- a/python/docs/OPTIONS.rst
+++ b/python/docs/OPTIONS.rst
@@ -24,7 +24,7 @@ Use the following command to see all options:
     cardano-get-address             Get Cardano address.
     cardano-get-public-key          Get Cardano public key.
     cardano-sign-tx                 Sign Cardano transaction.
-    change-pin                      Change new PIN or remove existing.
+    change-pin                      Set, change or remove PIN.
     clear-session                   Clear session (remove cached PIN, passphrase, etc.).
     cosi-commit                     Ask device to commit to CoSi signing.
     cosi-sign                       Ask device to sign using CoSi.
@@ -66,6 +66,7 @@ Use the following command to see all options:
     reset-device                    Perform device setup and generate new seed.
     ripple-get-address              Get Ripple address
     ripple-sign-tx                  Sign Ripple transaction
+    sd-protect                      Secure the device with SD card protection.
     self-test                       Perform a self-test.
     set-auto-lock-delay             Set auto-lock delay (in seconds).
     set-flags                       Set device flags.
diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py
index a78893d3e8..c23cd1582a 100755
--- a/python/src/trezorlib/cli/trezorctl.py
+++ b/python/src/trezorlib/cli/trezorctl.py
@@ -123,6 +123,14 @@ CHOICE_RESET_DEVICE_TYPE = ChoiceType(
     }
 )
 
+CHOICE_SD_PROTECT_OPERATION_TYPE = ChoiceType(
+    {
+        "enable": proto.SdProtectOperationType.ENABLE,
+        "disable": proto.SdProtectOperationType.DISABLE,
+        "refresh": proto.SdProtectOperationType.REFRESH,
+    }
+)
+
 
 class UnderscoreAgnosticGroup(click.Group):
     """Command group that normalizes dashes and underscores.
@@ -261,13 +269,35 @@ def get_features(connect):
 #
 
 
-@cli.command(help="Change new PIN or remove existing.")
+@cli.command(help="Set, change or remove PIN.")
 @click.option("-r", "--remove", is_flag=True)
 @click.pass_obj
 def change_pin(connect, remove):
     return device.change_pin(connect(), remove)
 
 
+@cli.command()
+@click.argument("operation", type=CHOICE_SD_PROTECT_OPERATION_TYPE)
+@click.pass_obj
+def sd_protect(connect, operation):
+    """Secure the device with SD card protection.
+
+    When SD card protection is enabled, a randomly generated secret is stored
+    on the SD card. During every PIN checking and unlocking operation this
+    secret is combined with the entered PIN value to decrypt data stored on
+    the device. The SD card will thus be needed every time you unlock the
+    device. The options are:
+
+    \b
+    enable - Generate SD card secret and use it to protect the PIN and storage.
+    disable - Remove SD card secret protection.
+    refresh - Replace the current SD card secret with a new one.
+    """
+    if connect().features.model == "1":
+        raise click.BadUsage("Trezor One does not support SD card protection.")
+    return device.sd_protect(connect(), operation)
+
+
 @cli.command(help="Enable passphrase.")
 @click.pass_obj
 def enable_passphrase(connect):
diff --git a/python/src/trezorlib/device.py b/python/src/trezorlib/device.py
index 0e4355bb0f..68cddfada5 100644
--- a/python/src/trezorlib/device.py
+++ b/python/src/trezorlib/device.py
@@ -90,6 +90,12 @@ def change_pin(client, remove=False):
     return ret
 
 
+@expect(proto.Success, field="message")
+def sd_protect(client, operation):
+    ret = client.call(proto.SdProtect(operation=operation))
+    return ret
+
+
 @expect(proto.Success, field="message")
 def set_u2f_counter(client, u2f_counter):
     ret = client.call(proto.SetU2FCounter(u2f_counter=u2f_counter))