Z-Link Firmware Tools

ztools.py

#!/usr/bin/env python3
import os
import mmap
import argparse
from collections import OrderedDict

__title__ = "Z-Link Firmware Tools"
__version__ = "0.1"
__author__ = "Gareth Bryan"
__license__ = "MIT"

FILELIST = OrderedDict([
    ('NBoot.bin',{'maxsize':131072}),
    ('uboot.bin',{'maxsize':786432}),
    ('bootlogo.jpg',{'maxsize':524288}),
    ('kernel',{'maxsize':4194304}),
    ('rootfs.img', {'maxsize':4194304}),
    ('app.img', {'maxsize':6946816}),
    ('updateOK.jpg',{'maxsize':524288})
])

def in_dir(dirname):
    if not os.path.isdir(dirname):
        msg = "{0} does not exist or is not a directory".format(dirname)
        raise argparse.ArgumentTypeError(msg)
    else:
        return dirname

def out_dir(dirname):
    if not os.path.isdir(dirname):
        try:
            os.mkdir(dirname)
        except OSError as err:
            msg = os.strerror(errno.ENOENT)
            raise argparse.ArgumentTypeError(msg)
    else:
        msg = "Directory '{0}' already exists".format(dirname)
        raise argparse.ArgumentTypeError(msg)
    return dirname

def encrypt(data):
    for i in range(len(data)):
        data[i] = data[i] - 106 & 255

def decrypt(data):
    for i in range(data.size()):
        data[i] = data[i] + 106 & 255

def pack(output, inputdir):
    for file_name in [x for x in FILELIST.keys() if x in os.listdir(inputdir)]:
        file_path = (os.path.join(inputdir, file_name))
        file_size = os.path.getsize(file_path)
        file_max_size = FILELIST[file_name]['maxsize']
        if file_size > file_max_size:
            print("Skipping  : {:<12} - File too big: {}>{}".format(
                file_name, file_size, file_max_size))
            continue
        file_header = '{} len={}\n'.format(file_name, file_size).encode()
        try:
            with open(file_path, 'rb') as file_in:
                print("Packing   : {:<12} - {:>7} bytes".format(file_name, file_size))
                mm = mmap.mmap(file_in.fileno(), 0, access=mmap.ACCESS_READ)
                output += file_header
                output += mm
        except IOError as err:
            print(err)

def unpack(mm, outputdir):
    for file_name in FILELIST.keys():
        header_offset =  mm.find(file_name.encode())
        if header_offset != -1:
            # Header is not fixed length/known; zbox parser assumes " len="+10
            file_header = mm[header_offset:header_offset+len(file_name)+15]
            # Header trailing byte is carriage return
            file_header = file_header[:file_header.find(b'\x0a')]
            if b"len=" in file_header:
                file_start = header_offset + len(file_header) + 1
                file_size = int(file_header.split(b'=')[1].decode())
                print("Unpacking : {:<12} - {:>7} bytes at 0x{:0>8x}".format(file_name, file_size, file_start))
                try:
                    with open(os.path.join(outputdir,file_name), 'wb') as file_out:
                        file_out.write(mm[file_start:file_start+file_size])
                except IOError as err:
                    print(err)

def main(args):
    banner = '{} v{}'.format(__title__,__version__)
    print('{}\n{}'.format(banner,'-'*55))
    if args.action == 'pack':
        output = bytearray()
        pack(output, args.inputdir)
        encrypt(output)
        args.updatebin.write(output)
    elif args.action == 'unpack':
        mm = mmap.mmap(args.updatebin.fileno(), 0, access=mmap.ACCESS_COPY)
        decrypt(mm)
        unpack(mm, args.outputdir)
    print('-'*55)
    print('Done')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
    description = "{} v{}".format(__title__,__version__))
    subparsers = parser.add_subparsers(dest='action')
    subparsers.required = True
    parser_pack = subparsers.add_parser('pack')
    parser_pack.add_argument(
                        "inputdir",
                        help = "inputdir",
                        metavar = "inputdir",
                        type=in_dir)
    parser_pack.add_argument(
                        "updatebin",
                        help = "update.bin",
                        metavar = "updatebin",
                        type=argparse.FileType('xb'))
    parser_unpack = subparsers.add_parser('unpack')
    parser_unpack.add_argument(
                        "updatebin",
                        help = "update.bin",
                        metavar = "updatebin",
                        type=argparse.FileType('rb'))
    parser_unpack.add_argument(
                        "outputdir",
                        help = "outputdir",
                        metavar = "outputdir",
                        type=out_dir)

args = parser.parse_args()
main(args)