Merge branch 'vfw' into main

src/main/python/firmware_flasher.py

@@ -1,17 +1,22 @@
 # SPDX-License-Identifier: GPL-2.0-or-later
 
 import datetime
+import hashlib
 import struct
 import time
 import threading
 
-from PyQt5.QtCore import pyqtSignal
+from PyQt5.QtCore import pyqtSignal, QCoreApplication
 from PyQt5.QtGui import QFontDatabase
 from PyQt5.QtWidgets import QHBoxLayout, QLineEdit, QToolButton, QPlainTextEdit, QProgressBar,QFileDialog, QDialog
 
 from basic_editor import BasicEditor
-from util import tr, chunks
+from unlocker import Unlocker
-from vial_device import VialBootloader
+from util import tr, chunks, find_vial_devices
+from vial_device import VialBootloader, VialKeyboard
+
+
+BL_SUPPORTED_VERSION = 0
 
 
 def send_retries(dev, data, retries=20):
@@ -31,36 +36,65 @@ def send_retries(dev, data, retries=20):
 CHUNK = 64
 
 
-def cmd_flash(device, firmware, log_cb, progress_cb, complete_cb, error_cb):
+def cmd_flash(device, firmware, enable_insecure, log_cb, progress_cb, complete_cb, error_cb):
-    while len(firmware) % CHUNK != 0:
+    if firmware[0:8] != b"VIALFW00":
-        firmware += b"\x00"
+        return error_cb("Error: Invalid signature")
+
+    fw_uid = firmware[8:16]
+    fw_ts = struct.unpack("<Q", firmware[16:24])[0]
+    log_cb("* Firmware build date: {} (UTC)".format(datetime.datetime.utcfromtimestamp(fw_ts)))
+
+    fw_hash = firmware[32:64]
+    fw_payload = firmware[64:]
+
+    if hashlib.sha256(fw_payload).digest() != fw_hash:
+        return error_cb("Error: Firmware failed integrity check\n\texpected={}\n\tgot={}".format(
+            fw_hash.hex(),
+            hashlib.sha256(fw_payload).hexdigest()
+        ))
 
     # Check bootloader is correct version
     device.send(b"VC\x00")
-    data = device.recv(8)
+    ver = device.recv(8)[0]
-    log_cb("* Bootloader version: {}".format(data[0]))
+    log_cb("* Bootloader version: {}".format(ver))
-    if data[0] != 0:
+    if ver != BL_SUPPORTED_VERSION:
         return error_cb("Error: Unsupported bootloader version")
 
-    # TODO: Check vial ID against firmware package
     device.send(b"VC\x01")
-    data = device.recv(8)
+    uid = device.recv(8)
-    log_cb("* Vial ID: {}".format(data.hex()))
+    log_cb("* Vial ID: {}".format(uid.hex()))
+
+    if uid == b"\xFF" * 8:
+        log_cb("\n\n\n!!! WARNING !!!\nBootloader UID is not set, make sure to configure it"
+               " before releasing production firmware\n!!! WARNING !!!\n\n")
+
+    if uid != fw_uid:
+        return error_cb("Error: Firmware package was built for different device\n\texpected={}\n\tgot={}".format(
+            fw_uid.hex(),
+            uid.hex()
+        ))
+
+    # OK all checks complete, we can flash now
+    while len(fw_payload) % CHUNK != 0:
+        fw_payload += b"\x00"
 
     # Flash
     log_cb("Flashing...")
-    device.send(b"VC\x02" + struct.pack("<H", len(firmware) // CHUNK))
+    device.send(b"VC\x02" + struct.pack("<H", len(fw_payload) // CHUNK))
     total = 0
-    for part in chunks(firmware, CHUNK):
+    for part in chunks(fw_payload, CHUNK):
         if len(part) < CHUNK:
             part += b"\x00" * (CHUNK - len(part))
         if not send_retries(device, part):
             return error_cb("Error while sending data, firmware is corrupted")
         total += len(part)
-        progress_cb(total / len(firmware))
+        progress_cb(total / len(fw_payload))
 
     # Reboot
     log_cb("Rebooting...")
+    # enable insecure mode on first boot in order to restore keymap/macros
+    if enable_insecure:
+        device.send(b"VC\x04")
     device.send(b"VC\x03")
 
     complete_cb("Done!")
@@ -109,22 +143,36 @@ class FirmwareFlasher(BasicEditor):
 
         self.device = None
 
+        self.layout_restore = self.uid_restore = None
+
     def rebuild(self, device):
         super().rebuild(device)
         self.txt_logger.clear()
+
+        if not self.valid():
+            return
+
         if isinstance(self.device, VialBootloader):
             self.log("Valid Vial Bootloader device at {}".format(self.device.desc["path"].decode("utf-8")))
+        elif isinstance(self.device, VialKeyboard):
+            self.log("Vial keyboard detected")
 
     def valid(self):
-        # TODO: it is also valid to flash a VialKeyboard which supports optional "vibl-integration" feature
+        return isinstance(self.device, VialBootloader) or\
-        return isinstance(self.device, VialBootloader)
+               isinstance(self.device, VialKeyboard) and self.device.keyboard.vibl
+
+    def find_device_with_uid(self, cls, uid):
+        devices = find_vial_devices()
+        for dev in devices:
+            if isinstance(dev, cls) and dev.get_uid() == uid:
+                return dev
+        return None
 
     def on_click_select_file(self):
         dialog = QFileDialog()
-        # TODO: this should be .vfw for Vial Firmware
+        dialog.setDefaultSuffix("vfw")
-        dialog.setDefaultSuffix("bin")
         dialog.setAcceptMode(QFileDialog.AcceptOpen)
-        dialog.setNameFilters(["Vial Firmware (*.bin)"])
+        dialog.setNameFilters(["Vial Firmware (*.vfw)"])
         if dialog.exec_() == QDialog.Accepted:
             self.selected_firmware_path = dialog.selectedFiles()[0]
             self.txt_file_selector.setText(self.selected_firmware_path)
@@ -144,8 +192,37 @@ class FirmwareFlasher(BasicEditor):
 
         self.log("Preparing to flash...")
         self.lock_ui()
+
+        self.layout_restore = self.uid_restore = None
+
+        if isinstance(self.device, VialKeyboard):
+            # back up current layout
+            self.log("Backing up current layout...")
+            self.layout_restore = self.device.keyboard.save_layout()
+
+            # keep track of which keyboard we should restore saved layout to
+            self.uid_restore = self.device.keyboard.get_uid()
+
+            Unlocker.get().perform_unlock(self.device.keyboard)
+
+            self.log("Restarting in bootloader mode...")
+            self.device.keyboard.reset()
+
+            # watch for bootloaders to appear and ask them for their UID, return one that matches the keyboard
+            found = None
+            while found is None:
+                self.log("Looking for devices...")
+                QCoreApplication.processEvents()
+                time.sleep(1)
+                found = self.find_device_with_uid(VialBootloader, self.uid_restore)
+
+            self.log("Found Vial Bootloader device at {}".format(found.desc["path"].decode("utf-8")))
+            found.open()
+            self.device = found
+
         threading.Thread(target=lambda: cmd_flash(
-            self.device, firmware, self.on_log, self.on_progress, self.on_complete, self.on_error)).start()
+            self.device, firmware, self.layout_restore is not None,
+            self.on_log, self.on_progress, self.on_complete, self.on_error)).start()
 
     def on_log(self, line):
         self.log_signal.emit(line)
@@ -168,6 +245,25 @@ class FirmwareFlasher(BasicEditor):
     def _on_complete(self, msg):
         self.log(msg)
         self.progress_bar.setValue(100)
+
+        # if we were asked to restore a layout, find keyboard with matching UID and restore the layout to it
+        if self.layout_restore:
+            found = None
+            while found is None:
+                self.log("Looking for devices...")
+                QCoreApplication.processEvents()
+                time.sleep(1)
+                found = self.find_device_with_uid(VialKeyboard, self.uid_restore)
+
+            self.log("Found Vial keyboard at {}".format(found.desc["path"].decode("utf-8")))
+            found.open()
+            self.device = found
+            self.log("Restoring saved layout...")
+            QCoreApplication.processEvents()
+            found.keyboard.restore_layout(self.layout_restore)
+            found.close()
+            self.log("Done!")
+
         self.unlock_ui()
 
     def _on_error(self, msg):

src/main/python/keyboard_comm.py

@@ -6,6 +6,7 @@ import lzma
 from collections import OrderedDict
 
 from kle_serial import Serial as KleSerial
+from unlocker import Unlocker
 from util import MSG_LEN, hid_send, chunks
 
 CMD_VIA_GET_KEYBOARD_VALUE = 0x02
@@ -27,6 +28,10 @@ CMD_VIAL_GET_SIZE = 0x01
 CMD_VIAL_GET_DEFINITION = 0x02
 CMD_VIAL_GET_ENCODER = 0x03
 CMD_VIAL_SET_ENCODER = 0x04
+CMD_VIAL_GET_UNLOCK_STATUS = 0x05
+CMD_VIAL_UNLOCK_START = 0x06
+CMD_VIAL_UNLOCK_POLL = 0x07
+CMD_VIAL_LOCK = 0x08
 
 # how much of a macro/keymap buffer we can read/write per packet
 BUFFER_FETCH_CHUNK = 28
@@ -50,10 +55,10 @@ class Keyboard:
         self.layout_options = -1
         self.keys = []
         self.encoders = []
-        self.sideload = False
         self.macro_count = 0
         self.macro_memory = 0
         self.macro = b""
+        self.vibl = False
 
         self.vial_protocol = self.keyboard_id = -1
 
@@ -80,7 +85,6 @@ class Keyboard:
 
         if sideload_json is not None:
             payload = sideload_json
-            self.sideload = True
         else:
             # get keyboard identification
             data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_KEYBOARD_ID))
@@ -103,6 +107,10 @@ class Keyboard:
 
             payload = json.loads(lzma.decompress(payload))
 
+        if "vial" in payload:
+            vial = payload["vial"]
+            self.vibl = vial.get("vibl", False)
+
         self.layouts = payload.get("layouts")
 
         self.rows = payload["matrix"]["rows"]
@@ -274,4 +282,64 @@ class Keyboard:
                 self.set_encoder(l, e, 1, encoder[1])
 
         self.set_layout_options(data["layout_options"])
-        self.set_macro(base64.b64decode(data["macro"]))
+
+        # we need to unlock the keyboard before we can restore the macros, lock it afterwards
+        # only do that if the user actually has macros defined
+        macro = base64.b64decode(data["macro"])
+        if macro:
+            Unlocker.get().perform_unlock(self)
+            self.set_macro(macro)
+            self.lock()
+
+    def reset(self):
+        self.usb_send(self.dev, struct.pack("B", 0xB))
+        self.dev.close()
+
+    def get_uid(self):
+        """ Retrieve UID from the keyboard, explicitly sending a query packet """
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_KEYBOARD_ID))
+        keyboard_id = data[4:12]
+        return keyboard_id
+
+    def get_unlock_status(self):
+        # VIA keyboards are always unlocked
+        if self.vial_protocol < 0:
+            return 1
+
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_UNLOCK_STATUS))
+        return data[0]
+
+    def get_unlock_keys(self):
+        """ Return keys users have to hold to unlock the keyboard as a list of rowcols """
+
+        # VIA keyboards don't have unlock keys
+        if self.vial_protocol < 0:
+            return []
+
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_UNLOCK_STATUS))
+        rowcol = []
+        for x in range(15):
+            row = data[2 + x * 2]
+            col = data[3 + x * 2]
+            if row != 255 and col != 255:
+                rowcol.append((row, col))
+        return rowcol
+
+    def unlock_start(self):
+        if self.vial_protocol < 0:
+            return
+
+        self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_UNLOCK_START))
+
+    def unlock_poll(self):
+        if self.vial_protocol < 0:
+            return b""
+
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_UNLOCK_POLL))
+        return data
+
+    def lock(self):
+        if self.vial_protocol < 0:
+            return
+
+        self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_LOCK))

src/main/python/keyboard_widget.py

@@ -10,6 +10,8 @@ from constants import KEY_WIDTH, KEY_SPACING, KEY_HEIGHT, KEYBOARD_WIDGET_PADDIN
 class KeyWidget:
 
     def __init__(self, desc, shift_x=0, shift_y=0):
+        self.active = False
+        self.masked = False
         self.desc = desc
         self.text = ""
         self.mask_text = ""
@@ -88,6 +90,9 @@ class KeyWidget:
     def setToolTip(self, tooltip):
         self.tooltip = tooltip
 
+    def setActive(self, active):
+        self.active = active
+
 
 class EncoderWidget(KeyWidget):
 
@@ -117,6 +122,10 @@ class KeyboardWidget(QWidget):
 
     def __init__(self, layout_editor):
         super().__init__()
+
+        self.enabled = True
+        self.scale = 1
+
         self.setMouseTracking(True)
 
         self.layout_editor = layout_editor
@@ -195,8 +204,8 @@ class KeyboardWidget(QWidget):
         max_w = max_h = 0
         for key in self.widgets:
             p = key.polygon.boundingRect().bottomRight()
-            max_w = max(max_w, p.x())
+            max_w = max(max_w, p.x() * self.scale)
-            max_h = max(max_h, p.y())
+            max_h = max(max_h, p.y() * self.scale)
 
         self.width = max_w + 2 * KEYBOARD_WIDGET_PADDING
         self.height = max_h + 2 * KEYBOARD_WIDGET_PADDING
@@ -234,11 +243,12 @@ class KeyboardWidget(QWidget):
 
         for idx, key in enumerate(self.widgets):
             qp.save()
+            qp.scale(self.scale, self.scale)
             qp.translate(key.rotation_x, key.rotation_y)
             qp.rotate(key.rotation_angle)
             qp.translate(-key.rotation_x, -key.rotation_y)
 
-            if self.active_key == key and not self.active_mask:
+            if key.active or (self.active_key == key and not self.active_mask):
                 qp.setPen(active_pen)
                 qp.setBrush(active_brush)
 
@@ -280,6 +290,9 @@ class KeyboardWidget(QWidget):
         return None, False
 
     def mousePressEvent(self, ev):
+        if not self.enabled:
+            return
+
         self.active_key, self.active_mask = self.hit_test(ev.pos())
         if self.active_key is not None:
             self.clicked.emit()
@@ -302,6 +315,9 @@ class KeyboardWidget(QWidget):
             self.update()
 
     def event(self, ev):
+        if not self.enabled:
+            super().event(ev)
+
         if ev.type() == QEvent.ToolTip:
             key = self.hit_test(ev.pos())[0]
             if key is not None:
@@ -309,3 +325,9 @@ class KeyboardWidget(QWidget):
             else:
                 QToolTip.hideText()
         return super().event(ev)
+
+    def set_enabled(self, val):
+        self.enabled = val
+
+    def set_scale(self, scale):
+        self.scale = scale

src/main/python/macro_recorder.py

@@ -11,6 +11,7 @@ from macro_action import ActionText, ActionTap, ActionDown, ActionUp, SS_TAP_COD
 from macro_key import KeyString, KeyDown, KeyUp, KeyTap
 from macro_line import MacroLine
 from macro_optimizer import macro_optimize
+from unlocker import Unlocker
 from util import tr
 from vial_device import VialKeyboard
 
@@ -329,5 +330,6 @@ class MacroRecorder(BasicEditor):
         self.deserialize(self.keyboard.macro)
 
     def on_save(self):
+        Unlocker.get().perform_unlock(self.device.keyboard)
         self.keyboard.set_macro(self.serialize())
         self.on_change()

src/main/python/main_window.py

@@ -10,13 +10,16 @@ from firmware_flasher import FirmwareFlasher
 from keymap_editor import KeymapEditor
 from layout_editor import LayoutEditor
 from macro_recorder import MacroRecorder
+from unlocker import Unlocker
 from util import tr, find_vial_devices
+from vial_device import VialKeyboard
 
 
 class MainWindow(QMainWindow):
 
     def __init__(self):
         super().__init__()
+
         self.current_device = None
         self.devices = []
         self.sideload_json = None
@@ -41,6 +44,7 @@ class MainWindow(QMainWindow):
 
         self.editors = [(self.keymap_editor, "Keymap"), (self.layout_editor, "Layout"), (self.macro_recorder, "Macros"),
                         (self.firmware_flasher, "Firmware updater")]
+        self.unlocker = Unlocker(self.layout_editor)
 
         self.tabs = QTabWidget()
         self.refresh_tabs()
@@ -81,6 +85,16 @@ class MainWindow(QMainWindow):
         file_menu.addSeparator()
         file_menu.addAction(exit_act)
 
+        keyboard_unlock_act = QAction(tr("MenuSecurity", "Unlock"), self)
+        keyboard_unlock_act.triggered.connect(self.unlock_keyboard)
+
+        keyboard_lock_act = QAction(tr("MenuSecurity", "Lock"), self)
+        keyboard_lock_act.triggered.connect(self.lock_keyboard)
+
+        self.security_menu = self.menuBar().addMenu(tr("Menu", "Security"))
+        self.security_menu.addAction(keyboard_unlock_act)
+        self.security_menu.addAction(keyboard_lock_act)
+
     def on_layout_load(self):
         dialog = QFileDialog()
         dialog.setDefaultSuffix("vil")
@@ -124,6 +138,9 @@ class MainWindow(QMainWindow):
         self.refresh_tabs()
 
     def rebuild(self):
+        # don't show "Security" menu for bootloader mode, as the bootloader is inherently insecure
+        self.security_menu.menuAction().setVisible(isinstance(self.current_device, VialKeyboard))
+
         for e in [self.layout_editor, self.keymap_editor, self.firmware_flasher, self.macro_recorder]:
             e.rebuild(self.current_device)
 
@@ -159,3 +176,11 @@ class MainWindow(QMainWindow):
         self.tabs.setEnabled(True)
         self.combobox_devices.setEnabled(True)
         self.btn_refresh_devices.setEnabled(True)
+
+    def unlock_keyboard(self):
+        if isinstance(self.current_device, VialKeyboard):
+            self.unlocker.perform_unlock(self.current_device.keyboard)
+
+    def lock_keyboard(self):
+        if isinstance(self.current_device, VialKeyboard):
+            self.current_device.keyboard.lock()

src/main/python/unlocker.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+import time
+
+from PyQt5.QtCore import QCoreApplication, Qt
+from PyQt5.QtWidgets import QWidget, QVBoxLayout, QLabel, QProgressBar
+
+from keyboard_widget import KeyboardWidget
+from util import tr
+
+
+class Unlocker(QWidget):
+
+    def __init__(self, layout_editor):
+        super().__init__()
+        self.keyboard = None
+
+        layout = QVBoxLayout()
+
+        self.progress = QProgressBar()
+
+        layout.addWidget(QLabel(tr("Unlocker", "In order to proceed, the keyboard must be set into unlocked mode.\n"
+                                               "You should only perform this operation on computers that you trust.")))
+        layout.addWidget(QLabel(tr("Unlocker", "To exit this mode, you will need to replug the keyboard\n"
+                                               "or select Security->Lock from the menu.")))
+        layout.addWidget(QLabel(tr("Unlocker", "Press and hold the following keys until the progress bar "
+                                               "below fills up:")))
+
+        self.keyboard_reference = KeyboardWidget(layout_editor)
+        self.keyboard_reference.set_enabled(False)
+        self.keyboard_reference.set_scale(0.5)
+        layout.addWidget(self.keyboard_reference)
+        layout.setAlignment(self.keyboard_reference, Qt.AlignHCenter)
+
+        layout.addWidget(self.progress)
+
+        self.setLayout(layout)
+        self.setWindowFlag(Qt.Dialog)
+
+        Unlocker.obj = self
+
+    @classmethod
+    def get(cls):
+        return cls.obj
+
+    def update_reference(self, keyboard):
+        """ Updates keycap reference image """
+
+        self.keyboard_reference.set_keys(keyboard.keys, keyboard.encoders)
+
+        # use "active" background to indicate keys to hold
+        lock_keys = keyboard.get_unlock_keys()
+        for w in self.keyboard_reference.widgets:
+            if (w.desc.row, w.desc.col) in lock_keys:
+                w.setActive(True)
+
+        self.keyboard_reference.update_layout()
+        self.keyboard_reference.update()
+        self.keyboard_reference.updateGeometry()
+
+    def perform_unlock(self, keyboard):
+        # if it's already unlocked, don't need to do anything
+        unlock = keyboard.get_unlock_status()
+        if unlock == 1:
+            return
+
+        self.update_reference(keyboard)
+
+        self.progress.setMaximum(1)
+        self.progress.setValue(0)
+
+        self.show()
+        self.keyboard = keyboard
+        self.keyboard.unlock_start()
+
+        while True:
+            data = self.keyboard.unlock_poll()
+            unlocked = data[0]
+            unlock_counter = data[2]
+
+            self.progress.setMaximum(max(self.progress.maximum(), unlock_counter))
+            self.progress.setValue(self.progress.maximum() - unlock_counter)
+
+            if unlocked == 1:
+                break
+
+            QCoreApplication.processEvents()
+            time.sleep(0.2)
+
+        # ok all done, the keyboard is now set to insecure state
+        self.hide()
+
+    def closeEvent(self, ev):
+        ev.ignore()

src/main/python/util.py

@@ -32,7 +32,7 @@ def is_rawhid(dev):
     return dev["interface_number"] == 1
 
 
-def find_vial_devices(sideload_vid, sideload_pid):
+def find_vial_devices(sideload_vid=None, sideload_pid=None):
     from vial_device import VialBootloader, VialKeyboard
 
     filtered = []

src/main/python/vial_device.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: GPL-2.0-or-later
 from hidproxy import hid
 from keyboard_comm import Keyboard
+from util import MSG_LEN
 
 
 class VialDevice:
@@ -19,8 +20,8 @@ class VialDevice:
         # add 00 at start for hidapi report id
         return self.dev.write(b"\x00" + data)
 
-    def recv(self, length):
+    def recv(self, length, timeout_ms=0):
-        return bytes(self.dev.read(length))
+        return bytes(self.dev.read(length, timeout_ms=timeout_ms))
 
     def close(self):
         self.dev.close()
@@ -44,8 +45,28 @@ class VialKeyboard(VialDevice):
             s += " [sideload]"
         return s
 
+    def get_uid(self):
+        try:
+            super().open()
+        except OSError:
+            return b""
+        self.send(b"\xFE\x00" + b"\x00" * 30)
+        data = self.recv(MSG_LEN, timeout_ms=500)
+        super().close()
+        return data[4:12]
+
 
 class VialBootloader(VialDevice):
 
     def title(self):
         return "Vial Bootloader [{:04X}:{:04X}]".format(self.desc["vendor_id"], self.desc["product_id"])
+
+    def get_uid(self):
+        try:
+            super().open()
+        except OSError:
+            return b""
+        self.send(b"VC\x01")
+        data = self.recv(8, timeout_ms=500)
+        super().close()
+        return data