From a6c42b513f1538c8d48246627f10cb6f877c75e2 Mon Sep 17 00:00:00 2001
From: Ilya Zhuravlev <whatever@xyz.is>
Date: Sun, 27 Dec 2020 08:03:45 -0500
Subject: [PATCH] unlocker: initial implementation of unlocking keyboard to
 perform security-sensitive actions

---
 src/main/python/firmware_flasher.py |  6 ++-
 src/main/python/keyboard_comm.py    | 14 +++++++
 src/main/python/unlocker.py         | 60 +++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 src/main/python/unlocker.py

diff --git a/src/main/python/firmware_flasher.py b/src/main/python/firmware_flasher.py
index 5d5e9f6..85bcd76 100644
--- a/src/main/python/firmware_flasher.py
+++ b/src/main/python/firmware_flasher.py
@@ -11,6 +11,7 @@ from PyQt5.QtGui import QFontDatabase
 from PyQt5.QtWidgets import QHBoxLayout, QLineEdit, QToolButton, QPlainTextEdit, QProgressBar,QFileDialog, QDialog
 
 from basic_editor import BasicEditor
+from unlocker import Unlocker
 from util import tr, chunks, find_vial_devices
 from vial_device import VialBootloader, VialKeyboard
 
@@ -141,6 +142,8 @@ class FirmwareFlasher(BasicEditor):
 
         self.layout_restore = self.uid_restore = None
 
+        self.unlocker = Unlocker()
+
     def rebuild(self, device):
         super().rebuild(device)
         self.txt_logger.clear()
@@ -191,7 +194,6 @@ class FirmwareFlasher(BasicEditor):
 
         self.layout_restore = self.uid_restore = None
 
-        # TODO: this needs to switch to the secure assisted-reset feature before public release
         if isinstance(self.device, VialKeyboard):
             # back up current layout
             self.log("Backing up current layout...")
@@ -200,6 +202,8 @@ class FirmwareFlasher(BasicEditor):
             # keep track of which keyboard we should restore saved layout to
             self.uid_restore = self.device.keyboard.get_uid()
 
+            self.unlocker.perform_unlock(self.device.keyboard)
+
             self.log("Restarting in bootloader mode...")
             self.device.keyboard.reset()
 
diff --git a/src/main/python/keyboard_comm.py b/src/main/python/keyboard_comm.py
index 04005d0..6335ccf 100644
--- a/src/main/python/keyboard_comm.py
+++ b/src/main/python/keyboard_comm.py
@@ -27,6 +27,9 @@ CMD_VIAL_GET_SIZE = 0x01
 CMD_VIAL_GET_DEFINITION = 0x02
 CMD_VIAL_GET_ENCODER = 0x03
 CMD_VIAL_SET_ENCODER = 0x04
+CMD_VIAL_GET_LOCK = 0x05
+CMD_VIAL_UNLOCK_START = 0x06
+CMD_VIAL_UNLOCK_POLL = 0x07
 
 # how much of a macro/keymap buffer we can read/write per packet
 BUFFER_FETCH_CHUNK = 28
@@ -290,3 +293,14 @@ class Keyboard:
         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_lock(self):
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_GET_LOCK))
+        return data[0]
+
+    def unlock_start(self):
+        self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_UNLOCK_START))
+
+    def unlock_poll(self):
+        data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_UNLOCK_POLL))
+        return data
diff --git a/src/main/python/unlocker.py b/src/main/python/unlocker.py
new file mode 100644
index 0000000..b3d9b96
--- /dev/null
+++ b/src/main/python/unlocker.py
@@ -0,0 +1,60 @@
+# 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 util import tr
+
+
+class Unlocker(QWidget):
+
+    def __init__(self):
+        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.")))
+        layout.addWidget(QLabel(tr("Unlocker", "Press and hold the following keys until the progress bar "
+                                               "below fills up:")))
+
+        # TODO: add image/text reference of keys user needs to hold
+
+        layout.addWidget(self.progress)
+
+        self.setLayout(layout)
+        self.setWindowFlag(Qt.Dialog)
+
+    def perform_unlock(self, keyboard):
+        # if it's already unlocked, don't need to do anything
+        if keyboard.get_lock() == 0:
+            return
+
+        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()