From 66e3e603701b0e9db4640d2077ed5dfe7d50f494 Mon Sep 17 00:00:00 2001
From: slush <slush@satoshilabs.com>
Date: Sat, 3 Mar 2018 01:25:19 +0100
Subject: [PATCH] trezor: Use composite transport for device detection.

---
 libagent/device/trezor.py      | 43 +++++++++++++++-------------------
 libagent/device/trezor_defs.py |  4 +---
 2 files changed, 20 insertions(+), 27 deletions(-)

diff --git a/libagent/device/trezor.py b/libagent/device/trezor.py
index f1a2861..04db53a 100644
--- a/libagent/device/trezor.py
+++ b/libagent/device/trezor.py
@@ -41,12 +41,6 @@ class Trezor(interface.Device):
     @property
     def _defs(self):
         from . import trezor_defs
-        # Allow using TREZOR bridge transport (instead of the HID default)
-        trezor_defs.Transport = {
-            'bridge': trezor_defs.BridgeTransport,
-            'udp': trezor_defs.UdpTransport,
-            'hid': trezor_defs.HidTransport,
-        }[os.environ.get('TREZOR_TRANSPORT', 'hid')]
         return trezor_defs
 
     required_version = '>=1.4.0'
@@ -122,26 +116,27 @@ class Trezor(interface.Device):
 
     def connect(self):
         """Enumerate and connect to the first USB HID interface."""
-        for transport in self._defs.Transport.enumerate():
-            log.debug('transport: %s', transport)
-            for _ in range(5):
-                connection = self._defs.Client(transport)
-                self._override_pin_handler(connection)
-                self._override_passphrase_handler(connection)
-                self._verify_version(connection)
+        transport = self._defs.TrezorDevice.enumerate()
+        if not transport:
+            raise interface.NotFoundError('{} not connected'.format(self))
 
-                try:
-                    connection.ping(msg='', pin_protection=True)  # unlock PIN
-                    return connection
-                except (self._defs.PinException, ValueError) as e:
-                    log.error('Invalid PIN: %s, retrying...', e)
-                    continue
-                except Exception as e:
-                    log.exception('ping failed: %s', e)
-                    connection.close()  # so the next HID open() will succeed
-                    raise
+        log.debug('transports: %s', transport)
+        for _ in range(5):
+            connection = self._defs.Client(transport[0])
+            self._override_pin_handler(connection)
+            self._override_passphrase_handler(connection)
+            self._verify_version(connection)
 
-        raise interface.NotFoundError('{} not connected'.format(self))
+            try:
+                connection.ping(msg='', pin_protection=True)  # unlock PIN
+                return connection
+            except (self._defs.PinException, ValueError) as e:
+                log.error('Invalid PIN: %s, retrying...', e)
+                continue
+            except Exception as e:
+                log.exception('ping failed: %s', e)
+                connection.close()  # so the next HID open() will succeed
+                raise
 
     def close(self):
         """Close connection."""
diff --git a/libagent/device/trezor_defs.py b/libagent/device/trezor_defs.py
index 62501ef..5804e4b 100644
--- a/libagent/device/trezor_defs.py
+++ b/libagent/device/trezor_defs.py
@@ -5,6 +5,4 @@
 from trezorlib.client import CallException, PinException
 from trezorlib.client import TrezorClient as Client
 from trezorlib.messages import IdentityType, PassphraseAck, PinMatrixAck
-from trezorlib.transport_bridge import BridgeTransport
-from trezorlib.transport_hid import HidTransport
-from trezorlib.transport_udp import UdpTransport
+from trezorlib.device import TrezorDevice