From 6c9c563323e160d8b4c736423460ef3295a20cc1 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Tue, 20 Mar 2018 16:44:23 +0100
Subject: [PATCH] pin: simplify label code

---
 src/apps/common/request_pin.py    | 22 ++++------------------
 src/apps/management/change_pin.py | 30 +++++++++++++-----------------
 src/trezor/ui/pin.py              | 11 ++++-------
 3 files changed, 21 insertions(+), 42 deletions(-)

diff --git a/src/apps/common/request_pin.py b/src/apps/common/request_pin.py
index dc5532059..dfd8badd1 100644
--- a/src/apps/common/request_pin.py
+++ b/src/apps/common/request_pin.py
@@ -1,5 +1,4 @@
 from trezor import loop, res, ui
-from trezor.messages import PinMatrixRequestType
 from trezor.ui.confirm import CONFIRMED, ConfirmDialog
 from trezor.ui.pin import PinMatrix
 
@@ -12,7 +11,7 @@ class PinCancelled(Exception):
 
 
 @ui.layout
-async def request_pin(code=None, cancellable: bool=True) -> str:
+async def request_pin(label=None, cancellable: bool=True) -> str:
 
     def onchange():
         c = dialog.cancel
@@ -36,8 +35,9 @@ def onchange():
                 c.taint()
         c.render()
 
-    label = _get_label(code)
-    matrix = PinMatrix(label, with_zero=True)
+    if label is None:
+        label = 'Enter your PIN'
+    matrix = PinMatrix(label)
     matrix.onchange = onchange
     dialog = ConfirmDialog(matrix)
     dialog.cancel.area = ui.grid(12)
@@ -58,17 +58,3 @@ def onchange():
             continue
         else:  # cancel
             raise PinCancelled()
-
-
-def _get_label(code):
-    if isinstance(code, str):
-        return code
-    if code is None:
-        code = PinMatrixRequestType.Current
-    if code == PinMatrixRequestType.NewFirst:
-        label = 'Enter new PIN'
-    elif code == PinMatrixRequestType.NewSecond:
-        label = 'Re-enter new PIN'
-    else:  # PinMatrixRequestType.Current
-        label = 'Enter your PIN'
-    return label
diff --git a/src/apps/management/change_pin.py b/src/apps/management/change_pin.py
index 12a5281e4..a921f8c40 100644
--- a/src/apps/management/change_pin.py
+++ b/src/apps/management/change_pin.py
@@ -1,14 +1,13 @@
 from trezor import config, loop, ui, wire
-from trezor.messages import FailureType, PinMatrixRequestType
+from trezor.messages import FailureType, wire_types
 from trezor.messages.ButtonRequest import ButtonRequest
 from trezor.messages.ButtonRequestType import Other
 from trezor.messages.Failure import Failure
 from trezor.messages.Success import Success
-from trezor.messages import wire_types
 from trezor.pin import pin_to_int, show_pin_timeout
 from trezor.ui.text import Text
 from apps.common.confirm import require_confirm
-from apps.common.request_pin import request_pin, PinCancelled
+from apps.common.request_pin import PinCancelled, request_pin
 
 
 async def change_pin(ctx, msg):
@@ -18,7 +17,7 @@ async def change_pin(ctx, msg):
 
     # get current pin, return failure if invalid
     if config.has_pin():
-        curpin = await request_pin_ack(ctx, PinMatrixRequestType.Current)
+        curpin = await request_pin_ack(ctx)
         if not config.check_pin(pin_to_int(curpin), show_pin_timeout):
             return Failure(code=FailureType.PinInvalid, message='PIN invalid')
     else:
@@ -62,26 +61,23 @@ def require_confirm_change_pin(ctx, msg):
             'set new PIN?'))
 
 
-async def request_pin_ack(ctx, *args, **kwargs):
-    # TODO: send PinMatrixRequest here, with specific code?
-    await ctx.call(ButtonRequest(code=Other), wire_types.ButtonAck)
-    try:
-        return await ctx.wait(request_pin(*args, **kwargs))
-    except PinCancelled:
-        raise wire.FailureError(FailureType.ActionCancelled, 'Cancelled')
-
-
 async def request_pin_confirm(ctx, *args, **kwargs):
     while True:
-        pin1 = await request_pin_ack(
-            ctx, code=PinMatrixRequestType.NewFirst, *args, **kwargs)
-        pin2 = await request_pin_ack(
-            ctx, code=PinMatrixRequestType.NewSecond, *args, **kwargs)
+        pin1 = await request_pin_ack(ctx, 'Enter new PIN', *args, **kwargs)
+        pin2 = await request_pin_ack(ctx, 'Re-enter new PIN', *args, **kwargs)
         if pin1 == pin2:
             return pin1
         await pin_mismatch()
 
 
+async def request_pin_ack(ctx, *args, **kwargs):
+    try:
+        await ctx.call(ButtonRequest(code=Other), wire_types.ButtonAck)
+        return await ctx.wait(request_pin(*args, **kwargs))
+    except PinCancelled:
+        raise wire.FailureError(FailureType.ActionCancelled, 'Cancelled')
+
+
 @ui.layout
 async def pin_mismatch():
     text = Text(
diff --git a/src/trezor/ui/pin.py b/src/trezor/ui/pin.py
index eb357bf70..37da9eaa8 100644
--- a/src/trezor/ui/pin.py
+++ b/src/trezor/ui/pin.py
@@ -11,22 +11,19 @@ def digit_area(i):
     return ui.grid(i + 3)  # skip the first line
 
 
-def generate_digits(with_zero):
-    if with_zero:
-        digits = list(range(0, 10))  # 0-9
-    else:
-        digits = list(range(1, 10))  # 1-9
+def generate_digits():
+    digits = list(range(0, 10))  # 0-9
     random.shuffle(digits)
     return digits
 
 
 class PinMatrix(ui.Widget):
 
-    def __init__(self, label, pin='', maxlength=9, with_zero=False):
+    def __init__(self, label, pin='', maxlength=9):
         self.label = label
         self.pin = pin
         self.maxlength = maxlength
-        self.digits = generate_digits(with_zero)
+        self.digits = generate_digits()
 
         # we lay out the buttons top-left to bottom-right, but the order of the
         # digits is defined as bottom-left to top-right (on numpad)