client: allow canceling pin/passphrase entry

trezorlib/client.py

@@ -102,7 +102,12 @@ class TrezorClient:
         return self.transport.read()
 
     def _callback_pin(self, msg):
-        pin = self.ui.get_pin(msg.type)
+        try:
+            pin = self.ui.get_pin(msg.type)
+        except exceptions.Cancelled:
+            self.call_raw(messages.Cancel())
+            raise
+
         if not pin.isdigit():
             raise ValueError("Non-numeric PIN provided")
 
@@ -120,7 +125,11 @@ class TrezorClient:
         if msg.on_device:
             passphrase = None
         else:
-            passphrase = self.ui.get_passphrase()
+            try:
+                passphrase = self.ui.get_passphrase()
+            except exceptions.Cancelled:
+                self.call_raw(messages.Cancel())
+                raise
 
         resp = self.call_raw(messages.PassphraseAck(passphrase=passphrase))
         if isinstance(resp, messages.PassphraseStateRequest):

trezorlib/ui.py

@@ -73,7 +73,10 @@ class ClickUI:
             echo(PIN_MATRIX_DESCRIPTION)
 
         while True:
-            pin = prompt("Please enter {}".format(desc), hide_input=True)
+            try:
+                pin = prompt("Please enter {}".format(desc), hide_input=True)
+            except click.Abort:
+                raise Cancelled from None
             if not pin.isdigit():
                 echo("Non-numerical PIN provided, please try again")
             else:
@@ -86,12 +89,15 @@ class ClickUI:
             return os.getenv("PASSPHRASE")
 
         while True:
-            passphrase = prompt("Passphrase required", hide_input=True)
-            second = prompt("Confirm your passphrase", hide_input=True)
-            if passphrase == second:
-                return passphrase
-            else:
-                echo("Passphrase did not match. Please try again.")
+            try:
+                passphrase = prompt("Passphrase required", hide_input=True)
+                second = prompt("Confirm your passphrase", hide_input=True)
+                if passphrase == second:
+                    return passphrase
+                else:
+                    echo("Passphrase did not match. Please try again.")
+            except click.Abort:
+                raise Cancelled from None
 
 
 def mnemonic_words(expand=False, language="english"):
@@ -119,7 +125,7 @@ def mnemonic_words(expand=False, language="english"):
                 return expand_word(word)
             except KeyError:
                 pass
-            except (KeyboardInterrupt, click.Abort):
+            except click.Abort:
                 raise Cancelled from None
 
     return get_word