summaryrefslogtreecommitdiffstats
path: root/lib/python/qmk/flashers.py
diff options
context:
space:
mode:
Diffstat (limited to 'lib/python/qmk/flashers.py')
-rw-r--r--lib/python/qmk/flashers.py203
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)