diff --git a/trezorlib/client.py b/trezorlib/client.py index d3553e489..0a238848a 100644 --- a/trezorlib/client.py +++ b/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): diff --git a/trezorlib/ui.py b/trezorlib/ui.py index 598b91d1b..cd042d3a6 100644 --- a/trezorlib/ui.py +++ b/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