firmware_flasher: automatically restart keyboard in bootloader mode and flash

src/main/python/firmware_flasher.py

@@ -6,13 +6,16 @@ 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 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):
@@ -32,6 +35,18 @@ def send_retries(dev, data, retries=20):
 CHUNK = 64
 
 
+def bl_get_version(dev):
+    dev.send(b"VC\x00")
+    data = dev.recv(8)
+    return data[0]
+
+
+def bl_get_uid(dev):
+    dev.send(b"VC\x01")
+    data = dev.recv(8)
+    return data
+
+
 def cmd_flash(device, firmware, log_cb, progress_cb, complete_cb, error_cb):
     if firmware[0:8] != b"VIALFW00":
         return error_cb("Error: Invalid signature")
@@ -50,24 +65,22 @@ def cmd_flash(device, firmware, log_cb, progress_cb, complete_cb, error_cb):
         ))
 
     # Check bootloader is correct version
-    device.send(b"VC\x00")
-    data = device.recv(8)
-    log_cb("* Bootloader version: {}".format(data[0]))
-    if data[0] != 0:
+    ver = bl_get_version(device)
+    log_cb("* Bootloader version: {}".format(ver))
+    if ver != BL_SUPPORTED_VERSION:
         return error_cb("Error: Unsupported bootloader version")
 
-    device.send(b"VC\x01")
-    data = device.recv(8)
-    log_cb("* Vial ID: {}".format(data.hex()))
+    uid = bl_get_uid(device)
+    log_cb("* Vial ID: {}".format(uid.hex()))
 
-    if data == b"\xFF" * 8:
+    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 data != fw_uid:
+    if uid != fw_uid:
         return error_cb("Error: Firmware package was built for different device\n\texpected={}\n\tgot={}".format(
             fw_uid.hex(),
-            data.hex()
+            uid.hex()
         ))
 
     # OK all checks complete, we can flash now
@@ -139,12 +152,18 @@ class FirmwareFlasher(BasicEditor):
     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)
+        return isinstance(self.device, VialBootloader) or\
+               isinstance(self.device, VialKeyboard) and self.device.keyboard.vibl
 
     def on_click_select_file(self):
         dialog = QFileDialog()
@@ -170,6 +189,34 @@ class FirmwareFlasher(BasicEditor):
 
         self.log("Preparing to flash...")
         self.lock_ui()
+
+        # TODO: this needs to switch to the secure assisted-reset feature before public release
+        if isinstance(self.device, VialKeyboard):
+            uid = self.device.keyboard.get_uid()
+
+            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
+            while True:
+                self.log("Looking for devices...")
+                QCoreApplication.processEvents()
+                time.sleep(1)
+                devices = find_vial_devices()
+                found = None
+                for dev in devices:
+                    if isinstance(dev, VialBootloader):
+                        dev.open()
+                        # TODO: update version check before release
+                        if bl_get_version(dev) != BL_SUPPORTED_VERSION or bl_get_uid(dev) != uid:
+                            dev.close()
+                        found = dev
+                        break
+                if found:
+                    self.log("Found Vial Bootloader device at {}".format(found.desc["path"].decode("utf-8")))
+                    self.device = found
+                    break
+
         threading.Thread(target=lambda: cmd_flash(
             self.device, firmware, self.on_log, self.on_progress, self.on_complete, self.on_error)).start()
 

src/main/python/keyboard_comm.py

@@ -54,6 +54,7 @@ class Keyboard:
         self.macro_count = 0
         self.macro_memory = 0
         self.macro = b""
+        self.vibl = False
 
         self.vial_protocol = self.keyboard_id = -1
 
@@ -103,6 +104,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"]
@@ -275,3 +280,13 @@ class Keyboard:
 
         self.set_layout_options(data["layout_options"])
         self.set_macro(base64.b64decode(data["macro"]))
+
+    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

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 = []