From 601a2b13366faaefeaf4fee4a4ec5d75187edfd6 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Wed, 7 Mar 2018 12:10:58 +0200
Subject: [PATCH] device: refactor PIN/passphrase UI into a separate class

This would allow easier customization.
---
 libagent/device/__init__.py    |  2 +-
 libagent/device/trezor.py      | 31 ++++++----------
 libagent/device/ui/__init__.py | 65 ++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 22 deletions(-)

diff --git a/libagent/device/__init__.py b/libagent/device/__init__.py
index eb320c4..4698ecd 100644
--- a/libagent/device/__init__.py
+++ b/libagent/device/__init__.py
@@ -1,3 +1,3 @@
 """Cryptographic hardware device management."""
 
-from . import interface
+from . import interface, ui
diff --git a/libagent/device/trezor.py b/libagent/device/trezor.py
index f41dafe..64c777b 100644
--- a/libagent/device/trezor.py
+++ b/libagent/device/trezor.py
@@ -2,24 +2,15 @@
 
 import binascii
 import logging
-import os
-import subprocess
-import sys
 
 import mnemonic
 import semver
 
 from . import interface
-from .ui import pinentry
 
 log = logging.getLogger(__name__)
 
 
-def _message_box(label):
-    """Launch an external process for PIN/passphrase entry GUI."""
-    return pinentry.interact(label.encode('ascii'))
-
-
 class Trezor(interface.Device):
     """Connection to TREZOR device."""
 
@@ -35,18 +26,15 @@ class Trezor(interface.Device):
 
     required_version = '>=1.4.0'
 
+    ui = None  # can be overridden by device's users
+
     def _override_pin_handler(self, conn):
-        cli_handler = conn.callback_PinMatrixRequest
+        if self.ui is None:
+            return
 
-        def new_handler(msg):
+        def new_handler(_):
             try:
-                scrambled_pin = _message_box(
-                    'Use the numeric keypad to describe number positions.\n'
-                    'The layout is:\n'
-                    '    7 8 9\n'
-                    '    4 5 6\n'
-                    '    1 2 3\n'
-                    'Please enter PIN:')
+                scrambled_pin = self.ui.get_pin()
                 result = self._defs.PinMatrixAck(pin=scrambled_pin)
                 if not set(scrambled_pin).issubset('123456789'):
                     raise self._defs.PinException(
@@ -61,15 +49,16 @@ class Trezor(interface.Device):
     cached_passphrase_ack = None
 
     def _override_passphrase_handler(self, conn):
-        cli_handler = conn.callback_PassphraseRequest
+        if self.ui is None:
+            return
 
-        def new_handler(msg):
+        def new_handler(_):
             try:
                 if self.__class__.cached_passphrase_ack:
                     log.debug('re-using cached %s passphrase', self)
                     return self.__class__.cached_passphrase_ack
 
-                passphrase = _message_box('Please enter passphrase:')
+                passphrase = self.ui.get_passphrase()
                 passphrase = mnemonic.Mnemonic.normalize_string(passphrase)
                 ack = self._defs.PassphraseAck(passphrase=passphrase)
 
diff --git a/libagent/device/ui/__init__.py b/libagent/device/ui/__init__.py
index 4313f96..9279d8e 100644
--- a/libagent/device/ui/__init__.py
+++ b/libagent/device/ui/__init__.py
@@ -1 +1,66 @@
 """UIs for PIN/passphrase entry."""
+
+import logging
+import os
+import subprocess
+
+from . import pinentry
+
+log = logging.getLogger(__name__)
+
+
+def _create_default_options_getter():
+    options = []
+    try:
+        ttyname = subprocess.check_output(args=['tty']).strip()
+        options.append(b'ttyname=' + ttyname)
+    except subprocess.CalledProcessError as e:
+        log.warning('no TTY found: %s', e)
+
+    display = os.environ.get('DISPLAY')
+    if display is not None:
+        options.append('display={}'.format(display).encode('ascii'))
+    else:
+        log.warning('DISPLAY not defined')
+
+    log.info('using %s for pinentry options', options)
+    return lambda: options
+
+
+class UI(object):
+    """UI for PIN/passphrase entry (for TREZOR devices)."""
+
+    def __init__(self):
+        """C-tor."""
+        self.options_getter = _create_default_options_getter()
+        self.pin_entry_binary = 'pinentry'
+        self.passphrase_entry_binary = 'pinentry'
+
+    @classmethod
+    def from_config_dict(cls, d):
+        """Simple c-tor from configuration dictionary."""
+        obj = cls()
+        obj.pin_entry_binary = d.get('pin_entry_binary',
+                                     obj.pin_entry_binary)
+        obj.passphrase_entry_binary = d.get('passphrase_entry_binary',
+                                            obj.passphrase_entry_binary)
+        return obj
+
+    def get_pin(self):
+        """Ask the user for (scrambled) PIN."""
+        return pinentry.interact(
+            'Use the numeric keypad to describe number positions.\n'
+            'The layout is:\n'
+            '    7 8 9\n'
+            '    4 5 6\n'
+            '    1 2 3\n'
+            'Please enter PIN:',
+            binary=self.pin_entry_binary,
+            options=self.options_getter())
+
+    def get_passphrase(self):
+        """Ask the user for passphrase."""
+        return pinentry.interact(
+            'Please enter passphrase:',
+            binary=self.passphrase_entry_binary,
+            options=self.options_getter())