From aec4908fd7c1d32a1a90d12d4fe3586b4a750ad6 Mon Sep 17 00:00:00 2001 From: slush0 Date: Thu, 10 Jul 2014 00:44:46 +0200 Subject: [PATCH] HidTransport raises ConnectionException when disconnected during HID session --- trezorlib/pinmatrix.py | 111 ------------------------------------- trezorlib/transport.py | 3 + trezorlib/transport_hid.py | 20 ++++++- 3 files changed, 20 insertions(+), 114 deletions(-) delete mode 100644 trezorlib/pinmatrix.py diff --git a/trezorlib/pinmatrix.py b/trezorlib/pinmatrix.py deleted file mode 100644 index 749c69b..0000000 --- a/trezorlib/pinmatrix.py +++ /dev/null @@ -1,111 +0,0 @@ -import sys -import math -import operator -from PyQt4.Qt import QApplication, QWidget, QGridLayout, QVBoxLayout, QHBoxLayout -from PyQt4.QtGui import QPushButton, QLineEdit, QSizePolicy, QRegExpValidator, QLabel -from PyQt4.QtCore import QObject, SIGNAL, QRegExp, Qt - -class PinButton(QPushButton): - def __init__(self, password, encoded_value): - super(PinButton, self).__init__('?') - self.password = password - self.encoded_value = encoded_value - - QObject.connect(self, SIGNAL('clicked()'), self._pressed) - - def _pressed(self): - self.password.setText(self.password.text() + str(self.encoded_value)) - print self.encoded_value - self.password.setFocus() - -class PinMatrixWidget(QWidget): - ''' - Displays widget with nine blank buttons and password box. - Encodes button clicks into sequence of numbers for passing - into PinAck messages of Trezor. - - show_strength=True may be useful for entering new PIN - ''' - def __init__(self, show_strength=True, parent=None): - super(PinMatrixWidget, self).__init__(parent) - - self.password = QLineEdit() - self.password.setValidator(QRegExpValidator(QRegExp('[1-9]+'), None)) - self.password.setEchoMode(QLineEdit.Password) - QObject.connect(self.password, SIGNAL('textChanged(QString)'), self._password_changed) - - self.strength = QLabel() - self.strength.setMinimumWidth(75) - self.strength.setAlignment(Qt.AlignCenter) - self._set_strength(0) - - grid = QGridLayout() - grid.setSpacing(0) - for y in range(3)[::-1]: - for x in range(3): - button = PinButton(self.password, x + y * 3 + 1) - button.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding) - button.setFocusPolicy(Qt.NoFocus) - grid.addWidget(button, 3 - y, x) - - hbox = QHBoxLayout() - hbox.addWidget(self.password) - if show_strength: - hbox.addWidget(self.strength) - - vbox = QVBoxLayout() - vbox.addLayout(grid) - vbox.addLayout(hbox) - self.setLayout(vbox) - - def _set_strength(self, strength): - if strength < 3000: - self.strength.setText('weak') - self.strength.setStyleSheet("QLabel { color : #d00; }") - elif strength < 60000: - self.strength.setText('fine') - self.strength.setStyleSheet("QLabel { color : #db0; }") - elif strength < 360000: - self.strength.setText('strong') - self.strength.setStyleSheet("QLabel { color : #0a0; }") - else: - self.strength.setText('ULTIMATE') - self.strength.setStyleSheet("QLabel { color : #000; font-weight: bold;}") - - def _password_changed(self, password): - self._set_strength(self.get_strength()) - - def get_strength(self): - digits = len(set(str(self.password.text()))) - strength = math.factorial(9) / math.factorial(9 - digits) - return strength - - def get_value(self): - return self.password.text() - -if __name__ == '__main__': - ''' - Demo application showing PinMatrix widget in action - ''' - a = QApplication(sys.argv) - - matrix = PinMatrixWidget() - - def clicked(): - print "PinMatrix value is", matrix.get_value() - print "Possible button combinations:", matrix.get_strength() - sys.exit() - - ok = QPushButton('OK') - QObject.connect(ok, SIGNAL('clicked()'), clicked) - - vbox = QVBoxLayout() - vbox.addWidget(matrix) - vbox.addWidget(ok) - - w = QWidget() - w.setLayout(vbox) - w.move(100, 100) - w.show() - - a.exec_() diff --git a/trezorlib/transport.py b/trezorlib/transport.py index f2aaeb0..afe59d4 100644 --- a/trezorlib/transport.py +++ b/trezorlib/transport.py @@ -4,6 +4,9 @@ import mapping class NotImplementedException(Exception): pass +class ConnectionError(Exception): + pass + class Transport(object): def __init__(self, device, *args, **kwargs): self.device = device diff --git a/trezorlib/transport_hid.py b/trezorlib/transport_hid.py index bb1f9ad..4a3cff5 100644 --- a/trezorlib/transport_hid.py +++ b/trezorlib/transport_hid.py @@ -3,7 +3,7 @@ import hid import time import platform -from transport import Transport, NotImplementedException +from transport import Transport, ConnectionError, NotImplementedException DEVICE_IDS = [ (0x10c4, 0xea80), # Shield @@ -17,11 +17,12 @@ class FakeRead(object): def read(self, size): return self.func(size) - + class HidTransport(Transport): def __init__(self, device, *args, **kwargs): self.hid = None self.buffer = '' + # self.read_timeout = kwargs.get('read_timeout') device = device[int(bool(kwargs.get('debug_link')))] super(HidTransport, self).__init__(device, *args, **kwargs) @@ -71,6 +72,13 @@ class HidTransport(Transport): # List of two-tuples (path_normal, path_debuglink) return devices.values() + + def is_connected(self): + # Check if the device is still connected + for d in hid.enumerate(0, 0): + if d['path'] == self.device: + return True + return False def _open(self): self.buffer = '' @@ -99,10 +107,16 @@ class HidTransport(Transport): (msg_type, datalen) = self._read_headers(FakeRead(self._raw_read)) return (msg_type, self._raw_read(datalen)) - def _raw_read(self, length): + def _raw_read(self, length): + start = time.time() while len(self.buffer) < length: data = self.hid.read(64) if not len(data): + if time.time() - start > 10 and not self.is_connected(): + # Over 10 of no response, let's check if + # device is still alive + raise ConnectionError("Connection failed") + time.sleep(0.05) continue