flashtool.py

#! /usr/bin/python3

from pwn import *
import os
import binascii
import usb

BSL_CMD_CONNECT = 0x0
BSL_CMD_START_DATA = 0x1
BSL_CMD_MIDST_DATA = 0x2
BSL_CMD_END_DATA = 0x3
BSL_CMD_EXEC_DATA = 0x4
BSL_CMD_READ_START = 0x10
BSL_CMD_READ_MIDST = 0x11
BSL_CMD_READ_END = 0x12
BSL_CMD_END_PROCESS = 0x7F
BSL_CMD_POWEROFF = 0x17

BSL_REP_ACK = 0x80
BSL_REP_VER = 0x81

SPRD_EP_IN = 0x85
SPRD_EP_OUT = 0x6
SPRD_MAX_LEN = 1024

DEFAULT_TIMEOUT = 200
UNISOC_VID = g_vid = 0x1782
UNISOC_PID = g_pid = 0x4d00

g_checksum_type = 0x1


class FlashTool:
    def __init__(self):
        self.usb_device = None
        self.timeout = DEFAULT_TIMEOUT
        self.set_checksum_type("brom")

    def device_connect(self):
        usb_dev = usb.core.find(idVendor=g_vid, idProduct=g_pid)
        if(usb_dev):
            usb_dev.set_configuration()
            self.usb_device = usb_dev
            return True
        else:
            return False


    # - Send a HDLC packet to the device
    def send_data(self, data, timeout = None):
        if(timeout == None):
            timeout = self.timeout
        try:
            self.usb_device.write(SPRD_EP_OUT, data, timeout)
            return True
        except usb.core.USBTimeoutError:
            return False
    
    def set_checksum_type(self, _type):
        global g_checksum_type
        if (_type == "brom"):
            g_checksum_type = 0x1
        elif (_type == "fdl"):
            g_checksum_type = 0x2
        else:
            log.error("Err: Incorrect checksum!")

    # - Read a HDLC packet from the device
    def read_data(self, max_len = SPRD_MAX_LEN, timeout = None):
        if (timeout == None): timeout = self.timeout
        try:
            return bytes(self.usb_device.read(SPRD_EP_IN, max_len, timeout))
        except usb.core.USBTimeoutError:
            return None
    
    def read_packet(self):
        try:
            return parse_packet(self.read_data())
        except:
            return None, None, None

    def check_resp(self, resp):
        if (resp == 0x82):
            return "BSL_REP_INVALID_CMD"
        elif(resp == 0x83):
            return "BSL_REP_UNKNOWN_CMD"
        elif(resp == 0x84):
            return "BSL_REP_OPERATION_FAILED"
        elif(resp == 0x85):
            return "BSL_REP_NOT_SUPPORT_BAUDRATE"
        elif(resp == 0x8b):
            return "BSL_REP_VERIFY_ERROR"

    def read_ack(self):
        resp, data, checksum_match = self.read_packet()
        if(resp != BSL_REP_ACK or len(data)):
            log.warn("Err: Failed with error {}".format(self.check_resp(resp)))
            return False
        return True

    # - check the version of software after execution i.e FDL??
    def read_version(self):
        resp, data, checksum_match = self.read_packet()
        if(resp != BSL_REP_VER):
            log.error("Err: Invalid version response!")
            return 
        return data.decode("utf-8", "ignore")

    def cmd_send_ping(self):
        return self.send_data(b"\x7e")
    
    def cmd_connect(self):
        return self.send_data(generate_packet(BSL_CMD_CONNECT))

    # - Send the command to start sending data to a specific address
    def cmd_data_start(self, addr, size):
        data = addr.to_bytes(4, byteorder="big", signed= False)
        data += size.to_bytes(4, byteorder="big", signed=False)
        
        return self.send_data(generate_packet(BSL_CMD_START_DATA, data))

    def cmd_data_send(self, data):
        return self.send_data(generate_packet(BSL_CMD_MIDST_DATA, data))

    def cmd_data_end(self):
        return self.send_data(genrate_packet(BSL_CMD_END_DATA))

    def cmd_data_exec(self):
        return self.send_data(generate_packet(BSL_CMD_EXEC_DATA))

    
    def cmd_poweroff(self):
        return self.send_data(generate_packet(BSL_CMD_POWEROFF))


def xmodem(data: bytes):
    crc = 0x0
    data = bytearray(data)
    _len = len(data)

    for ii in range(0, 1, 2):
        if ii + 1 == _len:
            crc += data[ii]
        else:
            crc += (data[ii] << 0x8) | data[ii + 0x1]

    crc = (crc >> 16) + (crc & 0xffff)
    crc += (crc >> 16)
    return ~crc & 0xffff

#calculate the checksum of the packet
def calc_checksum(data: bytes):
    if (g_checksum_type == 0x1):
        return binascii.crc_hqx(data, 0)
    elif(g_checksum_type == 0x2):
        return xmodem(data)

def hdlc_decode(data: bytes):
    _data = list(data)

    if(_data[0] != 0x7e or _data[-1] != 0x7e):
        log.error("Err: Malformed Packet!")
        return 
    del _data[0]; del _data[-1]

    ii = 0x0
    __data = b""

    while(ii < len(_data)):
        if(_data[ii] == 0x7e and _data[ii + 0x1] == 0x5e):
            __data += bytes([0x7e])
            ii += 0x2
        elif(_data[ii] == 0x7d and _data[ii + 0x1] == 0x5d):
            __data += bytes([0x7d])
            ii += 0x2
        else:
            __data += bytes([_data[ii]])
            ii += 0x1
    
    return __data

def hdlc_encode(data: bytes):
    _data = bytes([0x7e])

    for char in data:
        if (char == 0x7e):
            _data += bytes([0x7d, (char ^ 0x20)])
        elif (char == 0x7d):
            _data += bytes([0x7d, (char ^ 0x29)])
        else:
            _data += bytes([char])
    _data += bytes([0x7e])
    return _data


# This is used to generate a HDLC packet that is signed --> The CRC value is checked XD
def generate_packet(cmd, data=b''):
    pckt = cmd.to_bytes(2, byteorder="big", signed=False)
    pckt += len(data).to_bytes(2, byteorder="big", signed=False)
    
    if(len(data)):
        pckt += data

    pckt += calc_checksum(pckt).to_bytes(2, byteorder="big", signed=False)

    log.info("data-out: {}".format(pckt))
    return hdlc_encode(pckt)


# HDLC decode the packet 
def parse_packet(packet):
    data = hdlc_decode(packet)
    log.info("data-in: {}".format(data))
    if(data == None):
        log.error("Err: Packet parsing failed!")
        return 
    
    command = int.from_bytes(data[0:2], byteorder="big", signed=False)
    _len = int.from_bytes(data[2:4], byteorder="big", signed = False)
    pckt = data[4:-2]

    if(_len != len(pckt)):
        log.error("Err: Packet length Error!")
        return 
    
    chksum = int.from_bytes(data[:-2], byteorder="big", signed=False)
    return command, pckt, chksum