write_boot0.py

import plib
import plib.mmc
from argparse import ArgumentParser
import config
import sys
from utils import switch_device_to_brom_mode

def select_default_port(brom: bool):
    try:
        if brom:
            return config.BR_DEV_PATH
        else:
            return config.PL_DEV_PATH
    except AttributeError:
        if brom:
            print('Define BR_DEV_PATH in config.py (see README.md)')
        else:
            print('Define PL_DEV_PATH in config.py (see README.md)')
        sys.exit(1)


def main():
    parser = ArgumentParser()
    parser.add_argument('file')
    parser.add_argument('-b', '--brom', action='store_true',
                        help='Use this if your device already is in BROM')
    parser.add_argument('--skip-handshake', action='store_true')
    args = parser.parse_args()

    if not args.brom:
        d = plib.Device(config.PL_DEV_PATH, True)
        if not args.skip_handshake:
            d.handshake()
        switch_device_to_brom_mode(d)

    device = plib.Device(config.BR_DEV_PATH, True, brom_mode=True)
    if not args.skip_handshake and args.brom:
        device.handshake()

    # disable caches
    device.run_ext_cmd(0xb1)

    mmc = plib.mmc.Mmc(device, config.MMC0_BASE, True)
    mmc.mmc_init()
    mmc.mmc_set_part(1)

    i = open(args.file, 'rb')
    assert config.BOOT0_SIZE % mmc.block_size == 0
    blocks_left = config.BOOT0_SIZE // mmc.block_size
    for block in range(config.BOOT0_SIZE // mmc.block_size):
        print('%d      ' % blocks_left, end='\r', flush=True)
        mmc.mmc_write_single_block(block, i.read(mmc.block_size))
        blocks_left -= 1
    i.close()

    print('\x1b[0;33mREBOOT\x1b[0m', flush=True)
    device.wdt_reboot()


if __name__ == '__main__':
    main()