from serial import Serial
from tqdm import tqdm
import binascii
import hashlib
import struct
import time
import sys
import os

def if_read(ser, data_len):
    data = bytearray(0)
    received = 0
    while received < data_len:
        tmp = ser.read(data_len - received)
        if len(tmp) == 0:
            break
        else:
            data += tmp
            received += len(tmp)

    if len(data) != data_len:
        return (0, data)
    return (1, data)

def reset(ser):
    ser.setRTS(0)
    time.sleep(0.2)
    reset_cnt = 2
    while reset_cnt > 0:
        ser.setRTS(1)
        time.sleep(0.005)
        ser.setRTS(0)
        time.sleep(0.1)
        ser.setRTS(1)
        time.sleep(0.005)
        ser.setRTS(0)
        time.sleep(0.005)
        reset_cnt -= 1

def handshake(ser):
    ser.setRTS(1)
    time.sleep(0.2)
    ser.setRTS(0)
    time.sleep(0.05)
    ser.setRTS(1)
    ser.setDTR(1)
    time.sleep(0.1)
    ser.setDTR(0)
    time.sleep(0.1)

def expect_ok(ser):
    data = ser.read(2)
    if data[0] != 0x4f or data[1] != 0x4b:
        err = ser.read(2)
        raise ValueError(binascii.hexlify(err))

def expect_data(ser):
    expect_ok(ser)
    len = ser.read(2)
    len = struct.unpack('<h', len)[0]
    data = ser.read(len)
    return data

def cmd_load_seg_header(ser, file):
    header = file.read(0x10)
    ser.write(b'\x17\x00\x10\x00' + header)
    data = expect_data(ser)
    seg_addr, seg_len = struct.unpack('<II', data[0:8])
    print(f'{seg_len} bytes @ {hex(seg_addr)}')
    return seg_len

def cmd_load_seg_data(ser, data):
    ser.write(b'\x18\x00' + struct.pack('<H', len(data)) + data)
    expect_ok(ser)

def cmd_load_boot_header(ser, file):
    header = file.read(0xb0)
    ser.write(b'\x11\x00\xb0\x00' + header)
    expect_ok(ser)

def cmd_check_image(ser):
    ser.write(b'\x19\x00\x00\x00')
    expect_ok(ser)

def cmd_run_image(ser):
    ser.write(b'\x1a\x00\x00\x00')
    expect_ok(ser)

def load_image(ser, file):
    image = open(file, 'rb')
    cmd_load_boot_header(ser, image)
    total = cmd_load_seg_header(ser, image)
    sent = 0
    with tqdm(total=total, unit='byte', unit_scale=True) as pbar:
        while sent != total:
            chunk = image.read(min(total-sent, 4080))
            cmd_load_seg_data(ser, chunk)
            sent = sent + len(chunk)
            pbar.update(len(chunk))
    cmd_check_image(ser)
    cmd_run_image(ser)

def empty_buffer(ser):
    timeout = ser.timeout
    ser.timeout = 0.1
    if_read(ser, 10000)
    ser.timeout = timeout

def send_sync(ser):
    empty_buffer(ser)
    ser.write(b'\x55' * int(0.006 * ser.baudrate / 10))
    expect_ok(ser)

def efl_write_cmd(ser, id, payload = b''):
    plen = len(payload)
    plen_data = struct.pack('<h', plen)
    checksum = struct.pack('<h', sum(plen_data + payload) & 0xff)[0:1]
    data = bytes([id]) + checksum + plen_data + payload
    ser.write(data)

def efl_cmd_read_memory(ser, addr):
    # there is a length parameter here but it doesn't seem to work correctly
    efl_write_cmd(ser, 0x51, struct.pack('<II', addr, 0x4))
    return expect_data(ser)

def efl_cmd_write_memory(ser, addr, data):
    efl_write_cmd(ser, 0x50, struct.pack('<I', len(data)) + data)
    expect_ok(ser)

def efl_cmd_read_jid(ser):
    efl_write_cmd(ser, 0x36)
    return expect_data(ser)

def efl_cmd_flash_erase(ser, addr, len):
    end_addr = addr + len - 1
    efl_write_cmd(ser, 0x30, struct.pack('<II', addr, end_addr))
    timeout = ser.timeout
    ser.timeout = 10.0
    expect_ok(ser)
    ser.timeout = timeout
    print(f'Erased {len} bytes @ {hex(addr)}')

def efl_cmd_flash_write(ser, addr, data):
    efl_write_cmd(ser, 0x31, struct.pack('<I', addr) + data)
    expect_ok(ser)

def efl_cmd_flash_write_check(ser):
    efl_write_cmd(ser, 0x3a)
    expect_ok(ser)

def efl_cmd_flash_xip_read_start(ser):
    efl_write_cmd(ser, 0x60)
    expect_ok(ser)

def efl_cmd_flash_xip_read_sha(ser, addr, len):
    efl_write_cmd(ser, 0x3e, struct.pack('<II', addr, len))
    return expect_data(ser)

def efl_cmd_flash_xip_read_finish(ser):
    efl_write_cmd(ser, 0x61)
    expect_ok(ser)

def efl_cmd_reset(ser):
    efl_write_cmd(ser, 0x21)
    expect_ok(ser)

def efl_program_img(ser, addr, data):
    data_len = len(data)
    efl_cmd_flash_erase(ser, addr, data_len)

    print(f'Programming {data_len} bytes @ {hex(addr)}')
    sent = 0
    with tqdm(total=data_len, unit='byte', unit_scale=True) as pbar:
        while sent != data_len:
            buf_len = min(2048, data_len - sent)
            buf = data[sent:sent + buf_len]
            efl_cmd_flash_write(ser, addr + sent, buf)
            sent = sent + buf_len
            pbar.update(buf_len)
    efl_cmd_flash_write_check(ser)

    sha256sum = hashlib.sha256(data).digest()
    efl_cmd_flash_xip_read_start(ser)
    device_sum = efl_cmd_flash_xip_read_sha(ser, addr, data_len)
    efl_cmd_flash_xip_read_finish(ser)
    if device_sum != sha256sum:
        print('Verification failed')
        print('Host SHA256:', binascii.hexlify(sha256sum))
        print('BL   SHA256:', binascii.hexlify(device_sum))
        return False
    print('Verified by XIP SHA256 hash')
    return True

def prepend_fw_header(img, header_file):
    if img[0:4] == b'BFNP':
        print('Image already has FW header')
        return img
    with open(header_file, 'rb') as f:
        header = f.read()
    img = header + (b'\xFF' * (4096-len(header))) + img
    return img

def get_contrib_path(name):
    sep = os.path.sep
    return os.path.dirname(os.path.realpath(__file__)) + sep + 'contrib' + sep + name

def main():
    if len(sys.argv) < 3:
        print(f'Usage: {sys.argv[0]} <serial port> <firmware bin>')
        sys.exit(1)

    ser = Serial(sys.argv[1], baudrate=500000, timeout=2)
    handshake(ser)
    reset(ser)
    send_sync(ser)
    time.sleep(0.1)
    print('Loading helper binary')
    load_image(ser, get_contrib_path('eflash_loader_40m.bin'))
    time.sleep(0.2)
    print()

    # at this point, the eflash loader binary is running with efl_ commands
    # (which seems to work with a higher baudrate)
    ser.baudrate = 2000000
    send_sync(ser)
    with open(sys.argv[2], 'rb') as f:
        data = f.read()
    data = prepend_fw_header(data, get_contrib_path('bootheader.bin'))
    efl_program_img(ser, 0x10000, data)
    efl_cmd_reset(ser)

if __name__ == "__main__":
    main()

[((27, 4, 27, 19), 'time.sleep', 'time.sleep', ({(27, 15, 27, 18): '(0.2)'}, {}), '(0.2)', False, 'import time\n'), ((42, 4, 42, 19), 'time.sleep', 'time.sleep', ({(42, 15, 42, 18): '(0.2)'}, {}), '(0.2)', False, 'import time\n'), ((44, 4, 44, 20), 'time.sleep', 'time.sleep', ({(44, 15, 44, 19): '(0.05)'}, {}), '(0.05)', False, 'import time\n'), ((47, 4, 47, 19), 'time.sleep', 'time.sleep', ({(47, 15, 47, 18): '(0.1)'}, {}), '(0.1)', False, 'import time\n'), ((49, 4, 49, 19), 'time.sleep', 'time.sleep', ({(49, 15, 49, 18): '(0.1)'}, {}), '(0.1)', False, 'import time\n'), ((68, 24, 68, 55), 'struct.unpack', 'struct.unpack', ({(68, 38, 68, 43): '"""<II"""', (68, 45, 68, 54): 'data[0:8]'}, {}), "('<II', data[0:8])", False, 'import struct\n'), ((116, 16, 116, 39), 'struct.pack', 'struct.pack', ({(116, 28, 116, 32): '"""<h"""', (116, 34, 116, 38): 'plen'}, {}), "('<h', plen)", False, 'import struct\n'), ((212, 10, 212, 57), 'serial.Serial', 'Serial', (), '', False, 'from serial import Serial\n'), ((216, 4, 216, 19), 'time.sleep', 'time.sleep', ({(216, 15, 216, 18): '(0.1)'}, {}), '(0.1)', False, 'import time\n'), ((219, 4, 219, 19), 'time.sleep', 'time.sleep', ({(219, 15, 219, 18): '(0.2)'}, {}), '(0.2)', False, 'import time\n'), ((31, 8, 31, 25), 'time.sleep', 'time.sleep', ({(31, 19, 31, 24): '(0.005)'}, {}), '(0.005)', False, 'import time\n'), ((33, 8, 33, 23), 'time.sleep', 'time.sleep', ({(33, 19, 33, 22): '(0.1)'}, {}), '(0.1)', False, 'import time\n'), ((35, 8, 35, 25), 'time.sleep', 'time.sleep', ({(35, 19, 35, 24): '(0.005)'}, {}), '(0.005)', False, 'import time\n'), ((37, 8, 37, 25), 'time.sleep', 'time.sleep', ({(37, 19, 37, 24): '(0.005)'}, {}), '(0.005)', False, 'import time\n'), ((60, 10, 60, 34), 'struct.unpack', 'struct.unpack', ({(60, 24, 60, 28): '"""<h"""', (60, 30, 60, 33): 'len'}, {}), "('<h', len)", False, 'import struct\n'), ((94, 9, 94, 56), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((123, 29, 123, 58), 'struct.pack', 'struct.pack', ({(123, 41, 123, 46): '"""<II"""', (123, 48, 123, 52): 'addr', (123, 54, 123, 57): '(4)'}, {}), "('<II', addr, 4)", False, 'import struct\n'), ((136, 29, 136, 63), 'struct.pack', 'struct.pack', ({(136, 41, 136, 46): '"""<II"""', (136, 48, 136, 52): 'addr', (136, 54, 136, 62): 'end_addr'}, {}), "('<II', addr, end_addr)", False, 'import struct\n'), ((156, 29, 156, 58), 'struct.pack', 'struct.pack', ({(156, 41, 156, 46): '"""<II"""', (156, 48, 156, 52): 'addr', (156, 54, 156, 57): 'len'}, {}), "('<II', addr, len)", False, 'import struct\n'), ((173, 9, 173, 59), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((210, 8, 210, 19), 'sys.exit', 'sys.exit', ({(210, 17, 210, 18): '(1)'}, {}), '(1)', False, 'import sys\n'), ((55, 25, 55, 46), 'binascii.hexlify', 'binascii.hexlify', ({(55, 42, 55, 45): 'err'}, {}), '(err)', False, 'import binascii\n'), ((144, 29, 144, 52), 'struct.pack', 'struct.pack', ({(144, 41, 144, 45): '"""<I"""', (144, 47, 144, 51): 'addr'}, {}), "('<I', addr)", False, 'import struct\n'), ((182, 16, 182, 36), 'hashlib.sha256', 'hashlib.sha256', ({(182, 31, 182, 35): 'data'}, {}), '(data)', False, 'import hashlib\n'), ((188, 30, 188, 57), 'binascii.hexlify', 'binascii.hexlify', ({(188, 47, 188, 56): 'sha256sum'}, {}), '(sha256sum)', False, 'import binascii\n'), ((189, 30, 189, 58), 'binascii.hexlify', 'binascii.hexlify', ({(189, 47, 189, 57): 'device_sum'}, {}), '(device_sum)', False, 'import binascii\n'), ((205, 27, 205, 53), 'os.path.realpath', 'os.path.realpath', ({(205, 44, 205, 52): '__file__'}, {}), '(__file__)', False, 'import os\n')]