1 files changed, 203 insertions, 0 deletions
diff --git a/lib/python/qmk/flashers.py b/lib/python/qmk/flashers.py
new file mode 100644
index 0000000000..a9cf726b44
--- /dev/null
+++ b/lib/python/qmk/flashers.py
@@ -0,0 +1,203 @@
+import shutil
+import time
+import os
+import signal
+
+import usb.core
+
+from qmk.constants import BOOTLOADER_VIDS_PIDS
+from milc import cli
+
+# yapf: disable
+_PID_TO_MCU = {
+    '2fef': 'atmega16u2',
+    '2ff0': 'atmega32u2',
+    '2ff3': 'atmega16u4',
+    '2ff4': 'atmega32u4',
+    '2ff9': 'at90usb64',
+    '2ffa': 'at90usb162',
+    '2ffb': 'at90usb128'
+}
+
+AVRDUDE_MCU = {
+    'atmega32a': 'm32',
+    'atmega328p': 'm328p',
+    'atmega328': 'm328',
+}
+# yapf: enable
+
+
+class DelayedKeyboardInterrupt:
+    # Custom interrupt handler to delay the processing of Ctrl-C
+    # https://stackoverflow.com/a/21919644
+    def __enter__(self):
+        self.signal_received = False
+        self.old_handler = signal.signal(signal.SIGINT, self.handler)
+
+    def handler(self, sig, frame):
+        self.signal_received = (sig, frame)
+
+    def __exit__(self, type, value, traceback):
+        signal.signal(signal.SIGINT, self.old_handler)
+        if self.signal_received:
+            self.old_handler(*self.signal_received)
+
+
+# TODO: Make this more generic, so cli/doctor/check.py and flashers.py can share the code
+def _check_dfu_programmer_version():
+    # Return True if version is higher than 0.7.0: supports '--force'
+    check = cli.run(['dfu-programmer', '--version'], combined_output=True, timeout=5)
+    first_line = check.stdout.split('\n')[0]
+    version_number = first_line.split()[1]
+    maj, min_, bug = version_number.split('.')
+    if int(maj) >= 0 and int(min_) >= 7:
+        return True
+    else:
+        return False
+
+
+def _find_bootloader():
+    # To avoid running forever in the background, only look for bootloaders for 10min
+    start_time = time.time()
+    while time.time() - start_time < 600:
+        for bl in BOOTLOADER_VIDS_PIDS:
+            for vid, pid in BOOTLOADER_VIDS_PIDS[bl]:
+                vid_hex = int(f'0x{vid}', 0)
+                pid_hex = int(f'0x{pid}', 0)
+                with DelayedKeyboardInterrupt():
+                    # PyUSB does not like to be interrupted by Ctrl-C
+                    # therefore we catch the interrupt with a custom handler
+                    # and only process it once pyusb finished
+                    dev = usb.core.find(idVendor=vid_hex, idProduct=pid_hex)
+                if dev:
+                    if bl == 'atmel-dfu':
+                        details = _PID_TO_MCU[pid]
+                    elif bl == 'caterina':
+                        details = (vid_hex, pid_hex)
+                    elif bl == 'hid-bootloader':
+                        if vid == '16c0' and pid == '0478':
+                            details = 'halfkay'
+                        else:
+                            details = 'qmk-hid'
+                    elif bl == 'stm32-dfu' or bl == 'apm32-dfu' or bl == 'gd32v-dfu' or bl == 'kiibohd':
+                        details = (vid, pid)
+                    else:
+                        details = None
+                    return (bl, details)
+        time.sleep(0.1)
+    return (None, None)
+
+
+def _find_serial_port(vid, pid):
+    if 'windows' in cli.platform.lower():
+        from serial.tools.list_ports_windows import comports
+        platform = 'windows'
+    else:
+        from serial.tools.list_ports_posix import comports
+        platform = 'posix'
+
+    start_time = time.time()
+    # Caterina times out after 8 seconds
+    while time.time() - start_time < 8:
+        for port in comports():
+            port, desc, hwid = port
+            if f'{vid:04x}:{pid:04x}' in hwid.casefold():
+                if platform == 'windows':
+                    time.sleep(1)
+                    return port
+                else:
+                    start_time = time.time()
+                    # Wait until the port becomes writable before returning
+                    while time.time() - start_time < 8:
+                        if os.access(port, os.W_OK):
+                            return port
+                        else:
+                            time.sleep(0.5)
+                return None
+    return None
+
+
+def _flash_caterina(details, file):
+    port = _find_serial_port(details[0], details[1])
+    if port:
+        cli.run(['avrdude', '-p', 'atmega32u4', '-c', 'avr109', '-U', f'flash:w:{file}:i', '-P', port], capture_output=False)
+        return False
+    else:
+        return True
+
+
+def _flash_atmel_dfu(mcu, file):
+    force = '--force' if _check_dfu_programmer_version() else ''
+    cli.run(['dfu-programmer', mcu, 'erase', force], capture_output=False)
+    cli.run(['dfu-programmer', mcu, 'flash', force, file], capture_output=False)
+    cli.run(['dfu-programmer', mcu, 'reset'], capture_output=False)
+
+
+def _flash_hid_bootloader(mcu, details, file):
+    if details == 'halfkay':
+        if shutil.which('teensy-loader-cli'):
+            cmd = 'teensy-loader-cli'
+        elif shutil.which('teensy_loader_cli'):
+            cmd = 'teensy_loader_cli'
+
+    # Use 'hid_bootloader_cli' for QMK HID and as a fallback for HalfKay
+    if not cmd:
+        if shutil.which('hid_bootloader_cli'):
+            cmd = 'hid_bootloader_cli'
+        else:
+            return True
+
+    cli.run([cmd, f'-mmcu={mcu}', '-w', '-v', file], capture_output=False)
+
+
+def _flash_dfu_util(details, file):
+    # STM32duino
+    if details[0] == '1eaf' and details[1] == '0003':
+        cli.run(['dfu-util', '-a', '2', '-d', f'{details[0]}:{details[1]}', '-R', '-D', file], capture_output=False)
+    # kiibohd
+    elif details[0] == '1c11' and details[1] == 'b007':
+        cli.run(['dfu-util', '-a', '0', '-d', f'{details[0]}:{details[1]}', '-D', file], capture_output=False)
+    # STM32, APM32, or GD32V DFU
+    else:
+        cli.run(['dfu-util', '-a', '0', '-d', f'{details[0]}:{details[1]}', '-s', '0x08000000:leave', '-D', file], capture_output=False)
+
+
+def _flash_isp(mcu, programmer, file):
+    programmer = 'usbasp' if programmer == 'usbasploader' else 'usbtiny'
+    # Check if the provide mcu has an avrdude-specific name, otherwise pass on what the user provided
+    mcu = AVRDUDE_MCU.get(mcu, mcu)
+    cli.run(['avrdude', '-p', mcu, '-c', programmer, '-U', f'flash:w:{file}:i'], capture_output=False)
+
+
+def _flash_mdloader(file):
+    cli.run(['mdloader', '--first', '--download', file, '--restart'], capture_output=False)
+
+
+def flasher(mcu, file):
+    bl, details = _find_bootloader()
+    # Add a small sleep to avoid race conditions
+    time.sleep(1)
+    if bl == 'atmel-dfu':
+        _flash_atmel_dfu(details, file.name)
+    elif bl == 'caterina':
+        if _flash_caterina(details, file.name):
+            return (True, "The Caterina bootloader was found but is not writable. Check 'qmk doctor' output for advice.")
+    elif bl == 'hid-bootloader':
+        if mcu:
+            if _flash_hid_bootloader(mcu, details, file.name):
+                return (True, "Please make sure 'teensy_loader_cli' or 'hid_bootloader_cli' is available on your system.")
+        else:
+            return (True, "Specifying the MCU with '-m' is necessary for HalfKay/HID bootloaders!")
+    elif bl == 'stm32-dfu' or bl == 'apm32-dfu' or bl == 'gd32v-dfu' or bl == 'kiibohd':
+        _flash_dfu_util(details, file.name)
+    elif bl == 'usbasploader' or bl == 'usbtinyisp':
+        if mcu:
+            _flash_isp(mcu, bl, file.name)
+        else:
+            return (True, "Specifying the MCU with '-m' is necessary for ISP flashing!")
+    elif bl == 'md-boot':
+        _flash_mdloader(file.name)
+    else:
+        return (True, "Known bootloader found but flashing not currently supported!")
+
+    return (False, None)