Allow specifying 'state' at ProtocolMixin c-tor

trezor · Mar 27, 2018 · a6d8c21 · a6d8c21
1 parent 8d296ed
commit a6d8c21
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/trezorctl b/trezorctl
@@ -87,7 +87,7 @@ def cli(ctx, path, verbose, is_json):
                     if path is not None:
                         click.echo("Using path: {}".format(path))
                     sys.exit(1)
-            return cls(device)
+            return cls(transport=device)
 
         ctx.obj = get_device
 

diff --git a/trezorlib/client.py b/trezorlib/client.py
@@ -496,16 +496,20 @@ class ProtocolMixin(object):
     PRIME_DERIVATION_FLAG = 0x80000000
     VENDORS = ('bitcointrezor.com', 'trezor.io')
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, state=None, *args, **kwargs):
         super(ProtocolMixin, self).__init__(*args, **kwargs)
+        self.state = state
         self.init_device()
         self.tx_api = None
 
     def set_tx_api(self, tx_api):
         self.tx_api = tx_api
 
     def init_device(self):
-        self.features = expect(proto.Features)(self.call)(proto.Initialize())
+        init_msg = proto.Initialize()
+        if self.state is not None:
+            init_msg.state = self.state
+        self.features = expect(proto.Features)(self.call)(init_msg)
         if str(self.features.vendor) not in self.VENDORS:
             raise RuntimeError("Unsupported device")