unlocker: allow re-locking keyboard

src/main/python/keyboard_comm.py

@@ -30,6 +30,7 @@ CMD_VIAL_SET_ENCODER = 0x04
 CMD_VIAL_GET_LOCK = 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
@@ -316,3 +317,6 @@ class Keyboard:
     def unlock_poll(self):
         data = self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_UNLOCK_POLL))
         return data
+
+    def lock(self):
+        self.usb_send(self.dev, struct.pack("BB", CMD_VIA_VIAL_PREFIX, CMD_VIAL_LOCK))

src/main/python/main_window.py

@@ -12,6 +12,7 @@ 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):
@@ -84,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")
@@ -127,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)
 
@@ -162,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

@@ -20,7 +20,8 @@ class Unlocker(QWidget):
 
         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", "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:")))