Pass 'state' during TREZOR initialization

romanz · Mar 27, 2018 · 7443fc6 · 7443fc6
1 parent 5efb752
commit 7443fc6
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/libagent/device/keepkey.py b/libagent/device/keepkey.py
@@ -32,6 +32,9 @@ def _defs(self):
 
     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
@@ -47,6 +47,7 @@ def new_handler(_):
         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 @@ def new_handler(msg):
 
         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 @@ def connect(self):
             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
@@ -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