keyboard stuff
auri.ee
1import platform
2import shutil
3import time
4import os
5import signal
6
7import usb.core
8
9from qmk.constants import BOOTLOADER_VIDS_PIDS
10from 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
23AVRDUDE_MCU = {
24 'atmega32a': 'm32',
25 'atmega328p': 'm328p',
26 'atmega328': 'm328',
27}
28# yapf: enable
29
30
31class 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
48def _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
60def _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
74def _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
80def _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
110def _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
139def _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
148def _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
155def _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
173def _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
185def _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
194def _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
201def _flash_mdloader(file):
202 cli.run(['mdloader', '--first', '--download', file, '--restart'], capture_output=False)
203
204
205def _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
213def 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)