dji_xv4_fwcon.py

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

""" DJI 'xV4' Firmware Container tool.

Extract and creates the firmware package files.
"""

# 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.3.2"
__author__ = "Mefistotelis @ Original Gangsters"
__license__ = "GPL"

import sys
import re
import os
import hashlib
import binascii
import argparse
import configparser
import itertools
from ctypes import c_char, c_int, c_ubyte, c_ushort, c_uint
from ctypes import sizeof, LittleEndianStructure
from time import gmtime, strftime, strptime
from calendar import timegm
from Crypto.Cipher import AES


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


class DjiModuleTarget():
    "Stores identification info on module for specific target"
    def __init__(self, kind, model, name, desc):
        self.kind = kind
        self.model = model
        self.name = name
        self.desc = desc


dji_targets = [
    DjiModuleTarget( 1,-1, "CAM",     "camera"),
    DjiModuleTarget( 1, 0, "FC300X",  "camera 'Ambarella A9SE' App"), # P3X
    DjiModuleTarget( 1, 1, "CAMLDR",  "camera 'Ambarella A9SE' Ldr"), # P3X
    DjiModuleTarget( 1, 2, "CAMBST",  "camera BST"),
    DjiModuleTarget( 1, 4, "CAMBCPU", "camera BCPU"),
    DjiModuleTarget( 1, 5, "CAMLCPU", "camera LCPU"),
    DjiModuleTarget( 1, 6, "ZQ7020",  "camera 'Xilinx Zynq 7020'"),
    DjiModuleTarget( 2,-1, "MBAPP",   "mobile app"),
    DjiModuleTarget( 3,-1, "MC",      "main controller"),
    DjiModuleTarget( 3, 5, "MCLDR",   "main controller 'A3' ldr"), # P3X
    DjiModuleTarget( 3, 6, "MCAPP",   "main controller 'A3' app"), # P3X
    DjiModuleTarget( 4,-1, "GIMBAL",  "gimbal"),
    DjiModuleTarget( 4, 0, "GIMBAL0", "gimbal mdl 0"), # P3X
    DjiModuleTarget( 5,-1, "CENTER",  "central board"),
    DjiModuleTarget( 5, 0, "CENTER0", "central board mdl 0"),
    DjiModuleTarget( 6,-1, "RMRAD",   "remote radio"),
    DjiModuleTarget( 7,-1, "WIFI",    "Wi-Fi"),
    DjiModuleTarget( 7, 0, "WIFI0",   "Wi-Fi mdl 0"),
    DjiModuleTarget( 8,-1, "VENC",    "video encoder in air"),
    DjiModuleTarget( 8, 0, "DM368",   "video encoder 'DaVinci Dm368 Linux'"), # P3X
    DjiModuleTarget( 8, 1, "IG810LB2","video encoder 'IG810 LB2_ENC'"),
    DjiModuleTarget( 9,-1, "LBMCA",   "lightbridge MCU in air"),
    DjiModuleTarget( 9, 0, "MCA1765", "lightbridge MCU 'STM32F103'"), # P3X, OSMO_X5R
    DjiModuleTarget(10,-1, "BATTFW",  "battery firmware"),
    DjiModuleTarget(11,-1, "BATTMGR", "battery controller"),
    DjiModuleTarget(11, 0, "BATTERY", "battery controller 1 app"), # P3X
    DjiModuleTarget(11, 1, "BATTERY2","battery controller 2 app"),
    DjiModuleTarget(12,-1, "ESC",     "electronic speed control"),
    DjiModuleTarget(12, 0, "ESC0",    "electronic speed control 0"), # P3X
    DjiModuleTarget(12, 1, "ESC1",    "electronic speed control 1"), # P3X
    DjiModuleTarget(12, 2, "ESC2",    "electronic speed control 2"), # P3X
    DjiModuleTarget(12, 3, "ESC3",    "electronic speed control 3"), # P3X
    DjiModuleTarget(13, 0, "VDEC",    "video decoder"),
    DjiModuleTarget(13, 0, "DM365M0", "video decoder 'DaVinci Dm365 Linux'"),
    DjiModuleTarget(13, 1, "DM365M1", "video decoder 'DaVinci Dm385 Linux'"),
    DjiModuleTarget(14,-1, "LBMCG",   "lightbridge MCU on ground"),
    DjiModuleTarget(14, 0, "MCG1765A","lightbridge MCU 'LPC1765 GROUND LB2'"),
    DjiModuleTarget(15,-1, "TXUSBC",  "transmitter usb controller"),
    DjiModuleTarget(15, 0, "TX68013", "transmitter usb 'IG810 LB2_68013_TX'"), # P3X
    DjiModuleTarget(16,-1, "RXUSBCG", "receiver usb controller"),
    DjiModuleTarget(16, 0, "RX68013", "receiver usb 'IG810 LB2_68013_RX ground'"), # GL300a
    DjiModuleTarget(16, 1, "RXCY2014","receiver usb 'IG810 LB2_CY2014_RX ground'"), # GL300b+
    DjiModuleTarget(17,-1, "MVOM",    "visual positioning"),
    DjiModuleTarget(17, 0, "MVOMC4",  "visual positioning module 'camera'"), # P3X
    DjiModuleTarget(17, 1, "MVOMS0",  "visual positioning module 'sonar'"), # P3X
    DjiModuleTarget(19,-1, "FPGAA",   "lightbridge FPGA on air"),
    DjiModuleTarget(19, 0, "FPGAA0",  "lightbridge FPGA on air model 0"), # P3X
    DjiModuleTarget(20,-1, "FPGAG",   "lightbridge FPGA on ground"),
    DjiModuleTarget(20, 3, "FPGAG3",  "lightbridge FPGA on ground 'LB2'"),
    DjiModuleTarget(25,-1, "IMU",     "inertial measurement unit"),
    DjiModuleTarget(25, 0, "IMUA3M0", "inertial measurement unit pt0"),
    DjiModuleTarget(25, 1, "IMUA3M1", "inertial measurement unit pt1"),
    DjiModuleTarget(26,-1, "RTK",     "real time kinematic"),
    DjiModuleTarget(26, 6, "RTKAPP",  "real time kinematic App"),
    DjiModuleTarget(26, 7, "RTKLDR",  "real time kinematic Ldr"),
    DjiModuleTarget(27,-1, "WIFIGND", "Wi-Fi ground"),
    DjiModuleTarget(29,-1, "PMU",     "power management unit"),
    DjiModuleTarget(29, 0, "PMUA3LDR","power management unit App"),
    DjiModuleTarget(29, 1, "PMUA3APP","power management unit Ldr"),
    DjiModuleTarget(30,-1, "TESTA",   "test A"),
    DjiModuleTarget(31,-1, "TESTB",   "test B")
]

encrypt_aes128_key = bytes([0x96, 0x70, 0x9a, 0xD3, 0x26, 0x67, 0x4A, 0xC3, 0x82, 0xB6, 0x69, 0x27, 0xE6, 0xd8, 0x84, 0x21])
encrypt_aes128_iv = bytes([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])


class FwPkgHeader(LittleEndianStructure):
  _pack_ = 1
  _fields_ = [('magic', c_uint),
              ('magic_ver', c_ushort),
              ('hdrend_offs', c_ushort),
              ('timestamp', c_uint),
              ('manufacturer', c_char * 16),
              ('model', c_char * 16),
              ('entry_count', c_ushort),
              ('ver_latest_enc', c_int),
              ('ver_rollbk_enc', c_int),
              ('padding', c_ubyte * 10)]

  def set_ver_latest(self, ver):
    self.ver_latest_enc = 0x5127A564 ^ ver ^ self.timestamp

  def set_ver_rollbk(self, ver):
    self.ver_rollbk_enc = 0x5127A564 ^ ver ^ self.timestamp

  def get_format_version(self):
    if self.magic == 0x12345678 and self.magic_ver == 0x0000:
        # The versions are expected to be set properly for this magic_ver; that's surprising,
        # as it means the values are invalid only for magic_ver == 1
        if (self.ver_latest_enc != 0):
            return 201412
        else:
            return 0
    elif self.magic == 0x12345678 and self.magic_ver == 0x0001:
        if (self.ver_latest_enc == 0 and self.ver_rollbk_enc == 0):
            return 201502
        else:
            return 201507
    # Higher magic_ver - allow any 16-bit, but hex digits must be within BCD range
    elif self.magic == 0x12345678 and self.magic_ver >= 0x0002 and self.magic_ver <= 0xFFFF and \
      all(((self.magic_ver >> n) & 0xF <= 9) for n in range(0,16,4)):
        return 201608
    else:
        return 0

  def set_format_version(self, ver):
    if ver == 201412:
        self.magic = 0x12345678
        self.magic_ver = 0x0000
        self.set_ver_latest(0)
        self.set_ver_rollbk(0)
    elif ver == 201502:
        self.magic = 0x12345678
        self.magic_ver = 0x0001
        self.ver_latest_enc = 0
        self.ver_rollbk_enc = 0
    elif ver == 201507:
        self.magic = 0x12345678
        self.magic_ver = 0x0001
        self.set_ver_latest(0)
        self.set_ver_rollbk(0)
    elif ver == 201608:
        self.magic = 0x12345678
        self.magic_ver = 0x1130 # storing most common val; expected to be overwritten
        self.set_ver_latest(0)
        self.set_ver_rollbk(0)
    else:
        raise ValueError("Unsupported package format version.")

  def dict_export(self):
    d = dict()
    for (varkey, vartype) in self._fields_:
        d[varkey] = getattr(self, varkey)
    varkey = 'ver_latest'
    d[varkey] = d['timestamp'] ^ d[varkey+"_enc"] ^ 0x5127A564
    varkey = 'ver_rollbk'
    d[varkey] = d['timestamp'] ^ d[varkey+"_enc"] ^ 0x5127A564
    varkey = 'padding'
    d[varkey] = "".join("{:02X}".format(x) for x in d[varkey])
    return d

  def ini_export(self, fp):
    d = self.dict_export()
    fp.write("# DJI Firmware Container main header file.\n")
    fp.write(strftime("# Generated on %Y-%m-%d %H:%M:%S\n", gmtime()))
    varkey = 'pkg_format'
    fp.write("{:s}={:d}\n".format(varkey,self.get_format_version()))
    if self.magic_ver >= 2:
        varkey = 'magic_ver'
        fp.write("{:s}={:04x}\n".format(varkey,self.magic_ver))
    varkey = 'manufacturer'
    fp.write("{:s}={:s}\n".format(varkey,d[varkey].decode("utf-8")))
    varkey = 'model'
    fp.write("{:s}={:s}\n".format(varkey,d[varkey].decode("utf-8")))
    varkey = 'timestamp'
    fp.write("{:s}={:s}\n".format(varkey,strftime("%Y-%m-%d %H:%M:%S",gmtime(d[varkey]))))
    varkey = 'ver_latest'
    fp.write("{:s}={:02d}.{:02d}.{:04d}\n".format(varkey, (d[varkey]>>24)&255, (d[varkey]>>16)&255, (d[varkey])&65535))
    varkey = 'ver_rollbk'
    fp.write("{:s}={:02d}.{:02d}.{:04d}\n".format(varkey, (d[varkey]>>24)&255, (d[varkey]>>16)&255, (d[varkey])&65535))
    #varkey = 'padding'
    #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 FwPkgEntry(LittleEndianStructure):
  _pack_ = 1
  _fields_ = [('target', c_ubyte),
              ('spcoding', c_ubyte),
              ('reserved2', c_ushort),
              ('version', c_uint),
              ('dt_offs', c_uint),
              ('stored_len', c_uint),
              ('decrypted_len', c_uint),
              ('stored_md5', c_ubyte * 16),
              ('decrypted_md5', c_ubyte * 16)]
  preencrypted = 0

  def get_encrypt_type(self):
      return (self.spcoding >> 4) & 0x0F

  def set_encrypt_type(self, enctype):
      self.spcoding = (self.spcoding & 0x0F) | ((enctype & 0x0F) << 4)

  def get_splvalue(self):
      return (self.spcoding) & 0x0F

  def set_splvalue(self, splval):
      self.spcoding = (self.spcoding & 0xF0) | (splval & 0x0F)

  def target_name(self):
    tg_kind = getattr(self, 'target') & 31
    tg_model = (getattr(self, 'target') >> 5) & 7
    module_info = next((mi for mi in dji_targets if mi.kind == tg_kind and mi.model == tg_model), None)
    if (module_info is not None):
        return module_info.desc
    # If not found, try getting category
    module_info = next((mi for mi in dji_targets if mi.kind == tg_kind and mi.model == -1), None)
    if (module_info is not None):
        return "{:s} model {:02d}".format(module_info.desc,tg_model)
    # If category also not found, return as unknown device
    return "device kind {:02} model {:02}".format(tg_kind,tg_model)

  def hex_stored_md5(self):
    varkey = 'stored_md5'
    return "".join("{:02x}".format(x) for x in getattr(self, varkey))

  def hex_decrypted_md5(self):
    varkey = 'decrypted_md5'
    return "".join("{:02x}".format(x) for x in getattr(self, varkey))

  def dict_export(self):
    d = dict()
    for (varkey, vartype) in self._fields_:
        d[varkey] = getattr(self, varkey)
    varkey = 'version'
    d[varkey] = "{:02d}.{:02d}.{:04d}".format((d[varkey]>>24)&255, (d[varkey]>>16)&255, (d[varkey])&65535)
    varkey = 'stored_md5'
    d[varkey] = "".join("{:02x}".format(x) for x in d[varkey])
    varkey = 'decrypted_md5'
    d[varkey] = "".join("{:02x}".format(x) for x in d[varkey])
    varkey = 'target'
    d[varkey] = "m{:02d}{:02d}".format(d[varkey]&31, (d[varkey]>>5)&7)
    varkey = 'encrypt_type'
    d[varkey] = self.get_encrypt_type()
    varkey = 'splvalue'
    d[varkey] = self.get_splvalue()
    varkey = 'target_name'
    d[varkey] = self.target_name()
    return d

  def ini_export(self, fp):
    d = self.dict_export()
    fp.write("# DJI Firmware Container module header file.\n")
    fp.write("# Stores firmware for {:s}\n".format(d['target_name']))
    fp.write(strftime("# Generated on %Y-%m-%d %H:%M:%S\n", gmtime()))
    varkey = 'target'
    fp.write("{:s}={:s}\n".format(varkey,d[varkey]))
    varkey = 'version'
    fp.write("{:s}={:s}\n".format(varkey,d[varkey]))
    varkey = 'encrypt_type'
    fp.write("{:s}={:d}\n".format(varkey,d[varkey]))
    if (d[varkey] != 0):
        # If we do not support encryption of given type, we may extract the file in encrypted form
        varkey = 'preencrypted'
        fp.write("{:s}={:d}\n".format(varkey,self.preencrypted))
    if (self.preencrypted):
        # We cannot compute decrypted MD5 for pre-encrypted files, so store them in INI file
        varkey = 'decrypted_md5'
        fp.write("{:s}={:s}\n".format(varkey,d[varkey]))
    varkey = 'splvalue'
    fp.write("{:s}={:d}\n".format(varkey,d[varkey]))
    varkey = 'reserved2'
    fp.write("{:s}={:04X}\n".format(varkey,d[varkey]))

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


crc16_tab = [
  0x0000, 0x1189, 0x2312, 0x329B, 0x4624, 0x57AD, 0x6536, 0x74BF,
  0x8C48, 0x9DC1, 0xAF5A, 0xBED3, 0xCA6C, 0xDBE5, 0xE97E, 0xF8F7,
  0x1081, 0x0108, 0x3393, 0x221A, 0x56A5, 0x472C, 0x75B7, 0x643E,
  0x9CC9, 0x8D40, 0xBFDB, 0xAE52, 0xDAED, 0xCB64, 0xF9FF, 0xE876,
  0x2102, 0x308B, 0x0210, 0x1399, 0x6726, 0x76AF, 0x4434, 0x55BD,
  0xAD4A, 0xBCC3, 0x8E58, 0x9FD1, 0xEB6E, 0xFAE7, 0xC87C, 0xD9F5,
  0x3183, 0x200A, 0x1291, 0x0318, 0x77A7, 0x662E, 0x54B5, 0x453C,
  0xBDCB, 0xAC42, 0x9ED9, 0x8F50, 0xFBEF, 0xEA66, 0xD8FD, 0xC974,
  0x4204, 0x538D, 0x6116, 0x709F, 0x0420, 0x15A9, 0x2732, 0x36BB,
  0xCE4C, 0xDFC5, 0xED5E, 0xFCD7, 0x8868, 0x99E1, 0xAB7A, 0xBAF3,
  0x5285, 0x430C, 0x7197, 0x601E, 0x14A1, 0x0528, 0x37B3, 0x263A,
  0xDECD, 0xCF44, 0xFDDF, 0xEC56, 0x98E9, 0x8960, 0xBBFB, 0xAA72,
  0x6306, 0x728F, 0x4014, 0x519D, 0x2522, 0x34AB, 0x0630, 0x17B9,
  0xEF4E, 0xFEC7, 0xCC5C, 0xDDD5, 0xA96A, 0xB8E3, 0x8A78, 0x9BF1,
  0x7387, 0x620E, 0x5095, 0x411C, 0x35A3, 0x242A, 0x16B1, 0x0738,
  0xFFCF, 0xEE46, 0xDCDD, 0xCD54, 0xB9EB, 0xA862, 0x9AF9, 0x8B70,
  0x8408, 0x9581, 0xA71A, 0xB693, 0xC22C, 0xD3A5, 0xE13E, 0xF0B7,
  0x0840, 0x19C9, 0x2B52, 0x3ADB, 0x4E64, 0x5FED, 0x6D76, 0x7CFF,
  0x9489, 0x8500, 0xB79B, 0xA612, 0xD2AD, 0xC324, 0xF1BF, 0xE036,
  0x18C1, 0x0948, 0x3BD3, 0x2A5A, 0x5EE5, 0x4F6C, 0x7DF7, 0x6C7E,
  0xA50A, 0xB483, 0x8618, 0x9791, 0xE32E, 0xF2A7, 0xC03C, 0xD1B5,
  0x2942, 0x38CB, 0x0A50, 0x1BD9, 0x6F66, 0x7EEF, 0x4C74, 0x5DFD,
  0xB58B, 0xA402, 0x9699, 0x8710, 0xF3AF, 0xE226, 0xD0BD, 0xC134,
  0x39C3, 0x284A, 0x1AD1, 0x0B58, 0x7FE7, 0x6E6E, 0x5CF5, 0x4D7C,
  0xC60C, 0xD785, 0xE51E, 0xF497, 0x8028, 0x91A1, 0xA33A, 0xB2B3,
  0x4A44, 0x5BCD, 0x6956, 0x78DF, 0x0C60, 0x1DE9, 0x2F72, 0x3EFB,
  0xD68D, 0xC704, 0xF59F, 0xE416, 0x90A9, 0x8120, 0xB3BB, 0xA232,
  0x5AC5, 0x4B4C, 0x79D7, 0x685E, 0x1CE1, 0x0D68, 0x3FF3, 0x2E7A,
  0xE70E, 0xF687, 0xC41C, 0xD595, 0xA12A, 0xB0A3, 0x8238, 0x93B1,
  0x6B46, 0x7ACF, 0x4854, 0x59DD, 0x2D62, 0x3CEB, 0x0E70, 0x1FF9,
  0xF78F, 0xE606, 0xD49D, 0xC514, 0xB1AB, 0xA022, 0x92B9, 0x8330,
  0x7BC7, 0x6A4E, 0x58D5, 0x495C, 0x3DE3, 0x2C6A, 0x1EF1, 0x0F78,
]

def dji_calculate_crc16_part(buf, pcrc):
    """ A non-standard crc16 hashing algorithm, looks like 32-bit one from Ambarella simply cut down in bits.
    """
    crc = pcrc
    for octet in buf:
        crc = crc16_tab[(crc ^ octet) & 0xff] ^ (crc >> 8)
    return crc & 0xffff


def dji_decrypt_block(cipher_buf, enc_key, enc_iv):
    """ Decrypts a buffer with AES in a DJI way.
    """
    block_sz = 256
    plain_buf = b""
    for cbpos in range(0, len(cipher_buf), block_sz):
        # Reinit the crypto for each block, this is how Dji does it
        crypto = AES.new(enc_key, AES.MODE_CBC, enc_iv)
        plain_buf += crypto.decrypt(cipher_buf[cbpos:cbpos+block_sz])
    return plain_buf, enc_iv


def dji_encrypt_block(cipher_buf, enc_key, enc_iv):
    """ Encrypts a buffer with AES in a DJI way.
    """
    block_sz = 256
    plain_buf = b""
    for cbpos in range(0, len(cipher_buf), block_sz):
        # Reinit the crypto for each block, this is how Dji does it
        crypto = AES.new(enc_key, AES.MODE_CBC, enc_iv)
        plain_buf += crypto.encrypt(cipher_buf[cbpos:cbpos+block_sz])
    return plain_buf, enc_iv


def dji_write_fwpkg_head(po, pkghead, minames):
    fname = "{:s}_head.ini".format(po.mdprefix)
    fwheadfile = open(fname, "w")
    pkghead.ini_export(fwheadfile)
    fwheadfile.write("{:s}={:s}\n".format("modules",' '.join(minames)))
    fwheadfile.close()


def dji_read_fwpkg_head(po):
    pkghead = FwPkgHeader()
    fname = "{:s}_head.ini".format(po.mdprefix)
    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)
    # Set magic fields properly
    pkgformat = parser.get("asection", "pkg_format").encode("utf-8")
    pkghead.set_format_version(int(pkgformat))
    if parser.has_option('asection', 'magic_ver'):
        magicver_s = parser.get('asection', 'magic_ver')
        pkghead.magic_ver = int(magicver_s,16)
    # Set the rest of the fields
    pkghead.manufacturer = parser.get("asection", "manufacturer").encode("utf-8")
    pkghead.model = parser.get("asection", "model").encode("utf-8")
    pkghead.timestamp = timegm(strptime(parser.get("asection", "timestamp"),"%Y-%m-%d %H:%M:%S"))
    ver_latest_s = parser.get("asection", "ver_latest")
    ver_latest_m = re.search('(?P<major>[0-9]+)[.](?P<minor>[0-9]+)[.](?P<svn>[0-9A-Fa-f]+)', ver_latest_s)
    pkghead.set_ver_latest(
       ((int(ver_latest_m.group("major"), 10) & 0xff) << 24) +
       ((int(ver_latest_m.group("minor"), 10) & 0xff) << 16) +
       (int(ver_latest_m.group("svn"), 10) & 0xffff)
    )
    ver_rollbk_s = parser.get("asection", "ver_rollbk")
    ver_rollbk_m = re.search('(?P<major>[0-9]+)[.](?P<minor>[0-9]+)[.](?P<svn>[0-9A-Fa-f]+)', ver_rollbk_s)
    pkghead.set_ver_rollbk(
      ((int(ver_rollbk_m.group("major"), 10) & 0xff) << 24) +
      ((int(ver_rollbk_m.group("minor"), 10) & 0xff) << 16) +
      (int(ver_rollbk_m.group("svn"), 10) & 0xffff)
    )
    minames_s = parser.get("asection", "modules")
    minames = minames_s.split(' ')
    pkghead.entry_count = len(minames)
    pkghead.hdrend_offs = sizeof(pkghead) + sizeof(FwPkgEntry)*pkghead.entry_count + sizeof(c_ushort)
    del parser
    return (pkghead, minames)


def dji_write_fwentry_head(po, i, e, miname):
    fname = "{:s}_{:s}.ini".format(po.mdprefix,miname)
    fwheadfile = open(fname, "w")
    e.ini_export(fwheadfile)
    fwheadfile.close()


def dji_read_fwentry_head(po, i, miname):
    hde = FwPkgEntry()
    fname = "{:s}_{:s}.ini".format(po.mdprefix,miname)
    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)
    target_s = parser.get("asection", "target")
    target_m = re.search('m(?P<kind>[0-9]{2})(?P<model>[0-9]{2})', target_s)
    hde.target = ((int(target_m.group("kind"),10)&0x1f)) + ((int(target_m.group("model"),10)&0x07)<<5)
    version_s = parser.get("asection", "version")
    version_m = re.search('(?P<major>[0-9]+)[.](?P<minor>[0-9]+)[.](?P<svn>[0-9]+)', version_s)
    hde.version = (
      ((int(version_m.group("major"), 10) & 0xff) << 24) +
      ((int(version_m.group("minor"), 10) % 0xff) << 16) +
      (int(version_m.group("svn"), 10) % 0xffff)
    )
    if parser.has_option("asection", "preencrypted"):
        hde.preencrypted = int(parser.get("asection", "preencrypted"),10)
    if (hde.preencrypted):
        decrypted_md5_s = parser.get("asection", "decrypted_md5")
        hde.decrypted_md5 = (c_ubyte * 16).from_buffer_copy(binascii.unhexlify(decrypted_md5_s))
    hde.set_encrypt_type( int(parser.get("asection", "encrypt_type"),10) )
    hde.set_splvalue( int(parser.get("asection", "splvalue"),10) )
    hde.reserved2 = int(parser.get("asection", "reserved2"),16)
    del parser
    return (hde)

def dji_extract(po, fwpkgfile):
    pkghead = FwPkgHeader()
    if fwpkgfile.readinto(pkghead) != sizeof(pkghead):
        raise EOFError("Could not read firmware package file header.")
    pkgformat = pkghead.get_format_version()
    if pkgformat == 0:
        if (not po.force_continue):
            eprint("{}: Error: Unexpected magic value in main header; input file is not a firmware package.".format(po.fwpkg))
            exit(1)
        eprint("{}: Warning: Unexpected magic value in main header; will try to extract anyway.".format(po.fwpkg))
    if (po.verbose > 1):
        print("{}: Package format version {:d} detected".format(po.fwpkg,pkgformat))
    if (pkghead.ver_latest_enc == 0 and pkghead.ver_rollbk_enc == 0):
        eprint("{}: Warning: Unversioned firmware package identified; this format is not fully supported.".format(po.fwpkg))
        # In this format, versions should be set from file name, and CRC16 of the header should be equal to values hard-coded in updater
    if (po.verbose > 1):
        print("{}: Header:".format(po.fwpkg))
        print(pkghead)
    curhead_checksum = dji_calculate_crc16_part((c_ubyte * sizeof(pkghead)).from_buffer_copy(pkghead), 0x3692)

    pkgmodules = []
    for i in range(pkghead.entry_count):
        hde = FwPkgEntry()
        if fwpkgfile.readinto(hde) != sizeof(hde):
            raise EOFError("Couldn't read firmware package file entry.")
        if (po.verbose > 1):
            print("{}: Module index {:d}".format(po.fwpkg,i))
            print(hde)
        curhead_checksum = dji_calculate_crc16_part((c_ubyte * sizeof(hde)).from_buffer_copy(hde), curhead_checksum)
        if hde.stored_len != hde.decrypted_len:
            eprint("{}: Warning: decrypted size differs from stored one, {:d} instead of {:d}; this is not supported."
              .format(po.fwpkg,hde.decrypted_len,hde.stored_len))
        chksum_enctype = hde.get_encrypt_type()
        if (chksum_enctype != 0):
            if (po.no_crypto):
                hde.preencrypted = 1
            elif (chksum_enctype == 1):
                encrypt_key = encrypt_aes128_key
                encrypt_iv  = encrypt_aes128_iv
            else:
                # Since we cannot decode the encryption, mark the entry as pre-encrypted to extract in encrypted form
                eprint("{}: Warning: Unknown encryption {:d} in module {:d}, extracting encrypted."
                  .format(po.fwpkg,chksum_enctype,i))
                hde.preencrypted = 1
        pkgmodules.append(hde)

    pkghead_checksum = c_ushort()
    if fwpkgfile.readinto(pkghead_checksum) != sizeof(pkghead_checksum):
        raise EOFError("Couldn't read firmware package file header checksum.")

    if curhead_checksum != pkghead_checksum.value:
        eprint("{}: Warning: Firmware package file header checksum did not match; should be {:04X}, found {:04X}."
          .format(po.fwpkg, pkghead_checksum.value, curhead_checksum))
    elif (po.verbose > 1):
        print("{}: Headers checksum {:04X} matches.".format(po.fwpkg,pkghead_checksum.value))

    if fwpkgfile.tell() != pkghead.hdrend_offs:
        eprint("{}: Warning: Header end offset does not match; should end at {}, ends at {}."
          .format(po.fwpkg,pkghead.hdrend_offs,fwpkgfile.tell()))

    # Prepare array of names; "0" will mean empty index
    minames = ["0"]*len(pkgmodules)
    # Name the modules after target component
    for i, hde in enumerate(pkgmodules):
        if hde.stored_len > 0:
            d = hde.dict_export()
            minames[i] = "{:s}".format(d['target'])
    # Rename targets in case of duplicates
    minames_seen = set()
    for i in range(len(minames)):
        miname = minames[i]
        if miname in minames_seen:
            # Add suffix a..z to multiple uses of the same module
            for miname_suffix in range(97,110):
                if miname+chr(miname_suffix) not in minames_seen:
                    break
            # Show warning the first time duplicate is found
            if (miname_suffix == 97):
                eprint("{}: Warning: Found multiple modules {:s}; invalid firmware.".format(po.fwpkg,miname))
            minames[i] = miname+chr(miname_suffix)
        minames_seen.add(minames[i])
    minames_seen = None

    dji_write_fwpkg_head(po, pkghead, minames)

    for i, hde in enumerate(pkgmodules):
        if minames[i] == "0":
            if (po.verbose > 0):
                print("{}: Skipping module index {}, {} bytes".format(po.fwpkg,i,hde.stored_len))
            continue
        if (po.verbose > 0):
            print("{}: Extracting module index {}, {} bytes".format(po.fwpkg,i,hde.stored_len))
        chksum_enctype = hde.get_encrypt_type()
        stored_chksum = hashlib.md5()
        decrypted_chksum = hashlib.md5()
        dji_write_fwentry_head(po, i, hde, minames[i])
        fwitmfile = open("{:s}_{:s}.bin".format(po.mdprefix,minames[i]), "wb")
        fwpkgfile.seek(hde.dt_offs)
        stored_n = 0
        decrypted_n = 0
        while stored_n < hde.stored_len:
            # read block limit must be a multiplication of encryption block size
            copy_buffer = fwpkgfile.read(min(1024 * 1024, hde.stored_len - stored_n))
            if not copy_buffer:
                break
            stored_n += len(copy_buffer)
            stored_chksum.update(copy_buffer)
            if (chksum_enctype != 0) and (not hde.preencrypted):
                copy_buffer, encrypt_iv = dji_decrypt_block(copy_buffer, encrypt_key, encrypt_iv)
            fwitmfile.write(copy_buffer)
            decrypted_n += len(copy_buffer)
            decrypted_chksum.update(copy_buffer)
        fwitmfile.close()
        if (stored_chksum.hexdigest() != hde.hex_stored_md5()):
            eprint("{}: Warning: Module index {:d} stored checksum mismatch; got {:s}, expected {:s}."
              .format(po.fwpkg,i,stored_chksum.hexdigest(),hde.hex_stored_md5()))
        if (not hde.preencrypted) and (decrypted_chksum.hexdigest() != hde.hex_decrypted_md5()):
            eprint("{}: Warning: Module index {:d} decrypted checksum mismatch; got {:s}, expected {:s}."
              .format(po.fwpkg,i,decrypted_chksum.hexdigest(),hde.hex_decrypted_md5()))
            eprint("{}: Module index {:d} may be damaged due to bad decryption; use no-crypto option to leave it as-is."
              .format(po.fwpkg,i))
        if (not hde.preencrypted) and (decrypted_n != hde.decrypted_len):
            eprint("{}: Warning: decrypted size mismatch, {:d} instead of {:d}."
              .format(po.fwpkg,decrypted_n,hde.decrypted_len))
        if (po.verbose > 1):
            print("{}: Module index {:d} stored checksum {:s}".format(po.fwpkg,i,stored_chksum.hexdigest()))
    return


def dji_create(po, fwpkgfile):
    # Read headers from INI files
    (pkghead, minames) = dji_read_fwpkg_head(po)
    pkgmodules = []
    # Create module entry for each partition
    for i, miname in enumerate(minames):
        if miname == "0":
            hde = FwPkgEntry()
        else:
            hde = dji_read_fwentry_head(po, i, miname)
        pkgmodules.append(hde)
    # Write the unfinished headers
    if (po.verbose > 2):
        print("{}: File map: 0x{:08x} FwPkgHeader".format(po.fwpkg,fwpkgfile.tell()))
    fwpkgfile.write((c_ubyte * sizeof(pkghead)).from_buffer_copy(pkghead))
    for hde in pkgmodules:
        if (po.verbose > 2):
            print("{}: File map: 0x{:08x} FwPkgEntry[m{:02d}{:02d}]".format(po.fwpkg,
              fwpkgfile.tell(), getattr(hde, 'target') & 31, (getattr(hde, 'target') >> 5) & 7))
        fwpkgfile.write((c_ubyte * sizeof(hde)).from_buffer_copy(hde))
    fwpkgfile.write((c_ubyte * sizeof(c_ushort))())
    # Write module data
    for i, miname in enumerate(minames):
        hde = pkgmodules[i]
        if miname == "0":
            if (po.verbose > 0):
                print("{}: Empty module index {:d}".format(po.fwpkg,i))
            continue
        if (po.verbose > 0):
            print("{}: Copying module index {:d}".format(po.fwpkg,i))
        fname = "{:s}_{:s}.bin".format(po.mdprefix,miname)
        # Skip unused pkgmodules
        if (os.stat(fname).st_size < 1):
            eprint("{}: Warning: module index {:d} empty".format(po.fwpkg,i))
            continue
        chksum_enctype = hde.get_encrypt_type()
        epos = fwpkgfile.tell()
        # Check for data encryption
        if (chksum_enctype != 0) and (not hde.preencrypted):
            if (po.no_crypto):
                if (not po.force_continue):
                    eprint("{}: Error: Module {:d} needs encryption {:d}, but crypto is disabled."
                      .format(po.fwpkg,chksum_enctype,i))
                    exit(1)
                eprint("{}: Warning: Module {:d} needs encryption {:d}, but crypto is disabled; switching to unencrypted."
                  .format(po.fwpkg,chksum_enctype,i))
                hde.set_encrypt_type(0)
                chksum_enctype = hde.get_encrypt_type()
            elif (chksum_enctype == 1):
                encrypt_key = encrypt_aes128_key
                encrypt_iv  = encrypt_aes128_iv
            else:
                if (not po.force_continue):
                    eprint("{}: Error: Unknown encryption {:d} in module {:d}; cannot encrypt.".format(po.fwpkg,chksum_enctype,i))
                    exit(1)
                eprint("{}: Warning: Unknown encryption {:d} in module {:d}; switching to unencrypted.".format(po.fwpkg,chksum_enctype,i))
                hde.set_encrypt_type(0)
                chksum_enctype = hde.get_encrypt_type()
        # Copy partition data and compute checksum
        if (po.verbose > 2):
            print("{}: File map: 0x{:08x} FwModuleData[m{:02d}{:02d}]".format(po.fwpkg,
              epos, getattr(hde, 'target') & 31, (getattr(hde, 'target') >> 5) & 7))
        fwitmfile = open(fname, "rb")
        stored_chksum = hashlib.md5()
        decrypted_chksum = hashlib.md5()
        decrypted_n = 0
        while True:
            # read block limit must be a multiplication of encryption block size
            copy_buffer = fwitmfile.read(1024 * 1024)
            if not copy_buffer:
                break
            decrypted_chksum.update(copy_buffer)
            decrypted_n += len(copy_buffer)
            if (chksum_enctype != 0) and (not hde.preencrypted):
                copy_buffer, encrypt_iv = dji_encrypt_block(copy_buffer, encrypt_key, encrypt_iv)
            stored_chksum.update(copy_buffer)
            fwpkgfile.write(copy_buffer)
        fwitmfile.close()
        hde.dt_offs = epos
        hde.stored_len = fwpkgfile.tell() - epos
        # We do not support pre-encryption which changes length of data
        # If we need it at some point, the only way is to store decrypted_len in INI file
        hde.decrypted_len = decrypted_n
        hde.stored_md5 = (c_ubyte * 16).from_buffer_copy(stored_chksum.digest())
        if (hde.preencrypted):
            # If the file is pre-encrypted, then it has to have encryption type and MD5 set from INI file
            if (chksum_enctype == 0):
                eprint("{}: Warning: Module {:d} marked as pre-encrypted, but with no encryption type.".format(po.fwpkg,i))
            if all([ v == 0 for v in hde.decrypted_md5 ]):
                eprint("{}: Warning: Module {:d} marked as pre-encrypted, but decrypted MD5 is zeros.".format(po.fwpkg,i))
            else:
                print("{}: Module {:d} marked as pre-encrypted; decrypted MD5 accepted w/o verification.".format(po.fwpkg,i))
        else:
            # If the file is not pre-encrypted, then we should just use the MD5 we've computed
            hde.decrypted_md5 = (c_ubyte * 16).from_buffer_copy(decrypted_chksum.digest())
        pkgmodules[i] = hde
    if (po.verbose > 2):
        print("{}: File map: 0x{:08x} FwDataEnd".format(po.fwpkg, fwpkgfile.tell()))
    # Write all headers again
    fwpkgfile.seek(0,os.SEEK_SET)
    fwpkgfile.write((c_ubyte * sizeof(pkghead)).from_buffer_copy(pkghead))
    curhead_checksum = dji_calculate_crc16_part((c_ubyte * sizeof(pkghead)).from_buffer_copy(pkghead), 0x3692)
    for hde in pkgmodules:
        fwpkgfile.write((c_ubyte * sizeof(hde)).from_buffer_copy(hde))
        curhead_checksum = dji_calculate_crc16_part((c_ubyte * sizeof(hde)).from_buffer_copy(hde), curhead_checksum)
    pkghead_checksum = c_ushort(curhead_checksum)
    fwpkgfile.write((c_ubyte * sizeof(c_ushort)).from_buffer_copy(pkghead_checksum))


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('-p', '--fwpkg', default="", type=str, required=True,
          help="name of the firmware package file")

    parser.add_argument('-m', '--mdprefix', default="", type=str,
          help=("directory and file name prefix for the single decomposed firmware modules "
           "(default is base name of fwpkg with extension stripped, in working dir)"))

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

    parser.add_argument('-c', '--no-crypto', action='store_true',
          help="disable cryptography - do not encrypt/decrypt modules")

    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()

    subparser.add_argument('-x', '--extract', action='store_true',
          help="extract firmware package into modules")

    subparser.add_argument('-a', '--add', action='store_true',
          help="add module files to firmware package")

    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.fwpkg) > 0 and len(po.mdprefix) == 0:
        po.mdprefix = os.path.splitext(os.path.basename(po.fwpkg))[0]

    if po.extract:
        if (po.verbose > 0):
            print("{}: Opening for extraction".format(po.fwpkg))
        with open(po.fwpkg, 'rb') as fwpkgfile:
            dji_extract(po, fwpkgfile)

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

    else:
        raise NotImplementedError("Unsupported command.")


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