From 1f2db3666b48f792e4a8980cda18c67192af2365 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Fri, 23 Nov 2018 14:47:47 +0100
Subject: [PATCH] ui: smarter ClickUI prompts only once

This also fixes #331 by moving the PIN matrix from trezorctl into the UI
class
---
 trezorctl       |  3 +--
 trezorlib/ui.py | 25 +++++++++++++++++--------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/trezorctl b/trezorctl
index 584ac9588..e3567d41d 100755
--- a/trezorctl
+++ b/trezorctl
@@ -137,7 +137,7 @@ def cli(ctx, path, verbose, is_json):
                 if path is not None:
                     click.echo("Using path: {}".format(path))
                 sys.exit(1)
-        return TrezorClient(transport=device, ui=ui.ClickUI)
+        return TrezorClient(transport=device, ui=ui.ClickUI())
 
     ctx.obj = get_device
 
@@ -231,7 +231,6 @@ def get_features(connect):
 @click.option("-r", "--remove", is_flag=True)
 @click.pass_obj
 def change_pin(connect, remove):
-    click.echo(ui.PIN_MATRIX_DESCRIPTION)
     return device.change_pin(connect(), remove)
 
 
diff --git a/trezorlib/ui.py b/trezorlib/ui.py
index cd042d3a6..1f663de30 100644
--- a/trezorlib/ui.py
+++ b/trezorlib/ui.py
@@ -56,12 +56,18 @@ def prompt(*args, **kwargs):
 
 
 class ClickUI:
-    @staticmethod
-    def button_request(code):
-        echo("Please confirm action on your Trezor device")
-
-    @staticmethod
-    def get_pin(code=None):
+    def __init__(self, always_prompt=False):
+        self.pinmatrix_shown = False
+        self.prompt_shown = False
+        self.always_prompt = always_prompt
+
+    def button_request(self, code):
+        if not self.prompt_shown:
+            echo("Please confirm action on your Trezor device")
+        if not self.always_prompt:
+            self.prompt_shown = True
+
+    def get_pin(self, code=None):
         if code == PIN_CURRENT:
             desc = "current PIN"
         elif code == PIN_NEW:
@@ -70,7 +76,11 @@ class ClickUI:
             desc = "new PIN again"
         else:
             desc = "PIN"
+
+        if not self.pinmatrix_shown:
             echo(PIN_MATRIX_DESCRIPTION)
+            if not self.always_prompt:
+                self.pinmatrix_shown = True
 
         while True:
             try:
@@ -82,8 +92,7 @@ class ClickUI:
             else:
                 return pin
 
-    @staticmethod
-    def get_passphrase():
+    def get_passphrase(self):
         if os.getenv("PASSPHRASE") is not None:
             echo("Passphrase required. Using PASSPHRASE environment variable.")
             return os.getenv("PASSPHRASE")