#!/usr/bin/env python

import usb
import sys, time
import os
import unicmd

# global params

UNISOC_VID = 0x1782
UNISOC_PID = 0x4d00
UNISOC_FLASH_BASE_ADDR = 0x10000000
MAX_PKT_SIZE = 1024
bSize = 512 # read block size
genTimeout = 4000

# all main procedures

def connect(vid, pid):
    while True:
        dev = usb.core.find(idVendor=vid, idProduct=pid)
        sys.stdout.write('.')
        sys.stdout.flush()
        if dev is not None:
            print('\nDevice connected')
            break
        time.sleep(0.1)
    dev.set_configuration()
    cfg = dev.get_active_configuration()
    intf = cfg[(0,0)]
    epIn = usb.util.find_descriptor(
        intf,
        custom_match = \
        lambda e: \
            usb.util.endpoint_direction(e.bEndpointAddress) == \
            usb.util.ENDPOINT_IN)
    epOut = usb.util.find_descriptor(
        intf,
        custom_match = \
        lambda e: \
            usb.util.endpoint_direction(e.bEndpointAddress) == \
            usb.util.ENDPOINT_OUT)
    assert epIn is not None
    assert epOut is not None
    return dev, epIn, epOut


def reqonly(packet, fdlBooted = False, noCrc = False):
    packet = unicmd.hdlc_encode(packet, fdlBooted, noCrc)
    dev.write(epOut, packet, genTimeout)

def reqresp(packet, fdlBooted = False, noCrc = False):
    reqonly(packet, fdlBooted, noCrc)
    resp = bytes(dev.read(epIn, bSize, genTimeout))
    return resp

def handshake(fdlBooted = False):
    resp = reqresp(unicmd.cmd_sync(), fdlBooted)
    rcode, rlen, r = unicmd.resp_decode(resp, fdlBooted)
    if len(r):
        print('>', r.decode())
    resp = reqresp(unicmd.cmd_connect(), fdlBooted)
    rcode, rlen, r = unicmd.resp_decode(resp, fdlBooted)
    if len(r):
        print('>', r.decode())

def send_file_to_addr(fname, faddr, fdlBooted = False, flashMode = False, fbs = 1024):
    pSize = MAX_PKT_SIZE
    f = open(fname, 'rb')
    fdata = f.read()
    f.close()
    flen = len(fdata)
    print('Initializing data transfer...')
    if flashMode:
        pSize = fbs
        # faddr is our flash offset in this case
        resp = reqresp(unicmd.cmd_data_start(UNISOC_FLASH_BASE_ADDR + faddr, flen), fdlBooted)
    else:
        resp = reqresp(unicmd.cmd_data_start(faddr, flen), fdlBooted)
    rcode, rlen, r = unicmd.resp_decode(resp, fdlBooted)
    assert rcode == unicmd.BSL_REP_ACK, 'Could not start data transfer, response code is %X' % rcode
    print('Starting data transfer...')
    while fdata:
        buf = fdata[:pSize]
        resp = reqresp(unicmd.cmd_data_send(buf), fdlBooted)
        rcode, rlen, r = unicmd.resp_decode(resp, fdlBooted)
        assert rcode == unicmd.BSL_REP_ACK, 'Something is wrong and response code is %X, block is %s' % (rcode, buf.hex())
        fdata = fdata[pSize:]
        sys.stdout.write('.')
        sys.stdout.flush()
    print('\nEnding data transfer...')
    resp = reqresp(unicmd.cmd_data_end(), fdlBooted)
    rcode, rlen, r = unicmd.resp_decode(resp, fdlBooted)
    if not flashMode or rcode != unicmd.BSL_FLASH_CFG_ERROR: # on flashing, ignore 0xA4 error
        assert rcode == unicmd.BSL_REP_ACK, 'Could not finalize data transfer, response code is %X' % rcode
    print('Data transfer successful')

# readback code implementation

def read_partdata(partid, size, offset):
    t = b''
    reqonly(unicmd.cmd_read_flash(partid, size, offset), True)
    while True:
        xr = bytes(dev.read(epIn, bSize, genTimeout))
        t += xr
        if len(xr) < bSize:
            break
    return t

def read_partition(partid, partsize, partoffset, outfile, rbblocksize):
    outf = open(outfile, 'wb')
    psize = partsize
    offset = partoffset
    print('Dumping %d bytes from partition 0x%X at offset 0x%X to %s...' % (partsize, partid, partoffset, outfile))
    bufsize = rbblocksize
    while psize > 0:
        if psize < bufsize:
            bufsize = psize
        resp = read_partdata(partid, bufsize, offset)
        rcode, rlen, r = unicmd.resp_decode(resp, True)
        outf.write(r)
        sys.stdout.write('.')
        sys.stdout.flush()
        psize -= rlen
        offset += rlen
    outf.flush()
    print('\nPartition dumped!')

def auto_int(x):
    return int(x,0)

# main code start

if __name__ == '__main__': # main app start
    from argparse import ArgumentParser
    rootdir = os.path.dirname(os.path.realpath(__file__))
    parser = ArgumentParser(description='UniFlash: an opensource Unisoc SC6531E/F/M phone flash reader/writer', epilog='(c) Luxferre 2021 --- No rights reserved <https://unlicense.org>')
    parser.add_argument('mode', help='Operation mode (flash/dump)')
    parser.add_argument('file', help='File to read the flash data from or write the dump into')
    parser.add_argument('-p','--partid', type=auto_int, default=0x80000003, help='partition ID for readback (defaults to 0x80000003 that can address full flash space on SC6531E/F/M)')
    parser.add_argument('-s','--start', type=auto_int, default=0, help='start position (in the partition, defaults to 0)')
    parser.add_argument('-l', '--length', type=auto_int, default=0x400000, help='data length in bytes to read/write, defaults to 0x400000')
    parser.add_argument('-bs','--block-size', type=auto_int, default=4096, help='Readback/write block size (in bytes), defaults to 4096')
    parser.add_argument('-dv','--device-vid', type=auto_int, default=UNISOC_VID, help='Override device vendor ID')
    parser.add_argument('-dp','--device-pid', type=auto_int, default=UNISOC_PID, help='Override device product ID')
    parser.add_argument('-fdl1','--fdl1-file', default=rootdir+'/fdls/sc6531efm/nor_fdl1.bin', help='Path to FDL1, defaults to the generic SC6531E/F/M FDL1 shipped with UniDump')
    parser.add_argument('-addr1','--fdl1-addr', type=auto_int, default=0x40004000, help='Address to load FDL1 into, defaults to 0x40004000')
    parser.add_argument('-fdl2','--fdl2-file', default=rootdir+'/fdls/sc6531efm/nor_fdl.bin', help='Path to FDL2, defaults to the generic SC6531E/F/M FDL2 shipped with UniDump')
    parser.add_argument('-addr2','--fdl2-addr', type=auto_int, default=0x14000000, help='Address to load FDL2 into, defaults to 0x14000000')
    args = parser.parse_args()

    is_flash = False
    if args.mode == 'flash':
        is_flash = True

    UNISOC_VID = args.device_vid
    UNISOC_PID = args.device_pid
    fdl1Addr = args.fdl1_addr
    fdl2Addr = args.fdl2_addr
    fdl1Name = args.fdl1_file
    fdl2Name = args.fdl2_file
    outfile = args.file
    partitionId = args.partid
    readbs = args.block_size
    readoffset = args.start
    readlen = args.length

    print('Using FDL1 %s, loading to 0x%X' % (fdl1Name, fdl1Addr))
    print('Using FDL2 %s, loading to 0x%X' % (fdl2Name, fdl2Addr))

    # initial connection
    print('Connect the device %X:%X while holding the bootkey...' % (UNISOC_VID, UNISOC_PID) )
    dev, epIn, epOut = connect(UNISOC_VID, UNISOC_PID)
    handshake()

    def reconnect():
        global dev
        if dev is not None:
            usb.util.dispose_resources(dev)
            time.sleep(0.5)
        dev, epIn, epOut = connect(UNISOC_VID, UNISOC_PID)

    print('Boot mode entered')

    print('Sending FDL1')
    send_file_to_addr(fdl1Name, fdl1Addr)
    print('Starting FDL1')
    resp = reqresp(unicmd.cmd_data_exec(fdl1Addr))
    rcode, rlen, r = unicmd.resp_decode(resp, False)
    if rcode == unicmd.BSL_REP_ACK:
        print('FDL1 started successfully, reconnecting...')
        reconnect()
        handshake(True)

        print('Protocol set up, sending FDL2')
        send_file_to_addr(fdl2Name, fdl2Addr, True)
        print('Starting FDL2')
        resp = reqresp(unicmd.cmd_data_exec(fdl2Addr), True)
        rcode, rlen, r = unicmd.resp_decode(resp, True)
        if rcode == unicmd.BSL_REP_ACK:
            print('FDL2 started successfully!')

            resp = reqresp(unicmd.cmd_sync_full(), True)
            rcode, rlen, r = unicmd.resp_decode(resp, True)
            assert rcode == unicmd.BSL_REP_ACK, 'Could not set the baudrate, response code is %X' % rcode

            print('FDL2 running, may start interacting with flash memory')

            if is_flash:
                print('Writing flash at offset 0x%X from %s...' % (readoffset, outfile))
                send_file_to_addr(outfile, readoffset, True, True, readbs)
                print('Flash memory written, disconnect the device!')
            else:
                read_partition(partitionId, readlen, readoffset, outfile, readbs)
                resp = reqresp(unicmd.cmd_reset(), True)
                rcode, rlen, r = unicmd.resp_decode(resp, True)
                assert rcode == unicmd.BSL_REP_ACK, 'Could not reset the device, response code is %X' % rcode

    if dev is not None: 
        usb.util.dispose_resources(dev)