trezorlib/transport: move TrezorDevice functionality to transport and make it better ^_^

trezorlib/device.py

@@ -1,68 +0,0 @@
-# This file is part of the TREZOR project.
-#
-# Copyright (C) 2012-2017 Marek Palatinus <slush@satoshilabs.com>
-# Copyright (C) 2012-2017 Pavol Rusnak <stick@satoshilabs.com>
-#
-# This library is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Lesser General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This library is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU Lesser General Public License for more details.
-#
-# You should have received a copy of the GNU Lesser General Public License
-# along with this library.  If not, see <http://www.gnu.org/licenses/>.
-
-
-from .transport_bridge import BridgeTransport
-from .transport_hid import HidTransport
-from .transport_udp import UdpTransport
-from .transport_webusb import WebUsbTransport
-
-
-class TrezorDevice(object):
-
-    @classmethod
-    def enumerate(cls):
-        devices = []
-
-        for d in BridgeTransport.enumerate():
-            devices.append(d)
-
-        for d in UdpTransport.enumerate():
-            devices.append(d)
-
-        for d in HidTransport.enumerate():
-            devices.append(d)
-
-        for d in WebUsbTransport.enumerate():
-            devices.append(d)
-
-        return devices
-
-    @classmethod
-    def find_by_path(cls, path):
-        if path is None:
-            try:
-                return cls.enumerate()[0]
-            except IndexError:
-                raise Exception("No TREZOR device found")
-
-        prefix = path.split(':')[0]
-
-        if prefix == BridgeTransport.PATH_PREFIX:
-            return BridgeTransport.find_by_path(path)
-
-        if prefix == UdpTransport.PATH_PREFIX:
-            return UdpTransport.find_by_path(path)
-
-        if prefix == WebUsbTransport.PATH_PREFIX:
-            return WebUsbTransport.find_by_path(path)
-
-        if prefix == HidTransport.PATH_PREFIX:
-            return HidTransport.find_by_path(path)
-
-        raise Exception("Unknown path prefix '%s'" % prefix)

trezorlib/transport/__init__.py

@@ -60,3 +60,37 @@ class Transport(object):
                 return device
 
         raise TransportException('{} device not found: {}'.format(cls.PATH_PREFIX, path))
+
+
+def all_transports():
+    from .bridge import BridgeTransport
+    from .hid import HidTransport
+    from .udp import UdpTransport
+    from .webusb import WebUsbTransport
+    return (BridgeTransport, HidTransport, UdpTransport, WebUsbTransport)
+
+
+def enumerate_devices():
+    return [device
+            for transport in all_transports()
+            for device in transport.enumerate()]
+
+
+def get_transport(path=None, prefix_search=False):
+    if path is None:
+        try:
+            return enumerate_devices()[0]
+        except IndexError:
+            raise Exception("No TREZOR device found")
+
+    # Find whether B is prefix of A (transport name is part of the path)
+    # or A is prefix of B (path is a prefix, or a name, of transport).
+    # This naively expects that no two transports have a common prefix.
+    def match_prefix(a,b):
+        return a.startswith(b) or b.startswith(a)
+
+    transports = [t for t in all_transports() if match_prefix(path, t.PATH_PREFIX)]
+    if transports:
+        return transports[0].find_by_path(path, prefix_search=prefix_search)
+
+    raise Exception("Unknown path prefix '%s'" % prefix)