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, 'time.sleep', 'time.sleep', False, 'import time\n'), (42, 'time.sleep', 'time.sleep', False, 'import time\n'), (44, 'time.sleep', 'time.sleep', False, 'import time\n'), (47, 'time.sleep', 'time.sleep', False, 'import time\n'), (49, 'time.sleep', 'time.sleep', False, 'import time\n'), (68, 'struct.unpack', 'struct.unpack', False, 'import struct\n'), (116, 'struct.pack', 'struct.pack', False, 'import struct\n'), (212, 'serial.Serial', 'Serial', False, 'from serial import Serial\n'), (216, 'time.sleep', 'time.sleep', False, 'import time\n'), (219, 'time.sleep', 'time.sleep', False, 'import time\n'), (31, 'time.sleep', 'time.sleep', False, 'import time\n'), (33, 'time.sleep', 'time.sleep', False, 'import time\n'), (35, 'time.sleep', 'time.sleep', False, 'import time\n'), (37, 'time.sleep', 'time.sleep', False, 'import time\n'), (60, 'struct.unpack', 'struct.unpack', False, 'import struct\n'), (94, 'tqdm.tqdm', 'tqdm', False, 'from tqdm import tqdm\n'), (123, 'struct.pack', 'struct.pack', False, 'import struct\n'), (136, 'struct.pack', 'struct.pack', False, 'import struct\n'), (156, 'struct.pack', 'struct.pack', False, 'import struct\n'), (173, 'tqdm.tqdm', 'tqdm', False, 'from tqdm import tqdm\n'), (210, 'sys.exit', 'sys.exit', False, 'import sys\n'), (55, 'binascii.hexlify', 'binascii.hexlify', False, 'import binascii\n'), (144, 'struct.pack', 'struct.pack', False, 'import struct\n'), (182, 'hashlib.sha256', 'hashlib.sha256', False, 'import hashlib\n'), (188, 'binascii.hexlify', 'binascii.hexlify', False, 'import binascii\n'), (189, 'binascii.hexlify', 'binascii.hexlify', False, 'import binascii\n'), (205, 'os.path.realpath', 'os.path.realpath', False, 'import os\n')]