flashers.py (9013B)
1 import platform 2 import shutil 3 import time 4 import os 5 import signal 6 7 import usb.core 8 9 from qmk.constants import BOOTLOADER_VIDS_PIDS 10 from milc import cli 11 12 # yapf: disable 13 _PID_TO_MCU = { 14 '2fef': 'atmega16u2', 15 '2ff0': 'atmega32u2', 16 '2ff3': 'atmega16u4', 17 '2ff4': 'atmega32u4', 18 '2ff9': 'at90usb64', 19 '2ffa': 'at90usb162', 20 '2ffb': 'at90usb128' 21 } 22 23 AVRDUDE_MCU = { 24 'atmega32a': 'm32', 25 'atmega328p': 'm328p', 26 'atmega328': 'm328', 27 } 28 # yapf: enable 29 30 31 class DelayedKeyboardInterrupt: 32 # Custom interrupt handler to delay the processing of Ctrl-C 33 # https://stackoverflow.com/a/21919644 34 def __enter__(self): 35 self.signal_received = False 36 self.old_handler = signal.signal(signal.SIGINT, self.handler) 37 38 def handler(self, sig, frame): 39 self.signal_received = (sig, frame) 40 41 def __exit__(self, type, value, traceback): 42 signal.signal(signal.SIGINT, self.old_handler) 43 if self.signal_received: 44 self.old_handler(*self.signal_received) 45 46 47 # TODO: Make this more generic, so cli/doctor/check.py and flashers.py can share the code 48 def _check_dfu_programmer_version(): 49 # Return True if version is higher than 0.7.0: supports '--force' 50 check = cli.run(['dfu-programmer', '--version'], combined_output=True, timeout=5) 51 first_line = check.stdout.split('\n')[0] 52 version_number = first_line.split()[1] 53 maj, min_, bug = version_number.split('.') 54 if int(maj) >= 0 and int(min_) >= 7: 55 return True 56 else: 57 return False 58 59 60 def _find_usb_device(vid_hex, pid_hex): 61 # WSL doesnt have access to USB - use powershell instead...? 62 if 'microsoft' in platform.uname().release.lower(): 63 ret = cli.run(['powershell.exe', '-command', 'Get-PnpDevice -PresentOnly | Select-Object -Property InstanceId']) 64 if f'USB\\VID_{vid_hex:04X}&PID_{pid_hex:04X}' in ret.stdout: 65 return (vid_hex, pid_hex) 66 else: 67 with DelayedKeyboardInterrupt(): 68 # PyUSB does not like to be interrupted by Ctrl-C 69 # therefore we catch the interrupt with a custom handler 70 # and only process it once pyusb finished 71 return usb.core.find(idVendor=vid_hex, idProduct=pid_hex) 72 73 74 def _find_uf2_devices(): 75 """Delegate to uf2conv.py as VID:PID pairs can potentially fluctuate more than other bootloaders 76 """ 77 return cli.run(['util/uf2conv.py', '--list']).stdout.splitlines() 78 79 80 def _find_bootloader(): 81 # To avoid running forever in the background, only look for bootloaders for 10min 82 start_time = time.time() 83 while time.time() - start_time < 600: 84 for bl in BOOTLOADER_VIDS_PIDS: 85 for vid, pid in BOOTLOADER_VIDS_PIDS[bl]: 86 vid_hex = int(f'0x{vid}', 0) 87 pid_hex = int(f'0x{pid}', 0) 88 dev = _find_usb_device(vid_hex, pid_hex) 89 if dev: 90 if bl == 'atmel-dfu': 91 details = _PID_TO_MCU[pid] 92 elif bl == 'caterina': 93 details = (vid_hex, pid_hex) 94 elif bl == 'hid-bootloader': 95 if vid == '16c0' and pid == '0478': 96 details = 'halfkay' 97 else: 98 details = 'qmk-hid' 99 elif bl in {'apm32-dfu', 'at32-dfu', 'gd32v-dfu', 'kiibohd', 'stm32-dfu'}: 100 details = (vid, pid) 101 else: 102 details = None 103 return (bl, details) 104 if _find_uf2_devices(): 105 return ('_uf2_compatible_', None) 106 time.sleep(0.1) 107 return (None, None) 108 109 110 def _find_serial_port(vid, pid): 111 if 'windows' in cli.platform.lower(): 112 from serial.tools.list_ports_windows import comports 113 platform = 'windows' 114 else: 115 from serial.tools.list_ports_posix import comports 116 platform = 'posix' 117 118 start_time = time.time() 119 # Caterina times out after 8 seconds 120 while time.time() - start_time < 8: 121 for port in comports(): 122 port, desc, hwid = port 123 if f'{vid:04x}:{pid:04x}' in hwid.casefold(): 124 if platform == 'windows': 125 time.sleep(1) 126 return port 127 else: 128 start_time = time.time() 129 # Wait until the port becomes writable before returning 130 while time.time() - start_time < 8: 131 if os.access(port, os.W_OK): 132 return port 133 else: 134 time.sleep(0.5) 135 return None 136 return None 137 138 139 def _flash_caterina(details, file): 140 port = _find_serial_port(details[0], details[1]) 141 if port: 142 cli.run(['avrdude', '-p', 'atmega32u4', '-c', 'avr109', '-U', f'flash:w:{file}:i', '-P', port], capture_output=False) 143 return False 144 else: 145 return True 146 147 148 def _flash_atmel_dfu(mcu, file): 149 force = '--force' if _check_dfu_programmer_version() else '' 150 cli.run(['dfu-programmer', mcu, 'erase', force], capture_output=False) 151 cli.run(['dfu-programmer', mcu, 'flash', force, file], capture_output=False) 152 cli.run(['dfu-programmer', mcu, 'reset'], capture_output=False) 153 154 155 def _flash_hid_bootloader(mcu, details, file): 156 cmd = None 157 if details == 'halfkay': 158 if shutil.which('teensy_loader_cli'): 159 cmd = 'teensy_loader_cli' 160 elif shutil.which('teensy-loader-cli'): 161 cmd = 'teensy-loader-cli' 162 163 # Use 'hid_bootloader_cli' for QMK HID and as a fallback for HalfKay 164 if not cmd: 165 if shutil.which('hid_bootloader_cli'): 166 cmd = 'hid_bootloader_cli' 167 else: 168 return True 169 170 cli.run([cmd, f'-mmcu={mcu}', '-w', '-v', file], capture_output=False) 171 172 173 def _flash_dfu_util(details, file): 174 # STM32duino 175 if details[0] == '1eaf' and details[1] == '0003': 176 cli.run(['dfu-util', '-a', '2', '-d', f'{details[0]}:{details[1]}', '-R', '-D', file], capture_output=False) 177 # kiibohd 178 elif details[0] == '1c11' and details[1] == 'b007': 179 cli.run(['dfu-util', '-a', '0', '-d', f'{details[0]}:{details[1]}', '-D', file], capture_output=False) 180 # STM32, APM32, AT32, or GD32V DFU 181 else: 182 cli.run(['dfu-util', '-a', '0', '-d', f'{details[0]}:{details[1]}', '-s', '0x08000000:leave', '-D', file], capture_output=False) 183 184 185 def _flash_wb32_dfu_updater(file): 186 if shutil.which('wb32-dfu-updater_cli'): 187 cmd = 'wb32-dfu-updater_cli' 188 else: 189 return True 190 191 cli.run([cmd, '-t', '-s', '0x08000000', '-D', file], capture_output=False) 192 193 194 def _flash_isp(mcu, programmer, file): 195 programmer = 'usbasp' if programmer == 'usbasploader' else 'usbtiny' 196 # Check if the provided mcu has an avrdude-specific name, otherwise pass on what the user provided 197 mcu = AVRDUDE_MCU.get(mcu, mcu) 198 cli.run(['avrdude', '-p', mcu, '-c', programmer, '-U', f'flash:w:{file}:i'], capture_output=False) 199 200 201 def _flash_mdloader(file): 202 cli.run(['mdloader', '--first', '--download', file, '--restart'], capture_output=False) 203 204 205 def _flash_uf2(file): 206 output = cli.run(['util/uf2conv.py', '--info', file]).stdout 207 if 'UF2 File' not in output: 208 return True 209 210 cli.run(['util/uf2conv.py', '--deploy', file], capture_output=False) 211 212 213 def flasher(mcu, file): 214 # Avoid "expected string or bytes-like object, got 'WindowsPath" issues 215 file = file.as_posix() 216 bl, details = _find_bootloader() 217 # Add a small sleep to avoid race conditions 218 time.sleep(1) 219 if bl == 'atmel-dfu': 220 _flash_atmel_dfu(details, file) 221 elif bl == 'caterina': 222 if _flash_caterina(details, file): 223 return (True, "The Caterina bootloader was found but is not writable. Check 'qmk doctor' output for advice.") 224 elif bl == 'hid-bootloader': 225 if mcu: 226 if _flash_hid_bootloader(mcu, details, file): 227 return (True, "Please make sure 'teensy_loader_cli' or 'hid_bootloader_cli' is available on your system.") 228 else: 229 return (True, "Specifying the MCU with '-m' is necessary for HalfKay/HID bootloaders!") 230 elif bl in {'apm32-dfu', 'at32-dfu', 'gd32v-dfu', 'kiibohd', 'stm32-dfu'}: 231 _flash_dfu_util(details, file) 232 elif bl == 'wb32-dfu': 233 if _flash_wb32_dfu_updater(file): 234 return (True, "Please make sure 'wb32-dfu-updater_cli' is available on your system.") 235 elif bl == 'usbasploader' or bl == 'usbtinyisp': 236 if mcu: 237 _flash_isp(mcu, bl, file) 238 else: 239 return (True, "Specifying the MCU with '-m' is necessary for ISP flashing!") 240 elif bl == 'md-boot': 241 _flash_mdloader(file) 242 elif bl == '_uf2_compatible_': 243 if _flash_uf2(file): 244 return (True, "Flashing only supports uf2 format files.") 245 else: 246 return (True, "Known bootloader found but flashing not currently supported!") 247 248 return (False, None)