Merge pull request #25 from Pieterv24/keylogger

Matrix tester

requirements.txt

@@ -8,4 +8,4 @@ pefile==2019.4.18
 PyInstaller==3.4
 PyQt5==5.9.2
 https://github.com/danthedeckie/simpleeval/archive/41c99b8e224a7a0ae0ac59c773598fe79a4470db.zip
-sip==4.19.8
+sip==4.19.8

src/main/python/keyboard_comm.py

@@ -26,6 +26,7 @@ CMD_VIA_KEYMAP_GET_BUFFER = 0x12
 CMD_VIA_VIAL_PREFIX = 0xFE
 
 VIA_LAYOUT_OPTIONS = 0x02
+VIA_SWITCH_MATRIX_STATE = 0x03
 
 CMD_VIAL_GET_KEYBOARD_ID = 0x00
 CMD_VIAL_GET_SIZE = 0x01
@@ -500,6 +501,13 @@ class Keyboard:
 
         self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_LOCK), retries=20)
 
+    def matrix_poll(self):
+        if self.via_protocol < 0:
+            return
+
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_GET_KEYBOARD_VALUE, VIA_SWITCH_MATRIX_STATE), retries=20)
+        return data
+
     def macro_serialize(self, macro):
         """
         Serialize a single macro, a macro is made out of macro actions (BasicAction)

src/main/python/keyboard_widget.py

@@ -12,6 +12,7 @@ class KeyWidget:
     def __init__(self, desc, scale, shift_x=0, shift_y=0):
         self.active = False
         self.masked = False
+        self.pressed = False
         self.desc = desc
         self.text = ""
         self.mask_text = ""
@@ -109,6 +110,9 @@ class KeyWidget:
     def setActive(self, active):
         self.active = active
 
+    def setPressed(self, pressed):
+        self.pressed = pressed
+
     def setColor(self, color):
         self.color = color
 
@@ -263,6 +267,16 @@ class KeyboardWidget(QWidget):
         active_brush.setColor(QApplication.palette().color(QPalette.Highlight))
         active_brush.setStyle(Qt.SolidPattern)
 
+        # for pressed keycaps
+        pressed_pen = qp.pen()
+        pressed_pen_color = QApplication.palette().color(QPalette.HighlightedText).lighter(75)
+        pressed_pen.setColor(pressed_pen_color)
+
+        pressed_brush = QBrush()
+        pressed_brush_color = QApplication.palette().color(QPalette.Highlight).lighter(75)
+        pressed_brush.setColor(pressed_brush_color)
+        pressed_brush.setStyle(Qt.SolidPattern)
+
         mask_font = qp.font()
         mask_font.setPointSize(mask_font.pointSize() * 0.8)
 
@@ -280,6 +294,12 @@ class KeyboardWidget(QWidget):
                 qp.setPen(active_pen)
                 qp.setBrush(active_brush)
 
+            if key.pressed:
+                # move key slightly down when pressed
+                qp.translate(0, 5)
+                qp.setPen(pressed_pen)
+                qp.setBrush(pressed_brush)
+
             # draw the keycap
             qp.drawPath(key.draw_path)
             qp.strokePath(key.draw_path2, regular_pen)

src/main/python/main_window.py

@@ -16,6 +16,7 @@ from macro_recorder import MacroRecorder
 from unlocker import Unlocker
 from util import tr, find_vial_devices, EXAMPLE_KEYBOARDS
 from vial_device import VialKeyboard
+from matrix_test import MatrixTest
 
 import themes
 
@@ -51,9 +52,10 @@ class MainWindow(QMainWindow):
         self.keymap_editor = KeymapEditor(self.layout_editor)
         self.firmware_flasher = FirmwareFlasher(self)
         self.macro_recorder = MacroRecorder()
+        self.matrix_tester = MatrixTest(self.layout_editor)
 
         self.editors = [(self.keymap_editor, "Keymap"), (self.layout_editor, "Layout"), (self.macro_recorder, "Macros"),
-                        (self.firmware_flasher, "Firmware updater")]
+                        (self.matrix_tester, "Matrix tester"), (self.firmware_flasher, "Firmware updater")]
         Unlocker.global_layout_editor = self.layout_editor
 
         self.tabs = QTabWidget()
@@ -234,7 +236,7 @@ class MainWindow(QMainWindow):
             Unlocker.unlock(self.current_device.keyboard)
             self.current_device.keyboard.reload()
 
-        for e in [self.layout_editor, self.keymap_editor, self.firmware_flasher, self.macro_recorder]:
+        for e in [self.layout_editor, self.keymap_editor, self.firmware_flasher, self.macro_recorder, self.matrix_tester]:
             e.rebuild(self.current_device)
 
     def refresh_tabs(self):

src/main/python/matrix_test.py

@@ -0,0 +1,115 @@
+from PyQt5.QtWidgets import QVBoxLayout, QPushButton
+from PyQt5.QtCore import Qt, QTimer
+import struct
+import math
+
+from basic_editor import BasicEditor
+from keyboard_widget import KeyboardWidget
+from vial_device import VialKeyboard
+from unlocker import Unlocker
+
+class MatrixTest(BasicEditor):
+    def __init__(self, layout_editor):
+        super().__init__()
+
+        self.layout_editor = layout_editor
+
+        self.keyboardWidget = KeyboardWidget(layout_editor)
+        self.keyboardWidget.set_enabled(False)
+
+        self.startButtonWidget = QPushButton("Start testing")
+        self.resetButtonWidget = QPushButton("Reset")
+
+        layout = QVBoxLayout()
+        layout.addWidget(self.keyboardWidget)
+        layout.setAlignment(self.keyboardWidget, Qt.AlignCenter)
+
+        self.addLayout(layout)
+        self.addWidget(self.resetButtonWidget)
+        self.addWidget(self.startButtonWidget)
+
+        self.keyboard = None
+        self.device = None
+        self.polling = False
+
+        self.timer = QTimer()
+        self.timer.timeout.connect(self.matrix_poller)
+
+        self.startButtonWidget.clicked.connect(self.start_poller)
+        self.resetButtonWidget.clicked.connect(self.reset_keyboard_widget)
+
+    def rebuild(self, device):
+        super().rebuild(device)
+        if self.valid():
+            self.keyboard = device.keyboard
+
+            self.keyboardWidget.set_keys(self.keyboard.keys, self.keyboard.encoders)
+
+    def valid(self):
+        # Check if vial protocol is v3 or later
+        return isinstance(self.device, VialKeyboard) and (self.device.keyboard and self.device.keyboard.vial_protocol >= 3)
+
+    def reset_keyboard_widget(self):
+        # reset keyboard widget
+        for w in self.keyboardWidget.widgets:
+            w.setPressed(False)
+            w.setActive(False)
+        
+        self.keyboardWidget.update_layout()
+        self.keyboardWidget.update()
+        self.keyboardWidget.updateGeometry()
+
+    def matrix_poller(self):
+        # Get size for matrix
+        rows = self.keyboard.rows
+        cols = self.keyboard.cols
+        # Generate 2d array of matrix
+        matrix = [ [None for y in range(cols)] for x in range(rows) ]
+
+        # Get matrix data from keyboard
+        data = self.keyboard.matrix_poll()
+        # Calculate the amount of bytes belong to 1 row, each bit is 1 key, so per 8 keys in a row, a byte is needed for the row.
+        row_size = math.ceil(cols / 8)
+
+
+        for row in range(rows):
+            # Make slice of bytes for the row (skip first 2 bytes, they're for VIAL)
+            row_data_start = 2 + (row * row_size)
+            row_data_end = row_data_start + row_size
+            row_data = data[row_data_start:row_data_end]
+
+            #Get each bit representing pressed state for col
+            for col in range(cols):
+                # row_data is array of bytes, calculate in which byte the col is located
+                col_byte = len(row_data) - 1 - math.floor(col / 8)
+                # since we select a single byte as slice of byte, mod 8 to get nth pos of byte
+                col_mod = (col % 8)
+                # write to matrix array
+                matrix[row][col] = (row_data[col_byte] >> col_mod) & 1
+
+        # write matrix state to keyboard widget
+        for w in self.keyboardWidget.widgets:
+            row = w.desc.row
+            col = w.desc.col
+
+            if row < len(matrix) and col < len(matrix[row]):
+                w.setPressed(matrix[row][col])
+                if matrix[row][col]:
+                    w.setActive(True)
+        
+        self.keyboardWidget.update_layout()
+        self.keyboardWidget.update()
+        self.keyboardWidget.updateGeometry()
+
+    def start_poller(self):
+        if not self.polling:
+            Unlocker.unlock(self.keyboard)
+            self.startButtonWidget.setText("Stop testing")
+            self.timer.start(20)
+            self.polling = True
+        else:
+            self.timer.stop()
+            self.keyboard.lock()
+            self.startButtonWidget.setText("Start testing")
+            self.polling = False
+