Added strength indicator (weak / fine / strong / ultimate)

trezorlib/pinmatrix.py

@@ -1,6 +1,8 @@
 import sys
-from PyQt4.Qt import QApplication, QWidget, QGridLayout, QVBoxLayout
-from PyQt4.QtGui import QPushButton, QLineEdit, QSizePolicy, QRegExpValidator
+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):
@@ -20,13 +22,21 @@ 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, parent=None):
+    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)
@@ -36,11 +46,38 @@ class PinMatrixWidget(QWidget):
             button.setFocusPolicy(Qt.NoFocus)
             grid.addWidget(button, x / 3, x % 3)
 
+        hbox = QHBoxLayout()
+        hbox.addWidget(self.password)
+        if show_strength:
+            hbox.addWidget(self.strength)
+
         vbox = QVBoxLayout()
         vbox.addLayout(grid)
-        vbox.addWidget(self.password)
+        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()
 
@@ -54,6 +91,7 @@ if __name__ == '__main__':
 
     def clicked():
         print "PinMatrix value is", matrix.get_value()
+        print "Possible button combinations:", matrix.get_strength()
         sys.exit()
 
     ok = QPushButton('OK')