sigrok-fwextract-kingst-la2016: handle compressed Qt resources

[sigrok-util.git] / firmware / kingst-la / sigrok-fwextract-kingst-la2016

#!/usr/bin/python3
##
## This file is part of the sigrok-util project.
##
## Copyright (C) 2020 Florian Schmidt <schmidt_florian@gmx.de>
##
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; if not, see <http://www.gnu.org/licenses/>.
##

# This utility extracts FX2 MCU firmware and FPGA bitstream images from
# the "KingstVIS" vendor software. The blobs are kept in Qt resources
# sections. The script was tested with several v3.5 software versions.

import argparse
import os
import sys
import re
import struct
import codecs
import importlib.util
import zlib

# Reuse the parseelf.py module from saleae-logic16.
fwdir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
parseelf_py = os.path.join(fwdir, "saleae-logic16", "parseelf.py")
spec = importlib.util.spec_from_file_location("parseelf", parseelf_py)
parseelf = importlib.util.module_from_spec(spec)
spec.loader.exec_module(parseelf)

class qt_resources(object):
    def __init__(self, program):
        self._elf = parseelf.elf(program)
        self._elf_sections = {} # idx -> data
        self._read_resources()

    def _get_elf_section(self, idx):
        s = self._elf_sections.get(idx)
        if s is None:
            shdr = self._elf.shdrs[idx]
            s = self._elf.read_section(shdr), shdr
            self._elf_sections[idx] = s
        return s

    def _get_elf_sym_value(self, sname):
        sym = self._elf.symtab[sname]
        section, shdr = self._get_elf_section(sym["st_shndx"])
        addr = sym["st_value"] - shdr["sh_addr"]
        value = section[addr:addr + sym["st_size"]]
        if len(value) != sym["st_size"]:
            print("warning: symbol %s should be %d bytes, but in section is only %d bytes" % (
                sname, sym["st_size"], len(value)))
        return value

    # Qt resource stuff.
    def _get_resource_name(self, offset):
        length, i = struct.unpack(">HI", self._res_names[offset:offset + 2 + 4])
        offset += 2 + 4
        name = self._res_names[offset:offset + 2 * length].decode("utf-16be")
        return name

    def _get_resource_data(self, offset):
        length = struct.unpack(">I", self._res_datas[offset:offset + 4])[0]
        offset += 4
        return self._res_datas[offset:offset + length]

    def _read_resources(self):
        RCCFileInfo_Directory = 0x02
        def read_table():
            table = []
            offset = 0
            while offset < len(self._res_struct):
                name_offset, flags = struct.unpack(">IH", self._res_struct[offset:offset+4+2])
                offset += 6
                name = self._get_resource_name(name_offset)
                if flags & RCCFileInfo_Directory:
                    child_count, first_child_offset = struct.unpack(">II", self._res_struct[offset:offset + 4 + 4])
                    offset += 4 + 4
                    table.append((name, flags, child_count, first_child_offset))
                else:
                    country, language, data_offset = struct.unpack(">HHI", self._res_struct[offset:offset + 2 + 2 + 4])
                    offset += 2 + 2 + 4
                    table.append((name, flags, country, language, data_offset))
            return table
        def read_dir_entries(table, which, parents=[]):
            name, flags = which[:2]
            if not flags & RCCFileInfo_Directory:
                raise Exception("not a directory!")
            child_count, first_child = which[2:]
            for i in range(child_count):
                child = table[first_child + i]
                flags = child[1]
                if flags & RCCFileInfo_Directory:
                    read_dir_entries(table, child, parents + [child[0]])
                else:
                    country, language, data_offset = child[2:]
                    full_name = "/".join(parents + [child[0]])
                    self._resources[full_name] = data_offset
                    self._resource_flags[full_name] = flags

        self._res_datas = self._get_elf_sym_value("_ZL16qt_resource_data")
        self._res_names = self._get_elf_sym_value("_ZL16qt_resource_name")
        self._res_struct = self._get_elf_sym_value("_ZL18qt_resource_struct")

        self._resources = {} # res_fn -> res_offset
        self._resource_flags = {} # res_fn -> RCC_flags
        table = read_table()
        read_dir_entries(table, table[0])

    def get_resource(self, res_fn):
        RCCFileInfo_Compressed = 1
        offset = self._resources[res_fn]
        flags = self._resource_flags[res_fn]
        data = self._get_resource_data(offset)
        if flags & RCCFileInfo_Compressed:
            data = zlib.decompress(data[4:])
        return data

    def find_resource_names(self, res_fn_re):
        for key in self._resources.keys():
            m = re.match(res_fn_re, key)
            if m is not None:
                yield key

class res_writer(object):
    def __init__(self, res):
        self.res = res

    def _decode_crc(self, data, decoder=None):
        if decoder is not None:
            data = decoder(data)
        data = bytearray(data)
        crc = zlib.crc32(data) & 0xffffffff
        return data, crc

    def _write_file(self, fn, data):
        with open(fn, "wb") as fp:
            fp.write(data)

    def extract_re(self, resource_pattern, fname_pattern, decoder=None):
        resources = sorted(res.find_resource_names(resource_pattern))
        for resource in resources:
            fname = re.sub(resource_pattern, fname_pattern, resource)
            fname = fname.lower()
            data = self.res.get_resource(resource)
            data, crc = self._decode_crc(data, decoder=decoder)
            self._write_file(fname, data)
            print("resource {rsc}, file {fname}, size {size}, checksum {crc:08x}".format(
                rsc = resource, fname = fname, size = len(data), crc = crc,
            ))

def decode_intel_hex(hexdata):
    """ return list of (address, data)
    """
    datas = []
    # Assume LF-only or CR-LF style end-of-line.
    for line in hexdata.split(b"\n"):
        line = line.strip()
        if chr(line[0]) != ":": raise Exception("invalid line: %r" % line)
        offset = 1
        record = codecs.decode(line[offset:], "hex")
        byte_count, address, record_type = struct.unpack(">BHB", record[:1 + 2 + 1])
        offset = 1 + 2 + 1
        if byte_count > 0:
            data = record[offset:offset + byte_count]
            offset += byte_count
        checksum = record[offset]
        ex_checksum = (~sum(record[:offset]) + 1) & 0xff
        if ex_checksum != checksum: raise Exception("invalid checksum %#x in %r" % (checksum, line))
        if record_type == 0:
            datas.append((address, data))
        elif record_type == 1:
            break
    return datas

def intel_hex_as_blob(hexdata):
    """ return continuous bytes sequence including all data
    (loosing start address here)
    """
    data = decode_intel_hex(hexdata)
    data.sort()
    last = data[-1]
    length = last[0] + len(last[1])
    img = bytearray(length)
    for off, part in data:
        img[off:off + len(part)] = part
    return img

def maybe_intel_hex_as_blob(data):
    if data[0] == ord(":") and max(data) < 127:
          return intel_hex_as_blob(data)
    return data # Keep binary data.

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description = "KingstVIS firmware extraction")
    parser.add_argument('executable', help = "KingstVIS executable file")
    options = parser.parse_args()
    exe_fn = options.executable

    res = qt_resources(exe_fn)
    writer = res_writer(res)

    # Extract all MCU firmware and FPGA bitstream images. The sigrok
    # project may not cover all KingstVIS supported devices. Users can
    # either just copy those files which are strictly required for their
    # specific device (diagnostics will identify those). Or just copy a
    # few more files while some of them remain unused later (their size
    # is small). Seeing which files would be contained is considered
    # valuable, to identify device variants or candidate models.
    writer.extract_re(r"fwusb/fw(.*)", r"kingst-la-\1.fw", decoder=maybe_intel_hex_as_blob)
    writer.extract_re(r"fwfpga/(.*)", r"kingst-\1-fpga.bitstream")