#!env python3
import time
import sys
import zlib
import os
try:
    import serial
except ImportError:
    print('PySerial must be installed, run `pip3 install pyserial`')
    sys.exit(1)

import struct
from enum import Enum
import binascii
import hashlib
from Crypto.Cipher import AES
import argparse
import math
import keyboard

current_len = 1
def printProgressBar (iteration, total, prefix = '', suffix = '', decimals = 1, length = 100, fill = '█'):
    global current_len
    """
    Call in a loop to create terminal progress bar
    @params:
        iteration   - Required  : current iteration (Int)
        total       - Required  : total iterations (Int)
        prefix      - Optional  : prefix string (Str)
        suffix      - Optional  : suffix string (Str)
        decimals    - Optional  : positive number of decimals in percent complete (Int)
        length      - Optional  : character length of bar (Int)
        fill        - Optional  : bar fill character (Str)
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filledLength = int(length * iteration // total)
    if filledLength == current_len:
        bar = fill * filledLength + '-' * (length - filledLength)
        print('\r%s |%s| %s%% %s' % (prefix, bar, percent, suffix), end = '\r')
        current_len += 1
    # Print New Line on Complete
    if iteration == total:
        print()

def slip_reader(port):
    partial_packet = None
    in_escape = False

    while True:
        waiting = port.inWaiting()
        read_bytes = port.read(1 if waiting == 0 else waiting)
        if read_bytes == b'':
            raise Exception("Timed out waiting for packet %s" % ("header" if partial_packet is None else "content"))
        for b in read_bytes:

            if type(b) is int:
                b = bytes([b])  # python 2/3 compat

            if partial_packet is None:  # waiting for packet header
                if b == b'\xc0':
                    partial_packet = b""
                else:
                    raise Exception('Invalid head of packet (%r)' % b)
            elif in_escape:  # part-way through escape sequence
                in_escape = False
                if b == b'\xdc':
                    partial_packet += b'\xc0'
                elif b == b'\xdd':
                    partial_packet += b'\xdb'
                else:
                    raise Exception('Invalid SLIP escape (%r%r)' % (b'\xdb', b))
            elif b == b'\xdb':  # start of escape sequence
                in_escape = True
            elif b == b'\xc0':  # end of packet
                yield partial_packet
                partial_packet = None
            else:  # normal byte in packet
                partial_packet += b


class UARTCTLResponse:
    class UARTCTLOperation(Enum):
        UC_ECHO = 0xC1
        UC_NOP = 0xC2
        UC_UARTHS_BAUDRATE_SET = 0xC3
        UC_TRUSTED_SHA256_WRITE = 0xC4
        UC_TRUSTED_SHA256_READ = 0xC5
        UC_TRUSTED_SHA256_PROTECT = 0xC6
        UC_GET_TRUSTED_SHA256_PROTECT = 0xC7
        UC_ENABLE_TRUSTED_SHA256_CHECKSUM = 0xC8
        UC_GET_ENABLE_TRUSTED_SHA256_CHECKSUM = 0xC9
        UC_WRITE_AES_KEY = 0xCA
        UC_AES_KEY_COMPARE = 0xCB
        UC_DIS_JTAG = 0xCC
        UC_GET_DIS_JTAG = 0xCD
        UC_FLASH_WRITE = 0xCE
        UC_FLASH_READ = 0xCF
        UC_AES_KEY_PROTECT = 0xD0
        UC_GET_AES_KEY_PROTECT = 0xD1
        UC_DIS_AES_KEY_COMPARE = 0xD2
        UC_GET_DIS_AES_KEY_COMPARE = 0xD3
        UC_DIS_ISP = 0xD4
        UC_GET_DIS_ISP = 0xD5
        UC_AES_IV_WRITE = 0xD6
        UC_AES_IV_READ = 0xD7
        UC_AES_IV_PROTECT = 0xD8
        UC_GET_AES_IV_PROTECT = 0xD9
        UC_FUNC_FIRMWARE_CIPHER_DISABLE = 0xDA
        UC_GET_FUNC_FIRMWARE_CIPHER_DISABLE = 0xDB
        UC_READ_OTP = 0xDC
        UC_DEBUG_INFO = 0xDD

    class ErrorCode(Enum):
        UC_RET_DEFAULT = 0
        UC_RET_OK = 0xE0
        UC_RET_BAD_DATA_LEN = 0xE1
        UC_RET_BAD_DATA_CHECKSUM = 0xE2
        UC_RET_INVALID_COMMAND = 0xE3
        UC_RET_OTP_ERROR = 0xE4

    @staticmethod
    def parse(data):
        op = data[0]
        reason = data[1]
        text = ''
        try:
            if UARTCTLResponse.UARTCTLOperation(op) == UARTCTLResponse.UARTCTLOperation.UC_DEBUG_INFO:
                text = data[2:].decode()
        except ValueError:
            print('Warning: recv unknown op', op)

        return (op, reason, text)

class ISPResponse:
    class ISPOperation(Enum):
        ISP_ECHO = 0xC1
        ISP_NOP = 0xC2
        ISP_MEMORY_WRITE = 0xC3
        ISP_MEMORY_READ = 0xC4
        ISP_MEMORY_BOOT = 0xC5
        ISP_DEBUG_INFO = 0xD1

    class ErrorCode(Enum):
        ISP_RET_DEFAULT = 0
        ISP_RET_OK = 0xE0
        ISP_RET_BAD_DATA_LEN = 0xE1
        ISP_RET_BAD_DATA_CHECKSUM = 0xE2
        ISP_RET_INVALID_COMMAND = 0xE3

    @staticmethod
    def parse(data):
        op = data[0]
        reason = data[1]
        text = ''
        try:
            if ISPResponse.ISPOperation(op) == ISPResponse.ISPOperation.ISP_DEBUG_INFO:
                text = data[2:].decode()
        except ValueError:
            print('Warning: recv unknown op', op)

        return (op, reason, text)

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]

class Uart_control:
    def __init__(self, port='/dev/ttyUSB1', baudrate=115200):
        # configure the serial connections (the parameters differs on the device you are connecting to)
        self._port = serial.Serial(
            port=port,
            baudrate=baudrate,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS
        )
        print('BAUDRATE', baudrate)

        self._port.isOpen()
        self._slip_reader = slip_reader(self._port)

    """ Read a SLIP packet from the serial port """

    def read(self):
        return next(self._slip_reader)

    """ Write bytes to the serial port while performing SLIP escaping """

    def write(self, packet):
        buf = b'\xc0' \
              + (packet.replace(b'\xdb', b'\xdb\xdd').replace(b'\xc0', b'\xdb\xdc')) \
              + b'\xc0'
        #print('[WRITE]', binascii.hexlify(buf))
        return self._port.write(buf)

    def read_loop(self):
        out = b''
        # while self._port.inWaiting() > 0:
        #     out += self._port.read(1)

        # print(out)
        while 1:
            sys.stdout.write('[RECV] raw data: ')
            sys.stdout.write(binascii.hexlify(self._port.read(1)).decode())
            sys.stdout.flush()

    def recv_one_return(self):
        data = b''
        # find start boarder
        #sys.stdout.write('[RECV one return] raw data: ')
        while 1:
            c = self._port.read(1)
            #sys.stdout.write(binascii.hexlify(c).decode())
            sys.stdout.flush()
            if c == b'\xc0':
                break

        in_escape = False
        while 1:
            c = self._port.read(1)
            #sys.stdout.write(binascii.hexlify(c).decode())
            sys.stdout.flush()
            if c == b'\xc0':
                break

            elif in_escape:  # part-way through escape sequence
                in_escape = False
                if c == b'\xdc':
                    data += b'\xc0'
                elif c == b'\xdd':
                    data += b'\xdb'
                else:
                    raise Exception('Invalid SLIP escape (%r%r)' % (b'\xdb', c))
            elif c == b'\xdb':  # start of escape sequence
                in_escape = True
            else:
                data += c

        #sys.stdout.write('\n')
        return data

    def recv_debug(self):
        op, reason, text = UARTCTLResponse.parse(self.recv_one_return())
        print('[RECV] op:', UARTCTLResponse.UARTCTLOperation(op).name, 'reason:', UARTCTLResponse.ErrorCode(reason).name)
        if text:
            print('-' * 30)
            print(text)
            print('-' * 30)
        if UARTCTLResponse.ErrorCode(reason) not in (UARTCTLResponse.ErrorCode.UC_RET_DEFAULT, UARTCTLResponse.ErrorCode.UC_RET_OK):
            print('Failed, retry, errcode=', hex(reason))
            return False
        return True

    def recv_debug_noprint(self):
        op, reason, text = UARTCTLResponse.parse(self.recv_one_return())
        if text:
            print('-' * 30)
            print(text)
            print('-' * 30)
        if UARTCTLResponse.ErrorCode(reason) not in (UARTCTLResponse.ErrorCode.UC_RET_DEFAULT, UARTCTLResponse.ErrorCode.UC_RET_OK):
            print('Failed, retry, errcode=', hex(reason))
            return False
        return True

    def isp_recv_debug(self):
        op, reason, text = ISPResponse.parse(self.recv_one_return())
        if text:
            print('-' * 30)
            print(text)
            print('-' * 30)
        if ISPResponse.ErrorCode(reason) not in (ISPResponse.ErrorCode.ISP_RET_DEFAULT, ISPResponse.ErrorCode.ISP_RET_OK):
            print('Failed, retry, errcode=', hex(reason))
            return False
        return True

    def flash_read(self, address=0x0, data_len=1024):
        out = struct.pack('II', address, data_len)
        sp_send_data = bytearray(data_len)
        crc32_checksum = struct.pack('I', binascii.crc32(out + sp_send_data) & 0xFFFFFFFF)
        out = struct.pack('HH', 0xcf, 0x00) + crc32_checksum + out + sp_send_data  # op: ISP_MEMORY_WRITE: 0xc3
        sent = self.write(out)
        data = self.recv_one_return()
        checksum = data[4:8]
        data_len = int.from_bytes(data[12:16], byteorder='little', signed=True)
        text = data[16:data_len+16]
        self.recv_debug_noprint()
        return text

    def change_baudrate(self, baudrate):
        out = struct.pack('III', 0, 4, baudrate)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        out = struct.pack('HH', 0xc3, 0x00) + crc32_checksum + out
        self.write(out)
        time.sleep(0.05)
        self._port.baudrate = baudrate

    def trusted_sha256_write(self, trusted_sha256: bytes = None):
        if trusted_sha256:
            sha256_len = len(trusted_sha256)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', sha256_len) + trusted_sha256
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xc4, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def trusted_sha256_read(self):
        out = struct.pack('I', 0x00) + struct.pack('I', 0x20) + bytearray(32)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xc5, 0x00) + crc32_checksum + out
        sent = self.write(data)
        data = self.recv_one_return()
        checksum = data[4:8]
        data_len = int.from_bytes(data[12:16], byteorder='little', signed=True)
        text = data[16:data_len + 16]
        self.recv_debug()
        return text

    def trusted_sha256_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xc6, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_trusted_sha256_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xc7, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def enable_trusted_sha256_checksum(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xc8, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_enable_trusted_sha256_checksum(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xc9, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def write_aes_key(self,aes_key: bytes = None):
        if aes_key:
            key_len = len(aes_key)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', key_len) + aes_key
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xca, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def aes_key_compare(self,aes_key: bytes = None):
        if aes_key:
            key_len = len(aes_key)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', key_len) + aes_key
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xcb, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def disable_jtag(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xcc, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_disable_jtag(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xcd, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def aes_key_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd0, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_aes_key_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd1, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def disable_aes_key_compare(self):
        if aes_key:
            key_len = len(aes_key)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xd2, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def get_disable_aes_key_compare(self):
        if aes_key:
            key_len = len(aes_key)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xd3, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def disable_isp(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd4, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_disable_isp(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd5, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def aes_iv_write(self, aes_iv: bytes = None):
        if aes_iv:
            iv_len = len(aes_iv)
            out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', iv_len) + aes_iv
            crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
            data = struct.pack('HH', 0xd6, 0x00) + crc32_checksum + out
            sent = self.write(data)
            self.recv_debug()

    def aes_iv_read(self):
        out = struct.pack('I', 0x00) + struct.pack('I', 0x10) + bytearray(16)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd7, 0x00) + crc32_checksum + out
        sent = self.write(data)
        data = self.recv_one_return()
        checksum = data[4:8]
        data_len = int.from_bytes(data[12:16], byteorder='little', signed=True)
        text = data[16:data_len+16]
        self.recv_debug()
        return text

    def aes_iv_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd8, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_aes_iv_protect(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xd9, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def func_firmware_cipher_disable(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xda, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

    def get_func_firmware_cipher_disable(self):
        out = struct.pack('HH', 0x00, 0x00) + struct.pack('I', 0x00)
        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)
        data = struct.pack('HH', 0xdb, 0x00) + crc32_checksum + out
        sent = self.write(data)
        self.recv_debug()

#define OTP_COMMON_DATA_ADDR    0x00000000U
#define OTP_SYSTEM_DATA_ADDR    0x00003AD0U
#define OTP_BISR_DATA_ADDR      0x00003DD0U
#define OTP_CHIPID_DATA_ADDR    0x00003D90U
    def chipid_read(self, address=0x3D90, data_len=32):
        out = struct.pack('II', address, data_len)
        sp_send_data = bytearray(data_len)
        crc32_checksum = struct.pack('I', binascii.crc32(out + sp_send_data) & 0xFFFFFFFF)
        out = struct.pack('HH', 0xdc, 0x00) + crc32_checksum + out + sp_send_data
        sent = self.write(out)
        data = self.recv_one_return()
        checksum = data[4:8]
        data_len = int.from_bytes(data[12:16], byteorder='little', signed=True)
        text = data[16:data_len+16]
        self.recv_debug_noprint()
        return text

    def firmware_burn(self, data, aes_key: bytes = None, aes_iv: bytes = None, address=0x4000):
        print('[DEBUG]firmware_burn: aeskey=', aes_key, 'aesiv=', aes_iv)
        global current_len
        current_len = 0
        DATAFRAME_SIZE = 4096
        # 固件加上头
        # 格式: SHA256(after)(32bytes) + AES_CIPHER_FLAG (1byte) + firmware_size(4bytes) + firmware_data

        aes_cipher_flag = b'\x01' if aes_key else b'\x00'
        # 加密
        firmware_bin = data

        if aes_key:
            firmware_len = len(firmware_bin)
            pad_len = firmware_len % 16
            if pad_len != 0:
                firmware_bin += bytearray(16 - pad_len)
            firmware_bin = AES.new(key=aes_key, mode=AES.MODE_CBC, iv=aes_iv).encrypt(firmware_bin)
        firmware_len = len(firmware_bin)
        data = aes_cipher_flag + struct.pack('I', firmware_len) + firmware_bin
        sha256_hash = hashlib.sha256(data).digest()
        firmware_with_header = data + sha256_hash
        total_chunk = math.ceil(len(firmware_with_header) / DATAFRAME_SIZE)
        data_chunks = chunks(firmware_with_header, DATAFRAME_SIZE)  # 1kb for a sector
        for i, chunk in enumerate(data_chunks):
            chunk = chunk.ljust(DATAFRAME_SIZE, b'\x00')  # align by 4kb
            while 1:
                #print('[INFO] sending chunk', i, '@address', hex(address), 'chunklen', len(chunk))
                out = struct.pack('II', address, len(chunk))

                crc32_checksum = struct.pack('I', binascii.crc32(out + chunk) & 0xFFFFFFFF)

                out = struct.pack('HH', 0xce, 0x00) + crc32_checksum + out + chunk
                sent = self.write(out)
                #print('[INFO]', 'sent', sent, 'bytes', 'checksum', binascii.hexlify(crc32_checksum).decode())

                address += len(chunk)

                if self.recv_debug_noprint():
                    break
            printProgressBar(i + 1, total_chunk, prefix='Downloading:', suffix='Complete', length=50)

    def firmware_read(self, path, data_len=0, aes_key: bytes = None, aes_iv:bytes = None, address=0x4000):
        print('[DEBUG]firmware_read: path=', path, 'data_len=', data_len)
        global current_len
        current_len = 0
        DATAFRAME_SIZE = 4096

        with open(path, 'wb')as fp:
            read_addr = address
            read_data = b''
            while data_len > 0:
                if data_len > 4096:
                    read_len = 4096
                else:
                    read_len = data_len
                read_data = self.flash_read(read_addr, read_len)                       #读取flash的数据，用于与烧录的文件比较，保证烧录的正确性

                fp.write(read_data)
                data_len -= read_len
                read_addr += read_len

    def model_burn(self, data, aes_key: bytes = None, aes_iv:bytes = None, address=0x800000):
        print('[DEBUG]model_burn: aeskey=', aes_key, 'aesiv=', aes_iv, 'address=%#x'%address)
        global current_len
        current_len = 0
        DATAFRAME_SIZE = 4096
        aes_cipher_flag = b'\x01' if aes_key else b'\x00'
        model_bin = data
        if aes_key:
            model_len = len(model_bin)
            pad_len = model_len % 16
            if pad_len != 0:
                model_bin += bytearray(16 - pad_len)
            model_bin = AES.new(key=aes_key, mode=AES.MODE_CBC, iv=aes_iv).encrypt(model_bin)
        model_len = len(model_bin)
        data = aes_cipher_flag + struct.pack('I', model_len) + model_bin
        total_chunk = math.ceil(len(data) / DATAFRAME_SIZE)
        data_chunks = chunks(data, DATAFRAME_SIZE)  # 1kb for a sector
        for i, chunk in enumerate(data_chunks):
            chunk = chunk.ljust(DATAFRAME_SIZE, b'\x00')  # align by 4kb
            while 1:
                #print('[INFO] sending chunk', i, '@address', hex(address), 'chunklen', len(chunk))
                out = struct.pack('II', address, len(chunk))

                crc32_checksum = struct.pack('I', binascii.crc32(out + chunk) & 0xFFFFFFFF)

                out = struct.pack('HH', 0xce, 0x00) + crc32_checksum + out + chunk
                sent = self.write(out)
                #print('[INFO]', 'sent', sent, 'bytes', 'checksum', binascii.hexlify(crc32_checksum).decode())
                address += len(chunk)
                if self.recv_debug_noprint():
                    break
            printProgressBar(i + 1, total_chunk, prefix='Downloading:', suffix='Complete', length=50)

    def model_read(self, path, data_len=0, aes_key: bytes = None, aes_iv:bytes = None, address=0x800000):
        print('[DEBUG]model_read: path=', path, 'data_len=', data_len, 'address=%#x'%address)
        global current_len
        current_len = 0
        DATAFRAME_SIZE = 4096

        with open(path, 'wb')as fp:
            read_addr = address
            read_data = b''
            while data_len > 0:
                if data_len > 4096:
                    read_len = 4096
                else:
                    read_len = data_len
                read_data = self.flash_read(read_addr, read_len)                       #读取flash的数据，用于与烧录的文件比较，保证烧录的正确性

                fp.write(read_data)
                data_len -= read_len
                read_addr += read_len


    def kboot_burn(self, kboot_bin: bytes, aes_key: bytes = None, trusted_sha256=None):
        print('[DEBUG] kboot_burn DEBUG: aeskey=', aes_key)
        address = 0x0
        DATAFRAME_SIZE = 4096
        aes_cipher_flag = b'\x01' if aes_key else b'\x00'
        # 加密
        kboot_len = len(kboot_bin)
        # 64bytes对齐
        pad_len = kboot_len % 64
        if pad_len != 59:
            if pad_len > 59:
                pad_len = 123 - pad_len
                kboot_bin += bytearray(pad_len)
            else:
                pad_len = 59 - pad_len
                kboot_bin += bytearray(pad_len)
            kboot_len += pad_len

        if aes_key:
            padded = kboot_bin + b'\x00' * 5  # zero pad
            kboot_bin = AES.new(key=aes_key, mode=AES.MODE_CBC, iv=b'\x00' * 16).encrypt(padded) #16bytes对齐

        kboot_len = len(kboot_bin)
        data = aes_cipher_flag + struct.pack('I', kboot_len) + kboot_bin

        sha256_hash = hashlib.sha256(data).digest()
        #这四条命令如果打开，则这个芯片就只能启动这个kboot了，因为bootrom里面会检查这个sha256的值。
        #trusted sha256的机制：
        #   1.芯片的bootrom启动时会读取otp的功能寄存器设置，判断是否使能了trusted sha256 check功能。
        #   2.如果使能了，则从otp的block48、47读取trusted sha256，与本次烧录的固件的sha256进行比较，如果一致则正常启动，否则启动失败。
        #   3.由此可见，这个trusted sha256可以保证启动的固件唯一。如果对当前kboot使能了该功能，在kboot中添加用户自定义的licence时绝对安全的。
        #   4.如果使能了这个功能，芯片的启动程序就不能再被修改了，所以在调试测试时，千万不要打开这个功能。
        #   5.这个功能的设置参考kboot_burn的注释。
        if trusted_sha256 == 'Y':
            print("before trusted_sha256_read: ", binascii.hexlify(sha256_hash))
            self.trusted_sha256_write(trusted_sha256=sha256_hash)
            print("after trusted_sha256_read: ", binascii.hexlify(self.trusted_sha256_read()))
            self.trusted_sha256_protect()
            self.enable_trusted_sha256_checksum()
        kboot_with_header = data + sha256_hash
        total_chunk = math.ceil(len(kboot_with_header) / DATAFRAME_SIZE)
        data_chunks = chunks(kboot_with_header, DATAFRAME_SIZE)  # 1kb for a sector
        for i, chunk in enumerate(data_chunks):
            chunk = chunk.ljust(DATAFRAME_SIZE, b'\x00')  # align by 4kb
            while 1:
                out = struct.pack('II', address, len(chunk))

                crc32_checksum = struct.pack('I', binascii.crc32(out + chunk) & 0xFFFFFFFF)

                out = struct.pack('HH', 0xce, 0x00) + crc32_checksum + out + chunk
                sent = self.write(out)

                address += len(chunk)

                if self.recv_debug_noprint():
                    break

    def flash_dataframe(self, data, address=0x80000000):
        DATAFRAME_SIZE = 1024
        data_chunks = chunks(data, DATAFRAME_SIZE)
        total_chunk = math.ceil(len(data)/DATAFRAME_SIZE)

        for n, chunk in enumerate(data_chunks):
            while 1:
                out = struct.pack('II', address, len(chunk))
                crc32_checksum = struct.pack('I', binascii.crc32(out + chunk) & 0xFFFFFFFF)
                out = struct.pack('HH', 0xc3, 0x00) + crc32_checksum + out + chunk  # op: ISP_MEMORY_WRITE: 0xc3
                sent = self.write(out)
                address += len(chunk)
                if self.isp_recv_debug():
                    break

    def install_flash_bootloader(self, data):
        self.flash_dataframe(data, address=0x80000000)

    def boot(self, address=0x80000000):
        out = struct.pack('II', address, 0)

        crc32_checksum = struct.pack('I', binascii.crc32(out) & 0xFFFFFFFF)

        out = struct.pack('HH', 0xc5, 0x00) + crc32_checksum + out  # op: ISP_MEMORY_WRITE: 0xc3
        self.write(out)

    def greeting(self):
        self._port.write(b'\xc0\xc2\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xc0')
        op, reason, text = ISPResponse.parse(self.recv_one_return())
        print('k210 return op:', ISPResponse.ISPOperation(op).name, 'reason:', ISPResponse.ErrorCode(reason).name)

    def reset_to_isp_kd233(self):
        # rts -> BOOT  dtr -> RST
        self._port.setDTR(False)
        self._port.setRTS(False)
        time.sleep(0.1)
        self._port.setRTS(True)
        self._port.setDTR(True)
        time.sleep(0.1)
        self._port.setDTR(False)
        self._port.setRTS(False)
        time.sleep(0.1)

    def reset_to_boot_kd233(self):
        # rts -> BOOT  dtr -> RST
        self._port.setDTR(False)
        self._port.setRTS(False)
        time.sleep(0.1)
        self._port.setRTS(False)
        self._port.setDTR(True)
        time.sleep(0.1)
        self._port.setDTR(False)
        self._port.setRTS(False)
        time.sleep(0.1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--device", help="UART device", default="/dev/ttyUSB0")
    parser.add_argument("-b", "--baudrate", type=int, help="UART baudrate when flash read/write", default=1500000)
    parser.add_argument("-u", "--uart_control", help="uart control bin path", required=True, default=None)
    parser.add_argument("-m", "--model", help="model bin path", required=False, default=None)
    parser.add_argument("-maddr", "--model_addr", help="model burn address", default=None)
    parser.add_argument("-rf", "--read_firmware", help="read firmware bin path", required=False, default=None)
    parser.add_argument("-rm", "--read_model", help="read model bin path", required=False, default=None)
    parser.add_argument("-rc", "--read_chipid", help="read chip id", required=False, default=None)
    parser.add_argument("-k", "--aes_key", help="aes key", required=False, default=None)
    parser.add_argument("-i", "--aes_iv", help="aes iv", required=False, default=None)
    parser.add_argument("-djtag", "--dis_jtag", help="disable jtag", required=False, default=None)
    parser.add_argument("-dcipher", "--dis_cipher", help="disable cipher", required=False, default=None)
    parser.add_argument("-disp", "--dis_isp", help="disable isp", required=False, default=None)
    parser.add_argument("-tsha", "--trusted_sha256", help="enable trusted sha256", required=False, default=None)

    args = parser.parse_args()
    uart_control = Uart_control(port=args.device, baudrate=115200)

    retry_count = 0
    while True:
        try:
            retry_count = retry_count + 1
            if retry_count > 15:
                print("retry overtime out")
            try:
                uart_control.reset_to_isp_kd233()
                uart_control.greeting()
                break
            except TimeoutError:
                pass
        except Exception as e:
            raise_exception(Exception("Greeting fail, check serial port (" + str(e) + ")"))
    # print("begin load uart control")
    #烧录uart_control.bin到内存运行
    if args.uart_control:
        print("load uart control")
        uart_control_bin = args.uart_control
        uart_control.install_flash_bootloader(open(uart_control_bin, 'rb').read())
        uart_control.boot()
        time.sleep(0.3)
    else:
        raise ValueError('uart control bin miss')
    #读取chipid
    if args.read_chipid == 'Y':
        print("chipid read: ", binascii.hexlify(uart_control.chipid_read()))
    #烧写aes key与aes iv
    if args.aes_key:
        aes_key = binascii.a2b_hex(args.aes_key)
        if len(aes_key) != 16:
            raise ValueError('AES key must by 16 bytes')
        if args.aes_iv:
            aes_iv = binascii.a2b_hex(args.aes_iv)
            if len(aes_iv) != 16:
                raise ValueError('AES iv must by 16 bytes')
        else:
            aes_iv = b'\x00'*16
    else:
        aes_key = None
        aes_iv = None

    if aes_key:
        uart_control.write_aes_key(aes_key=aes_key)   #写入aes key
        uart_control.aes_iv_write(aes_iv=aes_iv)      #写入aes iv
        uart_control.aes_key_protect()                #aes key写保护：otp的值可以从0改成1，不能从1写成0，如果写保护了，则otp的值固定不能被改变了。
        uart_control.disable_aes_key_compare()        #禁用aes key比较
        uart_control.aes_iv_protect()                 #aes iv写保护
        print("aes_iv_read: ", binascii.hexlify(uart_control.aes_iv_read()))    #读取设置的aes iv，用于判断保证烧写的iv正确。
    #禁用JTAG功能
    if args.dis_jtag == 'Y':
        uart_control.disable_jtag()                   #禁用jtag
    #禁用加密功能
    if args.dis_cipher == 'Y':
        uart_control.func_firmware_cipher_disable()   #禁用解密，如果设置上，则芯片的bootrom启动时不再执行解密操作。


    #设置串口波特率为115200
    if args.baudrate != 115200:
        uart_control.change_baudrate(115200)

    retry_count = 0
    while True:
        try:
            retry_count = retry_count + 1
            if retry_count > 15:
                print("retry overtime out")
            try:
                uart_control.reset_to_boot_kd233()
                uart_control.greeting()
                break
            except TimeoutError:
                pass
        except Exception as e:
            raise_exception(Exception("Greeting fail, check serial port (" + str(e) + ")"))

    time.sleep(3)
    retry_count = 0
    while True:
        try:
            retry_count = retry_count + 1
            if retry_count > 15:
                print("retry overtime out")
            try:
                uart_control.reset_to_isp_kd233()
                uart_control.greeting()
                break
            except TimeoutError:
                pass
        except Exception as e:
            raise_exception(Exception("Greeting fail, check serial port (" + str(e) + ")"))

    #烧录uart_control.bin到内存运行，用于double check，因为otp的设置需要重启才能生效。
    if args.uart_control:
        uart_control_bin = args.uart_control
        uart_control.install_flash_bootloader(open(uart_control_bin, 'rb').read())
        uart_control.boot()
        time.sleep(0.3)

    #读取aes key与aes iv相关的otp设置
    if aes_key:
        print("aes_iv_read: ", binascii.hexlify(uart_control.aes_iv_read()))    #读取设置的aes iv，用于判断保证烧写的iv正确。
        uart_control.get_aes_key_protect()                    #读取aes key写保护，如果返回ISP_RET_OK则说明该功能已经设置正确。
        uart_control.get_disable_aes_key_compare()            #读取aes key比较禁用，如果返回ISP_RET_OK则说明该功能已经设置正确。
        uart_control.get_aes_iv_protect()                     #读取aes iv写保护，如果返回ISP_RET_OK则说明该功能已经设置正确。
    #读取jtag禁用
    if args.dis_jtag == 'Y':
        uart_control.get_disable_jtag()                       #如果返回ISP_RET_OK则说明该功能已经设置正确。
    #读取禁用解密，
    if args.dis_cipher == 'Y':
        uart_control.get_func_firmware_cipher_disable()     #如果返回ISP_RET_OK则说明该功能已经设置正确。
    #读取trusted sha256校验使能/写保护，在烧录kboot.bin时会设置trusted sha256
    if args.trusted_sha256 == 'Y':
        uart_control.get_enable_trusted_sha256_checksum()  #如果返回ISP_RET_OK则说明该功能已经设置正确。
        uart_control.get_trusted_sha256_protect()          #如果返回ISP_RET_OK则说明该功能已经设置正确。

    uart_control.reset_to_boot_kd233()
