flash_lk_payload.py

from argparse import ArgumentParser
import os
import sys
from utils import switch_device_to_brom_mode

BLOCK_SIZE = 512
MAX_PAYLOAD_LEN = 16384*3


def main():
    script_path = os.path.dirname(os.path.abspath(__file__))
    working_dir = os.path.join(script_path, '.')
    os.chdir(working_dir)
    sys.path.append(working_dir)

    import plib
    import plib.mmc
    import config
    assert config.LK_PAYLOAD_OFFSET % BLOCK_SIZE == 0

    parser = ArgumentParser()
    parser.add_argument('payload')
    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()

    with open(args.payload, 'rb') as file:
        payload = file.read()

    #assert len(payload) % BLOCK_SIZE == 0
    while len(payload) % BLOCK_SIZE != 0:
        payload += b'\x00'
    assert len(payload) <= MAX_PAYLOAD_LEN

    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)

    d = plib.Device(config.BR_DEV_PATH, True, brom_mode=True)
    d.handshake()

    d.run_ext_cmd(0xB1)

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

    blocks_left = len(payload) // BLOCK_SIZE
    for block in range(len(payload) // BLOCK_SIZE):
        print('%d       ' % blocks_left, flush=True, end='\r')
        mmc.mmc_write_single_block(
            block + (config.LK_PAYLOAD_OFFSET // BLOCK_SIZE), payload[:BLOCK_SIZE])
        payload = payload[BLOCK_SIZE:]
        blocks_left -= 1

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


if __name__ == '__main__':
    main()