From 7443fc651211a6b7ccb6be6bc62ecb6f11a85a6e Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Tue, 27 Mar 2018 15:58:42 +0300
Subject: [PATCH] Pass 'state' during TREZOR initialization

---
 libagent/device/keepkey.py     |  3 +++
 libagent/device/trezor.py      | 15 +++++++++++++--
 libagent/device/trezor_defs.py |  2 +-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/libagent/device/keepkey.py b/libagent/device/keepkey.py
index 511d4f5..3a086e5 100644
--- a/libagent/device/keepkey.py
+++ b/libagent/device/keepkey.py
@@ -32,6 +32,9 @@ class KeepKey(trezor.Trezor):
 
     required_version = '>=1.0.4'
 
+    def _override_state_handler(self, _):
+        """No support for `state` handling on Keepkey."""
+
     def pubkey(self, identity, ecdh=False):
         """Return public key."""
         _verify_support(identity, ecdh)
diff --git a/libagent/device/trezor.py b/libagent/device/trezor.py
index 1d39b7b..0fbfc2b 100644
--- a/libagent/device/trezor.py
+++ b/libagent/device/trezor.py
@@ -47,6 +47,7 @@ class Trezor(interface.Device):
         conn.callback_PinMatrixRequest = new_handler
 
     cached_passphrase_ack = None
+    cached_state = None
 
     def _override_passphrase_handler(self, conn):
         if self.ui is None:
@@ -77,6 +78,14 @@ class Trezor(interface.Device):
 
         conn.callback_PassphraseRequest = new_handler
 
+    def _override_state_handler(self, conn):
+        def callback_PassphraseStateRequest(msg):
+            log.debug('caching state from %r', msg)
+            self.__class__.cached_state = msg.state
+            return self._defs.PassphraseStateAck()
+
+        conn.callback_PassphraseStateRequest = callback_PassphraseStateRequest
+
     def _verify_version(self, connection):
         f = connection.features
         log.debug('connected to %s %s', self, f.device_id)
@@ -100,10 +109,12 @@ class Trezor(interface.Device):
             raise interface.NotFoundError('{} not connected'.format(self))
 
         log.debug('transports: %s', transports)
-        for _ in range(5):
-            connection = self._defs.Client(transports[0])
+        for _ in range(5):  # Retry a few times in case of PIN failures
+            connection = self._defs.Client(transport=transports[0],
+                                           state=self.__class__.cached_state)
             self._override_pin_handler(connection)
             self._override_passphrase_handler(connection)
+            self._override_state_handler(connection)
             self._verify_version(connection)
 
             try:
diff --git a/libagent/device/trezor_defs.py b/libagent/device/trezor_defs.py
index a4daf81..49662d3 100644
--- a/libagent/device/trezor_defs.py
+++ b/libagent/device/trezor_defs.py
@@ -4,7 +4,7 @@
 
 from trezorlib.client import CallException, PinException
 from trezorlib.client import TrezorClient as Client
-from trezorlib.messages import IdentityType, PassphraseAck, PinMatrixAck
+from trezorlib.messages import IdentityType, PassphraseAck, PinMatrixAck, PassphraseStateAck
 from trezorlib.device import TrezorDevice