import argparse import os.path import threading import time import subprocess import colorama from colorama import Fore import serial def printc(line, mode): if isinstance(line, bytes): try: line = line.decode() except: pass print(Fore.BLUE + mode.ljust(10) + Fore.RESET + line.rstrip()) def printe(event, line): print(Fore.YELLOW + event.ljust(10) + Fore.RESET + line.rstrip()) class Flasher(threading.Thread): def __init__(self, source, target, control): threading.Thread.__init__(self) self.source = source self.target = target self.control = control def run(self): # Start dd to flash the image to the device printe('flasher', f'Write {self.source} to {self.target}') while not os.path.exists(self.target): time.sleep(0.2) printe('flasher', f'{self.target} appeared') subprocess.run([ 'sudo', 'dd', f'if={self.source}', f'of={self.target}', 'bs=1M', 'oflag=direct,sync', 'status=progress' ]) # Run a filesystem sync to make sure printe('flasher', f'Running sync') subprocess.run(['sudo', 'sync']) # Press the power button printe('flasher', f'Flashing complete, shutting down') self.control.press_power(True) time.sleep(8) self.control.press_power(False) time.sleep(1) printe('flasher', f'Powering on') self.control.press_power(True) class ButtonController: def __init__(self, device): self.device = device self.port = serial.Serial(self.device, 115200) # Disable all gpios self.port.write(b'pbr') self.power_key = False self.bootloader_key = False self.power = False def press_power(self, state): self.power_key = state self.port.write(b'B' if state else b'b') def press_bootloader(self, state): self.bootloader_key = state self.port.write(b'R' if state else b'r') def set_power(self, state): self.power = state self.port.write(b'P' if state else b'p') def main(): parser = argparse.ArgumentParser() parser.add_argument('uartdevice') parser.add_argument('controldevice') parser.add_argument('diskdevice') parser.add_argument('image') parser.add_argument('--baudrate', '-b', default=115200, type=int) args = parser.parse_args() colorama.init() boot_to_ums = True uart = serial.Serial(args.uartdevice, args.baudrate, timeout=5) control = ButtonController(args.controldevice) time.sleep(1) # Hard reset device and launch tow-boot UMS mode. printc('Hard reset (15 seconds)', 'sys') control.press_power(True) time.sleep(15) uart.flushInput() control.press_power(False) printc('Hard reset completed, starting device...', 'sys') time.sleep(1) control.press_power(True) mode = 'unknown' while True: oldmode = mode line = uart.readline() if line.startswith(b'U-Boot SPL '): if oldmode == 'spl': printe("error", "Boot loop detected") mode = 'spl' if control.power_key: control.press_power(False) if boot_to_ums: control.press_bootloader(True) elif line.startswith(b'Tow-Boot '): mode = 'towboot' elif mode == 'towboot' and line.startswith(b'Starting kernel'): mode = 'kernel' elif mode == 'towboot' and line.startswith(b'Allwinner mUSB OTG'): mode = 'ums' # Don't hold the bootloader key on the next boot to get into the new kernel boot_to_ums = False control.press_bootloader(False) # Start writing the image in a second thread thread = Flasher(args.image, args.diskdevice, control) thread.daemon = True thread.start() elif line.startswith(b'[ 0.000000] Booting Linux on physical CPU'): mode = 'kernel' if mode != oldmode: if mode == 'spl': printe("state", "Reset detected") else: printe("state", "Moved from [" + Fore.BLUE + oldmode + Fore.RESET + "] to [" + Fore.BLUE + mode + Fore.RESET + "]") if mode == 'towboot' and b'to enter the boot menu' in line: printe("state", "Boot menu prompt") if mode == 'ums' and len(line) < 4: # Don't try to print the spinner animation in UMS mode continue printc(line, mode=mode) if __name__ == '__main__': main()