diff options
Diffstat (limited to 'lib/python/qmk/flashers.py')
| -rw-r--r-- | lib/python/qmk/flashers.py | 203 |
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) |