amba_fwpak.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

""" Ambarella Firmware Packer tool.

Extracts and re-packs partitions from Ambarella firmware module.
"""

# Copyright (C) 2016,2017 Mefistotelis <mefistotelis@gmail.com>
# Copyright (C) 2018 Original Gangsters <https://dji-rev.slack.com/>
#
# 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/>.

__version__ = "0.1.2"
__author__ = "Mefistotelis @ Original Gangsters"
__license__ = "GPL"

import argparse
import sys
if sys.version_info < (3, 0):
    # All checksums would have to be computed differently on Python 2.x
    # due to differences in types
    raise NotImplementedError('Python version 3 or newer is required.')
import mmap
import os
import re
import zlib
import configparser
import itertools
from ctypes import c_char, c_ubyte, c_uint, sizeof, LittleEndianStructure
from time import gmtime, strftime


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def raise_or_warn(po, ex):
    """ Raise exception, unless force-continue parameter was used.
    """
    if (po.force_continue):
        eprint("{:s}: Warning: {:s} Continuing anyway.".format(po.fwmdlfile, str(ex)))
    else:
        raise ex


part_entry_type_id = ["sys", "dsp_fw", "rom_fw", "lnx", "rfs"]
part_entry_type_name = ["System Software", "DSP uCode", "System ROM Data", "Linux Kernel", "Linux Root FS"]

# The Ambarella firmware file consists of 3 elements:
# 1. Main header, containing array of partitions
# 2. Partition header, before each partition
# 3. Partition data, for each partition
#
# The Main header is made of:
# - model_name - text description of the device model
# - ver_info - version info, set to 0 in DJI camera FW
# - crc32 - cummulative checksum of all modules with headers, equal to last
#   module cummulative checksum xor -1
# - fw module entries - array of FwModEntry, with amount of entries hard-coded
#   for specific component; the crc32 here is a cummulative checksum of data
#   with header, and initial value of -1
# - partition sizes - array of int, sizes of the partitions in partition table,
#   with 15 entries (amount of partitions is larger than the amount of modules)
#
# For a specific component, main header of Ambarella firmware has constant
# size. But DJI uses several camera types all developed on Ambarella - amount
# of module entries is different for each of them. To guess the amount of
# modules in a file, we're assuming the partition sizes are multiplication of
# 1024. This way we can detect beginning of the sizes array, as crc value and
# partition length are very unlikely to both divide by 1024.


class FwModA9Header(LittleEndianStructure):
    _pack_ = 1
    _fields_ = [
      ('model_name', c_char * 32),
      ('ver_info', c_uint),
      ('crc32', c_uint),
    ]

    def dict_export(self):
        d = dict()
        for (varkey, vartype) in self._fields_:
            d[varkey] = getattr(self, varkey)
        varkey = 'ver_info'
        d[varkey] = "{:d}.{:d}-{:04X}".format((d[varkey] >> 24) & 255,
          (d[varkey] >> 16) & 255, (d[varkey]) & 65535)
        varkey = 'crc32'
        d[varkey] = "{:08X}".format(d[varkey])
        return d

    def ini_export(self, fp):
        d = self.dict_export()
        fp.write("# Ambarella Firmware Packer module header file. Loosly based on AFT format.\n")
        fp.write(strftime("# Generated on %Y-%m-%d %H:%M:%S\n", gmtime()))
        varkey = 'model_name'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey].decode(encoding='utf-8', errors='ignore')));
        varkey = 'ver_info'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))

    def __repr__(self):
        d = self.dict_export()
        from pprint import pformat
        return pformat(d, indent=4, width=1)


class FwModEntry(LittleEndianStructure):
    _pack_ = 1
    _fields_ = [
      ('dt_len', c_uint),
      ('crc32', c_uint),
    ]

    def dict_export(self):
        d = dict()
        for (varkey, vartype) in self._fields_:
            d[varkey] = getattr(self, varkey)
        varkey = 'crc32'
        d[varkey] = "{:08X}".format(d[varkey])
        return d

    def __repr__(self):
        d = self.dict_export()
        from pprint import pformat
        return pformat(d, indent=4, width=1)


class FwModA9PostHeader(LittleEndianStructure):
    _pack_ = 1
    _fields_ = [
      ('part_size', c_uint * 15),
    ]

    def dict_export(self):
        d = dict()
        for (varkey, vartype) in self._fields_:
            d[varkey] = getattr(self, varkey)
        varkey = 'part_size'
        d[varkey] = " ".join("{:08x}".format(x) for x in d[varkey])
        return d

    def ini_export(self, fp):
        d = self.dict_export()
        # No header - this is a continuation of FwModA9Header export
        varkey = 'part_size'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))

    def __repr__(self):
        d = self.dict_export()
        from pprint import pformat
        return pformat(d, indent=4, width=1)


class FwModPartHeader(LittleEndianStructure):
    _pack_ = 1
    _fields_ = [
      ('crc32', c_uint),
      ('version', c_uint),
      ('build_date', c_uint),
      ('dt_len', c_uint),
      ('mem_addr', c_uint),
      ('flag1', c_uint),
      ('magic', c_uint),
      ('flag2', c_uint),
      ('padding', c_uint * 56),
    ]

    def build_date_year(self):
        return (self.build_date >> 16) & 65535

    def build_date_month(self):
        return (self.build_date >> 8) & 255

    def build_date_day(self):
        return (self.build_date) & 255

    def version_major(self):
        return (self.version >> 16) & 65535

    def version_minor(self):
        return (self.version) & 65535

    def dict_export(self):
        d = dict()
        for (varkey, vartype) in self._fields_:
            d[varkey] = getattr(self, varkey)
        varkey = 'mem_addr'
        d[varkey] = "{:08X}".format(d[varkey])
        varkey = 'version'
        d[varkey] = "{:d}.{:d}".format(self.version_major(), self.version_minor())
        varkey = 'build_date'
        d[varkey] = "{:d}-{:02d}-{:02d}".format(self.build_date_year(), self.build_date_month(), self.build_date_day())
        varkey = 'flag1'
        d[varkey] = "{:08X}".format(d[varkey])
        varkey = 'flag2'
        d[varkey] = "{:08X}".format(d[varkey])
        varkey = 'magic'
        d[varkey] = "{:08X}".format(d[varkey])
        varkey = 'crc32'
        d[varkey] = "{:08X}".format(d[varkey])
        varkey = 'padding'
        d[varkey] = " ".join("{:08x}".format(x) for x in d[varkey])
        return d

    def ini_export(self, fp, i):
        d = self.dict_export()
        if (i < len(part_entry_type_name)):
            ptyp_name = part_entry_type_name[i]
        else:
            ptyp_name = "type {:02d}".format(i)
        fp.write("# Ambarella Firmware Packer section header file. Loosly based on AFT format.\n")
        fp.write("# Stores partition with {:s}\n".format(ptyp_name))
        fp.write(strftime("# Generated on %Y-%m-%d %H:%M:%S\n", gmtime()))
        varkey = 'mem_addr'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))
        varkey = 'version'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))
        varkey = 'build_date'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))
        varkey = 'flag1'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))
        varkey = 'flag2'
        fp.write("{:s}={:s}\n".format(varkey, d[varkey]))

    def __repr__(self):
        d = self.dict_export()
        from pprint import pformat
        return pformat(d, indent=4, width=1)


crc32_tab = [
  0x00000000, 0x77073096, 0xee0e612c, 0x990951ba, 0x076dc419, 0x706af48f, 0xe963a535, 0x9e6495a3,
  0x0edb8832, 0x79dcb8a4, 0xe0d5e91e, 0x97d2d988, 0x09b64c2b, 0x7eb17cbd, 0xe7b82d07, 0x90bf1d91,
  0x1db71064, 0x6ab020f2, 0xf3b97148, 0x84be41de, 0x1adad47d, 0x6ddde4eb, 0xf4d4b551, 0x83d385c7,
  0x136c9856, 0x646ba8c0, 0xfd62f97a, 0x8a65c9ec, 0x14015c4f, 0x63066cd9, 0xfa0f3d63, 0x8d080df5,
  0x3b6e20c8, 0x4c69105e, 0xd56041e4, 0xa2677172, 0x3c03e4d1, 0x4b04d447, 0xd20d85fd, 0xa50ab56b,
  0x35b5a8fa, 0x42b2986c, 0xdbbbc9d6, 0xacbcf940, 0x32d86ce3, 0x45df5c75, 0xdcd60dcf, 0xabd13d59,
  0x26d930ac, 0x51de003a, 0xc8d75180, 0xbfd06116, 0x21b4f4b5, 0x56b3c423, 0xcfba9599, 0xb8bda50f,
  0x2802b89e, 0x5f058808, 0xc60cd9b2, 0xb10be924, 0x2f6f7c87, 0x58684c11, 0xc1611dab, 0xb6662d3d,
  0x76dc4190, 0x01db7106, 0x98d220bc, 0xefd5102a, 0x71b18589, 0x06b6b51f, 0x9fbfe4a5, 0xe8b8d433,
  0x7807c9a2, 0x0f00f934, 0x9609a88e, 0xe10e9818, 0x7f6a0dbb, 0x086d3d2d, 0x91646c97, 0xe6635c01,
  0x6b6b51f4, 0x1c6c6162, 0x856530d8, 0xf262004e, 0x6c0695ed, 0x1b01a57b, 0x8208f4c1, 0xf50fc457,
  0x65b0d9c6, 0x12b7e950, 0x8bbeb8ea, 0xfcb9887c, 0x62dd1ddf, 0x15da2d49, 0x8cd37cf3, 0xfbd44c65,
  0x4db26158, 0x3ab551ce, 0xa3bc0074, 0xd4bb30e2, 0x4adfa541, 0x3dd895d7, 0xa4d1c46d, 0xd3d6f4fb,
  0x4369e96a, 0x346ed9fc, 0xad678846, 0xda60b8d0, 0x44042d73, 0x33031de5, 0xaa0a4c5f, 0xdd0d7cc9,
  0x5005713c, 0x270241aa, 0xbe0b1010, 0xc90c2086, 0x5768b525, 0x206f85b3, 0xb966d409, 0xce61e49f,
  0x5edef90e, 0x29d9c998, 0xb0d09822, 0xc7d7a8b4, 0x59b33d17, 0x2eb40d81, 0xb7bd5c3b, 0xc0ba6cad,
  0xedb88320, 0x9abfb3b6, 0x03b6e20c, 0x74b1d29a, 0xead54739, 0x9dd277af, 0x04db2615, 0x73dc1683,
  0xe3630b12, 0x94643b84, 0x0d6d6a3e, 0x7a6a5aa8, 0xe40ecf0b, 0x9309ff9d, 0x0a00ae27, 0x7d079eb1,
  0xf00f9344, 0x8708a3d2, 0x1e01f268, 0x6906c2fe, 0xf762575d, 0x806567cb, 0x196c3671, 0x6e6b06e7,
  0xfed41b76, 0x89d32be0, 0x10da7a5a, 0x67dd4acc, 0xf9b9df6f, 0x8ebeeff9, 0x17b7be43, 0x60b08ed5,
  0xd6d6a3e8, 0xa1d1937e, 0x38d8c2c4, 0x4fdff252, 0xd1bb67f1, 0xa6bc5767, 0x3fb506dd, 0x48b2364b,
  0xd80d2bda, 0xaf0a1b4c, 0x36034af6, 0x41047a60, 0xdf60efc3, 0xa867df55, 0x316e8eef, 0x4669be79,
  0xcb61b38c, 0xbc66831a, 0x256fd2a0, 0x5268e236, 0xcc0c7795, 0xbb0b4703, 0x220216b9, 0x5505262f,
  0xc5ba3bbe, 0xb2bd0b28, 0x2bb45a92, 0x5cb36a04, 0xc2d7ffa7, 0xb5d0cf31, 0x2cd99e8b, 0x5bdeae1d,
  0x9b64c2b0, 0xec63f226, 0x756aa39c, 0x026d930a, 0x9c0906a9, 0xeb0e363f, 0x72076785, 0x05005713,
  0x95bf4a82, 0xe2b87a14, 0x7bb12bae, 0x0cb61b38, 0x92d28e9b, 0xe5d5be0d, 0x7cdcefb7, 0x0bdbdf21,
  0x86d3d2d4, 0xf1d4e242, 0x68ddb3f8, 0x1fda836e, 0x81be16cd, 0xf6b9265b, 0x6fb077e1, 0x18b74777,
  0x88085ae6, 0xff0f6a70, 0x66063bca, 0x11010b5c, 0x8f659eff, 0xf862ae69, 0x616bffd3, 0x166ccf45,
  0xa00ae278, 0xd70dd2ee, 0x4e048354, 0x3903b3c2, 0xa7672661, 0xd06016f7, 0x4969474d, 0x3e6e77db,
  0xaed16a4a, 0xd9d65adc, 0x40df0b66, 0x37d83bf0, 0xa9bcae53, 0xdebb9ec5, 0x47b2cf7f, 0x30b5ffe9,
  0xbdbdf21c, 0xcabac28a, 0x53b39330, 0x24b4a3a6, 0xbad03605, 0xcdd70693, 0x54de5729, 0x23d967bf,
  0xb3667a2e, 0xc4614ab8, 0x5d681b02, 0x2a6f2b94, 0xb40bbe37, 0xc30c8ea1, 0x5a05df1b, 0x2d02ef8d,
]


def amba_a9_part_entry_type_id(i):
    if (i >= len(part_entry_type_id)):
        return "{:02d}".format(i)
    return part_entry_type_id[i]


def amba_calculate_crc32h_part(buf, pcrc):
    """ A twist on crc32 hashing algorithm, probably different from original CRC32 due to a programming mistake.
    """
    crc = pcrc
    for octet in buf:
        crc = crc32_tab[(crc ^ octet) & 0xff] ^ (crc >> 8)
    return crc & 0xffffffff


def amba_calculate_crc32b_part(buf, pcrc):
    """ A standard crc32b hashing algorithm, the same as used in ZIP/PNG.
    """
    return zlib.crc32(buf, pcrc) & 0xffffffff


def amba_calculate_crc32(buf):
    return amba_calculate_crc32b_part(buf, 0)


def amba_detect_format(po, fwmdlfile):
    """ Detects which binary format the firmware module file has.
    """
    #TODO make multiple formats support
    # FC220 has different format (2016 - FwModA9Header longer 4 butes, 319 ints in FwModA9PostHeader)
    return '2014'


# We really need both i and ptyp params
def amba_extract_partition_head(po, e, i, ptyp):
    fwpartfile = open("{:s}_part_{:s}.a9h".format(po.ptprefix, ptyp), "w")
    e.ini_export(fwpartfile, i)
    fwpartfile.close()


def amba_read_part_head(po, i, ptyp):
    e = FwModPartHeader()
    e.magic = 0xA324EB90
    fname = "{:s}_part_{:s}.a9h".format(po.ptprefix, ptyp)
    parser = configparser.ConfigParser()
    with open(fname, "r") as lines:
        lines = itertools.chain(("[asection]",), lines)  # This line adds section header to ini
        parser.read_file(lines)
        e.mem_addr = int(parser.get("asection", "mem_addr"), 16)
        e.flag1 = int(parser.get("asection", "flag1"), 16)
        e.flag2 = int(parser.get("asection", "flag2"), 16)
        version_s = parser.get("asection", "version")
        version_m = re.search('(?P<major>[0-9]+)[.](?P<minor>[0-9]+)', version_s)
        e.version = ((int(version_m.group("major"), 10) & 0xffff) << 16) + (int(version_m.group("minor"), 10) % 0xffff)
        build_date_s = parser.get("asection", "build_date")
        build_date_m = re.search('(?P<year>[0-9]+)[-](?P<month>[0-9]+)[-](?P<day>[0-9]+)', build_date_s)
        e.build_date = (
          ((int(build_date_m.group("year"), 10) & 0xffff) << 16) +
          ((int(build_date_m.group("month"), 10) & 0xff) << 8) +
          (int(build_date_m.group("day"), 10) & 0xff) )
    del parser
    return e


def amba_extract_mod_head(po, modhead, ptyp_names, modposthd):
    fwpartfile = open("{:s}_header.a9h".format(po.ptprefix), "w")
    modhead.ini_export(fwpartfile)
    fwpartfile.write("part_load={:s}\n".format(",".join("{:s}".format(x) for x in ptyp_names)))
    modposthd.ini_export(fwpartfile)
    fwpartfile.close()


def amba_read_mod_head(po):
    modhead = FwModA9Header()
    modposthd = FwModA9PostHeader()
    fname = "{:s}_header.a9h".format(po.ptprefix)
    parser = configparser.ConfigParser()
    with open(fname, "r") as lines:
        lines = itertools.chain(("[asection]",), lines)  # This line adds section header to ini
        parser.read_file(lines)
    ptyp_names = parser.get("asection", "part_load").split(",")
    part_sizes_s = parser.get("asection", "part_size").split(" ")
    part_sizes = [int(n, 16) for n in part_sizes_s]
    ver_info_s = parser.get("asection", "ver_info")
    ver_info_m = re.search('(?P<major>[0-9]+)[.](?P<minor>[0-9]+)[-](?P<svn>[0-9A-Fa-f]+)', ver_info_s)
    modhead.model_name = parser.get("asection", "model_name").encode("utf-8")
    modhead.ver_info = (
      ((int(ver_info_m.group("major"), 10) & 0xff) << 24) +
      ((int(ver_info_m.group("minor"), 10) % 0xff) << 16) +
      (int(ver_info_m.group("svn"), 16) % 0xffff)
    )
    for i, n in enumerate(part_sizes):
        modposthd.part_size[i] = n
    del parser
    return (modhead, ptyp_names, modposthd)


def amba_extract_partition_data(po, fwmdlfile, e, i, ptyp, hdcrc):
    """ Extract partiton data to a file, using file input.

    The initial cummulative checksum `hdcrc` can be None, which disables its computation.
    """
    fwpartfile = open("{:s}_part_{:s}.a9s".format(po.ptprefix, ptyp), "wb")
    if (po.binhead):
        fwpartfile.write((c_ubyte * sizeof(e)).from_buffer_copy(e))
    ptcrc = 0
    n = 0
    while n < e.dt_len:
        copy_buffer = fwmdlfile.read(min(1024 * 1024, e.dt_len - n))
        if not copy_buffer:
            break
        n += len(copy_buffer)
        fwpartfile.write(copy_buffer)
        ptcrc = amba_calculate_crc32b_part(copy_buffer, ptcrc)
        if hdcrc is not None:
            hdcrc = amba_calculate_crc32h_part(copy_buffer, hdcrc)
    fwpartfile.close()
    if (n < e.dt_len):
        raise_or_warn(po, ValueError("Partition {:d} truncated, {:d} out of {:d} bytes".format(i, n, e.dt_len)))
    if (ptcrc != e.crc32):
        raise_or_warn(po, ValueError("Entry {:d} data checksum mismatch; got {:08X}, expected {:08X}."
          .format(i, ptcrc, e.crc32)))
    elif (po.verbose > 1):
        print("{}: Entry {:2d} data checksum {:08X} matched OK".format(po.fwmdlfile, i, ptcrc))
    return hdcrc


def amba_extract_partition_data_mm(po, fwmdlmm, epos, e, i, ptyp, hdcrc):
    """ Extract partiton data to a file, using memory mapped input file.

    The initial cummulative checksum `hdcrc` can be None, which disables its computation.
    """
    fwpartfile = open("{:s}_part_{:s}.a9s".format(po.ptprefix, ptyp), "wb")
    dtpos_beg = epos + sizeof(FwModPartHeader)
    dtpos_end = epos + sizeof(FwModPartHeader) + e.dt_len
    if (po.binhead):
        dtpos_beg = epos
    fwpartfile.write(fwmdlmm[dtpos_beg:dtpos_end])
    fwpartfile.close()
    dtpos_beg = epos + sizeof(FwModPartHeader)
    ptcrc = amba_calculate_crc32(fwmdlmm[dtpos_beg:dtpos_end])
    if hdcrc is not None:
        hdcrc = amba_calculate_crc32h_part(fwmdlmm[dtpos_beg:dtpos_end], hdcrc)
    if (ptcrc != e.crc32):
        raise_or_warn(po, ValueError("Entry {:d} checksum mismatch; got {:08X}, expected {:08X}.".format(i, ptcrc, e.crc32)))
    elif (po.verbose > 1):
        print("{}: Entry {:2d} data checksum {:08X} matched OK".format(po.fwmdlfile, i, ptcrc))
    return hdcrc


def amba_extract(po, fwmdlfile):
    modhead = FwModA9Header()
    fwmdlfile.seek(0, os.SEEK_END)
    fwmdlfile_len = fwmdlfile.tell()
    fwmdlfile.seek(0, os.SEEK_SET)
    if fwmdlfile.readinto(modhead) != sizeof(modhead):
        raise EOFError("Could not read firmware package file header.")
    if (po.verbose > 1):
        print("{}: Header:".format(po.fwmdlfile))
        print(modhead)
    hdcrc = 0xffffffff
    i = 0
    modentries = []
    ptyp_names = []
    while (True):
        hde = FwModEntry()
        if fwmdlfile.readinto(hde) != sizeof(hde):
            raise EOFError("Could not read firmware package file header entries.")
        # If both values are multiplications of 1024, and 2nd is non-zero, then assume we're past end
        # of entries array. Beyond entries array, there's an array of memory load addresses - and
        # load addresses are always rounded to multiplication of a power of 2.
        # Since specific Ambarella firmwares always have set number of partitions, we have to do
        # such guessing if we want one tool to support all Ambarella firmwares.
        if ((hde.dt_len & 0x3ff) == 0) and ((hde.crc32 & 0x3ff) == 0) and (hde.crc32 != 0):
            fwmdlfile.seek(-sizeof(hde), os.SEEK_CUR)
            break
        if (sizeof(modhead) + i * sizeof(hde) + hde.dt_len >= fwmdlfile_len):
            if (po.verbose > 1):
                print("{}: Detection finished with entry larger than file; expecting {:d} entries"
                  .format(po.fwmdlfile, len(modentries)))
            raise_or_warn(po, ValueError("Detection finished with unusual entry sizes, verify files."))
            fwmdlfile.seek(-sizeof(hde), os.SEEK_CUR)
            break
        modentries.append(hde)
        if (hde.dt_len > 0):
            ptyp_names.append(amba_a9_part_entry_type_id(i))
        i += 1
        if (i > 128):
            raise EOFError("Could not find header entries end marking.")
    if (po.verbose > 1):
        print("{}: After detection, expecting {:d} entries".format(po.fwmdlfile, len(modentries)))
    if (po.verbose > 1):
        print("{}: Entries:".format(po.fwmdlfile))
        print(modentries)

    modposthd = FwModA9PostHeader()
    if fwmdlfile.readinto(modposthd) != sizeof(modposthd):
        raise EOFError("Could not read firmware package file header.")
    if (po.verbose > 1):
        print("{}: Post Header:".format(po.fwmdlfile))
        print(modposthd)
    amba_extract_mod_head(po, modhead, ptyp_names, modposthd)
    i = -1
    while True:
        i += 1
        # Skip unused modentries
        if (i < len(modentries)):
            hde = modentries[i]
            if (hde.dt_len < 1):
                continue
        else:
            # Do not show warning yet - maybe the file is at EOF
            hde = FwModEntry()
        epos = fwmdlfile.tell()
        e = FwModPartHeader()
        n = fwmdlfile.readinto(e)
        if (n is None) or (n == 0):
            # End Of File, correct ending
            break
        if n != sizeof(e):
            raise EOFError("Could not read firmware package partition header, got {:d} out of {:d}."
              .format(n, sizeof(e)))
        if e.magic != 0xA324EB90:
            raise_or_warn(po, ValueError("Invalid magic value in partition {:d} header.".format(i)))
        if (po.verbose > 1):
            print("{}: Entry {}".format(po.fwmdlfile, i))
            print(e)
        hdcrc = amba_calculate_crc32h_part((c_ubyte * sizeof(e)).from_buffer_copy(e), hdcrc)
        if (e.dt_len < 16) or (e.dt_len > 128*1024*1024):
            raise_or_warn(po, ValueError("Entry at {:d} has bad size, {:d} bytes".format(epos, e.dt_len)))
        # Warn if no more module entries were expected
        if (i >= len(modentries)):
            raise_or_warn(po, ValueError("Data continues after parsing all {:d} known partitions; header inconsistent.".format(i)))
        print("{}: Extracting entry {:2d}, pos {:8d}, len {:8d} bytes".format(po.fwmdlfile, i, epos, e.dt_len))
        ptyp = amba_a9_part_entry_type_id(i)
        amba_extract_partition_head(po, e, i, ptyp)
        hdcrc = amba_extract_partition_data(po, fwmdlfile, e, i, ptyp, hdcrc)
        if (hdcrc != hde.crc32):
            raise_or_warn(po, ValueError("Entry {:d} cummulative checksum mismatch; got {:08X}, expected {:08X}."
              .format(i, hdcrc, hde.crc32)))
        elif (po.verbose > 1):
            print("{}: Entry {:2d} cummulative checksum {:08X} matched OK".format(po.fwmdlfile, i, hdcrc))
        # Check if the date makes sense
        if ((e.build_date_year() < 1970) or (e.build_date_month() < 1) or
          (e.build_date_month() > 12) or (e.build_date_day() < 1) or (e.build_date_day() > 31)):
            raise_or_warn(po, ValueError("Entry {:d} date makes no sense.".format(i)))
        elif (e.build_date_year() < 2004):
            raise_or_warn(po, ValueError("Entry {:d} date is from before Ambarella formed as company.".format(i)))
        # verify if padding area is completely filled with 0x00000000
        if (e.padding[0] != 0x00000000) or (len(set(e.padding)) != 1):
            raise_or_warn(po, ValueError("Partition {:d} header uses values from padded area in an unknown manner.".format(i)))
    # Now verify checksum in main header
    hdcrc = hdcrc ^ 0xffffffff
    if (hdcrc != modhead.crc32):
        raise_or_warn(po, ValueError("Total cummulative checksum mismatch; got {:08X}, expected {:08X}.".format(hdcrc, modhead.crc32)))
    elif (po.verbose > 1):
        print("{}: Total cummulative checksum {:08X} matched OK".format(po.fwmdlfile, hdcrc))
    return


def amba_search_extract(po, fwmdlfile):
    fwmdlmm = mmap.mmap(fwmdlfile.fileno(), length=0, access=mmap.ACCESS_READ)
    epos = -sizeof(FwModPartHeader)
    prev_dtlen = 0
    prev_dtpos = 0
    i = 0
    while True:
        epos = fwmdlmm.find(b'\x90\xEB\x24\xA3', epos+sizeof(FwModPartHeader))
        if (epos < 0):
            break
        epos -= 24  # pos of 'magic' within FwModPartHeader
        if (epos < 0):
            continue
        dtpos = epos + sizeof(FwModPartHeader)
        e = FwModPartHeader.from_buffer_copy(fwmdlmm[epos:dtpos])
        if (e.dt_len < 16) or (e.dt_len > 128*1024*1024) or (e.dt_len > fwmdlmm.size()-dtpos):
            print("{}: False positive - entry at {:d} has bad size, {:d} bytes".format(po.fwmdlfile, epos, e.dt_len))
            continue
        print("{}: Extracting entry {:2d}, pos {:8d}, len {:8d} bytes".format(po.fwmdlfile, i, epos, e.dt_len))
        if (prev_dtpos+prev_dtlen > epos):
            raise_or_warn(po, ValueError("Partition {:d} overlaps with previous by {:d} bytes"
              .format(i, prev_dtpos + prev_dtlen - epos)))
        ptyp = "{:02d}".format(i)
        amba_extract_partition_head(po, e, i, ptyp)
        amba_extract_partition_data_mm(po, fwmdlmm, epos, e, i, ptyp, None)
        prev_dtlen = e.dt_len
        prev_dtpos = dtpos
        i += 1
    return


def get_amba_max_modentry(po, ptyp_names):
    modentry_max = 0
    for ptyp in ptyp_names:
        if ptyp not in part_entry_type_id:
            raise ValueError("Unrecognized partition name in 'part_load' option.")
        i = part_entry_type_id.index(ptyp)
        if (modentry_max < i):
            modentry_max = i
    return modentry_max


def amba_partition_exists(po, i, ptyp, ptyp_names):
    """ Checks if there are any partition data for given index and type.
    """
    fname = "{:s}_part_{:s}.a9s".format(po.ptprefix, ptyp)
    # Skip unused modentries
    if ptyp not in ptyp_names:
        if (po.verbose > 1):
            print("{}: Entry {:2d} empty".format(po.fwmdlfile, i))
        return False
    # Also skip nonexisting ones
    if (os.stat(fname).st_size < 1):
        raise_or_warn(po, ValueError("Partition {:d} marked as existing but empty".format(i)))
        return False
    return True


def amba_merge_partition_data(po, fwmdlfile, e, i, ptyp):
    """ Opens the partition file for given `ptyp`, and copies the data to `fwmdlfile`.

    This also updates properties within the `e` entry class.
    """
    fname = "{:s}_part_{:s}.a9s".format(po.ptprefix, ptyp)
    fwpartfile = open(fname, 'rb')
    ptcrc = 0
    n = 0
    while True:
        copy_buffer = fwpartfile.read(1024 * 1024)
        if not copy_buffer:
            break
        n += len(copy_buffer)
        fwmdlfile.write(copy_buffer)
        ptcrc = amba_calculate_crc32b_part(copy_buffer, ptcrc)
    e.dt_len = n
    e.crc32 = ptcrc
    if (po.verbose > 1):
        print("{}: Entry {:2d} checksum {:08X}".format(po.fwmdlfile, i, ptcrc))
    return


def amba_create(po, fwmdlfile):
    # Read headers from INI files
    (modhead, ptyp_names, modposthd) = amba_read_mod_head(po)
    modentries = []
    # Get amount of partition slots to allocate
    modentry_max = get_amba_max_modentry(po, ptyp_names)
    # Create module entry for each partition
    for i in range(modentry_max+1):
        hde = FwModEntry()
        modentries.append(hde)
    # Write the unfinished headers
    fwmdlfile.write((c_ubyte * sizeof(modhead)).from_buffer_copy(modhead))
    for hde in modentries:
        fwmdlfile.write((c_ubyte * sizeof(hde)).from_buffer_copy(hde))
    fwmdlfile.write((c_ubyte * sizeof(modposthd)).from_buffer_copy(modposthd))
    # Write the partitions
    part_heads = []
    i = -1
    while True:
        i += 1
        if (i >= len(modentries)):
            break
        hde = modentries[i]
        ptyp = amba_a9_part_entry_type_id(i)
        if not amba_partition_exists(po, i, ptyp, ptyp_names):
            e = FwModPartHeader()
            part_heads.append(e)
            continue
        e = amba_read_part_head(po, i, ptyp)
        epos = fwmdlfile.tell()
        # Write unfinished header
        fwmdlfile.write((c_ubyte * sizeof(e)).from_buffer_copy(e))
        # Copy partition data and compute CRC
        amba_merge_partition_data(po, fwmdlfile, e, i, ptyp)
        part_heads.append(e)
        # Write final partition header
        npos = fwmdlfile.tell()
        fwmdlfile.seek(epos, os.SEEK_SET)
        fwmdlfile.write((c_ubyte * sizeof(e)).from_buffer_copy(e))
        fwmdlfile.seek(npos, os.SEEK_SET)
        # Store file beginning header
        hde.dt_len = sizeof(e) + e.dt_len
        modentries[i] = hde
    # Compute cummulative CRC32
    if (po.verbose > 1):
        print("{}: Recomputing checksums".format(po.fwmdlfile))
    hdcrc = 0xffffffff
    i = -1
    while True:
        i += 1
        if (i >= len(modentries)):
            break
        hde = modentries[i]
        ptyp = amba_a9_part_entry_type_id(i)
        fname = "{:s}_part_{:s}.a9s".format(po.ptprefix, ptyp)
        if (hde.dt_len < 1):
            continue
        fwpartfile = open(fname, 'rb')
        e = part_heads[i]
        hdcrc = amba_calculate_crc32h_part((c_ubyte * sizeof(e)).from_buffer_copy(e), hdcrc)
        n = 0
        while n < e.dt_len:
            copy_buffer = fwpartfile.read(min(1024 * 1024, e.dt_len - n))
            if not copy_buffer:
                break
            n += len(copy_buffer)
            hdcrc = amba_calculate_crc32h_part(copy_buffer, hdcrc)
        hde.crc32 = hdcrc
        modentries[i] = hde
    hdcrc = hdcrc ^ 0xffffffff
    modhead.crc32 = hdcrc
    if (po.verbose > 1):
        print("{}: Total cummulative checksum {:08X}".format(po.fwmdlfile, hdcrc))
    # Write all headers again
    fwmdlfile.seek(0, os.SEEK_SET)
    fwmdlfile.write((c_ubyte * sizeof(modhead)).from_buffer_copy(modhead))
    for hde in modentries:
        fwmdlfile.write((c_ubyte * sizeof(hde)).from_buffer_copy(hde))
    fwmdlfile.write((c_ubyte * sizeof(modposthd)).from_buffer_copy(modposthd))


def main():
    """ Main executable function.

    Its task is to parse command line options and call a function which performs requested command.
    """
    parser = argparse.ArgumentParser(description=__doc__.split('.')[0])

    parser.add_argument('-m', '--fwmdlfile', type=str, required=True,
          help="name of the firmware module file")

    parser.add_argument('-t', '--ptprefix', type=str,
          help=("file name prefix for the single decomposed partitions "
           "(defaults to base name of firmware module file)"))

    parser.add_argument('--binfmt', type=str, default='auto',
          help=("set binary format version"
           "(default is to detect it (auto); valid formats are 2014 and 2016)"))

    parser.add_argument('--binhead', action='store_true',
          help=("leave (`-x`) or use (`-a`) binary header in front of partition"
           "this leaves the original binary header before each partition"
           "on extraction, and uses that header on module file creation;"
           "you normally should have no need to use it"))

    parser.add_argument('-f', '--force-continue', action='store_true',
          help="force continuing execution despite warning signs of issues")

    parser.add_argument('-v', '--verbose', action='count', default=0,
          help="increases verbosity level; max level is set by `-vvv`")

    subparser = parser.add_mutually_exclusive_group(required=True)

    subparser.add_argument('-x', '--extract', action='store_true',
          help="extract firmware module file into partitions")

    subparser.add_argument('-s', '--search', action='store_true',
          help=("search for partitions within firmware module and extract them "
           "(works similar to `-x`, but uses brute-force search for partitions)"))

    subparser.add_argument('-a', '--add', action='store_true',
          help=("add partition files to firmware module file "
           "(works only on data created with `-x`; the `-s` is insufficient)"))

    subparser.add_argument('--version', action='version', version="%(prog)s {version} by {author}"
            .format(version=__version__, author=__author__),
          help="display version information and exit")

    po = parser.parse_args()

    if len(po.fwmdlfile) > 0 and po.ptprefix is None:
        po.ptprefix = os.path.splitext(os.path.basename(po.fwmdlfile))[0]

    if po.extract:
        if (po.verbose > 0):
            print("{}: Opening for extraction".format(po.fwmdlfile))
        with open(po.fwmdlfile, 'rb') as fwmdlfile:
            if po.binfmt == 'auto':
                po.binfmt = amba_detect_format(po, fwmdlfile)
            amba_extract(po, fwmdlfile)

    elif po.search:
        if (po.verbose > 0):
            print("{}: Opening for search".format(po.fwmdlfile))
        with open(po.fwmdlfile, 'rb') as fwmdlfile:
            amba_search_extract(po, fwmdlfile)

    elif po.add:
        if (po.verbose > 0):
            print("{}: Opening for creation".format(po.fwmdlfile))
        with open(po.fwmdlfile, 'wb') as fwmdlfile:
            amba_create(po, fwmdlfile)

    else:
        raise NotImplementedError("Unsupported command.")


if __name__ == '__main__':
    try:
        main()
    except Exception as ex:
        eprint("Error: "+str(ex))
        if 0: raise
        sys.exit(10)