# Copyright (c) 2013-2024 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation and its affiliates
# (the "Company") and all right, title, and interest in and to the software
# product, including all associated intellectual property rights, are and
# shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
#
import sys
import os
import re
import subprocess
import platform
import time
import tempfile
import struct
import argparse
import shutil
import json

global mfa_version
global start_time
global work_dir
global tmpdir
global srcdirs
global psid2bin
global use_xz
global mxz
global debug
global mf_archive
global txz_archive
global exedir
global currdir
start_time = time.time()
mfa_version = 0x00000001
fails = []
inis = None
use_existing_dir = 1
use_existing_txz = 0
simulate = 0
work_dir = sys.argv[0]
tmpdir = None
srcdirs = []
use_xz = 1
psid2bin = []
debug = 0
mf_archive = ''
txz_archive = None
currdir = None
mxz = 1
exedir = os.path.abspath(__file__)
exedir = os.path.normpath(exedir)
# remove last directory from a file path
exedir = os.path.dirname(exedir)
BIG_ENDIAN_4_BYTE = "!I".encode('utf-8')
BIG_ENDIAN_2_BYTE = "!H".encode('utf-8')
BIG_ENDIAN_1_BYTE = "!B".encode('utf-8')
LITTLE_ENDIAN_4_BYTE = "<I".encode('utf-8')


def xzcmd():
    global debug
    xzRunCmd = None
    if platform.system() == 'Linux':
        xzRunCmd = 'xz'
    else:
        xzRunCmd = os.path.normpath(exedir + "/xz.exe")
    # set compression preset to 8 to allow for efficient compression of 16MB images
    xzRunCmd += ' -8'
    if mxz:
        xzRunCmd += ' -T' + str(mxz)
    if debug:
        print("DEBUG xzcmd: " + xzRunCmd)
    return xzRunCmd


def copycmd():
    global debug
    copycmd = None
    if platform.system() == 'Linux':
        copycmd = 'cp'
    else:
        copycmd = 'copy'
    if debug:
        print("DEBUG copycmd: " + copycmd)
    return copycmd


def getCurrentDir():
    global debug
    currdir = None
    currdir = os.getcwd()
    if debug:
        print("DEBUG getCurrentDir: " + currdir)
    return currdir


def getDirListing(path):
    global debug
    cmd = ""
    if platform.system() == 'Linux':
        cmd = '/bin/ls ' + path
    else:
        cmd = 'dir /B ' + path
    dir_list = os.popen(cmd).readlines()
    dir_list = [entry.strip() for entry in dir_list]
    for i in range(len(dir_list)):
        while '/' in dir_list[i]:
            dir_list[i] = dir_list[i].split('/', 1)[-1]
    if debug:
        print("DEBUG getDirListing: CMD: " + cmd)
        print("DEBUG getDirListing: " + str(dir_list))
    return dir_list


def mktempDir():
    tmpdir = tempfile.mkdtemp()
    return tmpdir


def usage():
    print("mlx_mfa_gen -p <package name> -s <source_dir>")
    print("The -s option may be used more than once to specify more than one source directory.")


def copy_bins(srcdir, tmpdir, offset, metadata):
    global orig_file
    global debug
    if debug:
        print("DEBUG copy_bins: " + srcdir + " " + tmpdir + " " + str(offset))
    files = getDirListing(os.path.normpath(srcdir))
    bins = []
    cntr = 0
    for i in range(0, len(files)):
        if files[i].endswith(".fwpkg"):
            if metadata is None or files[i] not in metadata:
                raise Exception("%s requires valid json metadata for info" % files[i])
        elif not files[i].endswith(".bin"):
            continue
        cntr += 1
        fname = '%04d' % (i + 1 + offset)
        orig_file[fname] = files[i]
        cpc = copycmd() + ' ' + os.path.join(os.path.normpath(srcdir), files[i]) + ' ' + os.path.join(os.path.normpath(tmpdir), fname)  # TODO - replace with shutil
        default_var = subprocess.Popen(cpc, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        _, stderr = default_var.communicate()
        if stderr:
            raise Exception("Failed to copy " + files[i] + " to " + os.path.normpath(tmpdir + "/" + fname) + '' + "\n Error is:" + stderr)
        else:
            sys.stdout.write("Files copied: %s\r" % str(cntr))
            sys.stdout.flush()
        bins.append(files[i])
    return bins


def rmpath(f):
    # check if the string f ends with a forward slash followed by zero or more characters that are not forward slashes
    res = re.match('\\/([^\\/]*)$', f)
    if res:
        f = res.group(1)
    return f


def query_images(tmpdir, metadata):
    def get_field_from_bin_info(info, field):
        val = info.get(field)
        if val is None:
            raise Exception("Field '%s' is missing in metadata" % field)
        return val

    global orig_file
    global debug
    global psid2bin
    psid2fl = {}
    res = 0
    bins = getDirListing(os.path.normpath(tmpdir + "/*"))
    offset = 0
    cntr = 0
    try:
        OUTocF = open(os.path.normpath(tmpdir + "/TOC_FILE"), "wb")
    except IOError:
        sys.stderr.write("Couldn't open " + os.path.normpath(tmpdir + "/TOC_FILE") + " for writing!")
    try:
        OUTF = open(os.path.normpath(tmpdir + "/DATA_FILE"), "wb")
    except IOError:
        sys.stderr.write("Couldn't open " + os.path.normpath(tmpdir + "/DATA_FILE") + " for writing!")
    for i in range(0, len(bins)):
        version = ''
        psid = ''
        pn = ''
        description = ''
        pxe_version = ''
        clp_version = ''
        uefi_version = ''
        fcode_version = ''
        branch = ''
        image_type = 1  # Legacy for FW image
        global orig_file
        if orig_file[bins[i]].endswith("fwpkg"):
            image_type = 1  # TODO - add new type for fwpkg once mlxfwmanager supports multiple image types
            bin_info = metadata[orig_file[bins[i]]]
            try:
                # Mandatory fields in json
                version = get_field_from_bin_info(bin_info, 'version')
                psid = get_field_from_bin_info(bin_info, 'psid')
                pn = get_field_from_bin_info(bin_info, "part_number")
                description = get_field_from_bin_info(bin_info, "description")
            except Exception as e:
                raise Exception("%s for file '%s'" % (str(e), orig_file[bins[i]]))
            # Optional fields in json
            if 'clp_version' in bin_info:
                clp_version = bin_info.get('clp_version')
            if 'uefi_version' in bin_info:
                uefi_version = bin_info.get('uefi_version')
            if 'fcode_version' in bin_info:
                fcode_version = bin_info.get('fcode_version')
            if 'pxe_version' in bin_info:
                pxe_version = bin_info.get('pxe_version')
            cntr += 1
            sys.stdout.write("\rFiles queried: %s" % str(cntr))
            sys.stdout.flush()
        else:  # FW image flow
            cmd = 'flint -i ' + os.path.normpath(tmpdir + "/" + bins[i]) + ' q'
            process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output, exit_code = process.communicate()
            if process.returncode:
                raise Exception("Failed to query " + bins[i] + "\n Error is: " + output)
            else:
                for t in output.splitlines():
                    # search for the occurrence of "fw version:", one or more whitespace characters (\s+),
                    # and then capture one or more non-whitespace characters (\S+).
                    # This will result in getting the fw version.
                    res = re.search(r'fw version:\s+(\S+)', t.decode(), re.IGNORECASE)
                    if res:
                        version = res.group(1)
                    # search for the occurrence of "psid:", one or more whitespace characters (\s+),
                    # and then capture one or more non-whitespace characters (\S+).
                    # This will result in getting the psid.
                    res = re.search(r'psid:\s+(\S+)', t.decode(), re.IGNORECASE)
                    if res:
                        psid = res.group(1)
                    # search for the occurrence of "type", zero or more whitespace characters (\s*),
                    # an '=' sign, zero or more whitespace characters (\s*), then 'CLP', while ignoring case and searching multiline text.
                    # This will result in getting the CLP type text.
                    res = re.search(r'type\s*=\s*CLP', t.decode(), re.IGNORECASE | re.MULTILINE)
                    if res:
                        # search for the occurrence of the word "version",
                        # followed by any number of characters that are not an equals sign ([^\=]*),
                        # then an equals sign (=), and finally captures one or more non-whitespace characters ((\S+))
                        searchRes = re.search(r'version[^\=]*\=(\S+)', t.decode())
                        if searchRes:
                            clp_version = searchRes.group(1)
                    # search for the occurrence of the word "type=UEFI", while ignoring case
                    res = re.search(r'type=UEFI', t.decode(), re.IGNORECASE)
                    if res:
                        # search for the occurrence of the word "version=", while ignoring case
                        # finally captures one or more non-whitespace characters (\S+)
                        searchRes = re.search(r'version\=(\S+)', t.decode())
                        if searchRes:
                            uefi_version = searchRes.group(1)
                    # search for the occurrence of the word "type=FCODE", while ignoring case
                    res = re.search(r'type=FCODE', t.decode(), re.IGNORECASE)
                    if res:
                        # search for the occurrence of the word "version=", while ignoring case
                        # finally captures one or more non-whitespace characters (\S+)
                        searchRes = re.search(r'version\=(\S+)', t.decode())
                        if searchRes:
                            fcode_version = searchRes.group(1)
                    # search for the occurrence of the word "type=PXE", while ignoring case
                    res = re.search(r'type=PXE', t.decode(), re.IGNORECASE)
                    if res:
                        # search for the occurrence of the word "version=", while ignoring case
                        # finally captures one or more non-whitespace characters (\S+)
                        searchRes = re.search(r'version\=(\S+)', t.decode())
                        if searchRes:
                            pxe_version = searchRes.group(1)
            cmd = 'flint -i ' + os.path.normpath(tmpdir + "/" + bins[i]) + ' dc'
            process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output, exit_code = process.communicate()
            if process.returncode:
                if debug:
                    print("flint dc command failed, trying query full\n")
                cmd = 'flint -i ' + os.path.normpath(tmpdir + "/" + bins[i]) + ' q full'
                process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                output, exit_code = process.communicate()
                if process.returncode:
                    raise Exception(sys.stderr.write("Failed to query %s for PN\n" % bins[i]))
                else:
                    # search for the occurrence of "Part Number:", one or more whitespace characters (\s+),
                    # and then capture one or more non-whitespace characters (\S+).
                    # This will result in getting the Part Number.
                    res = re.search(r'Part Number:\s+(\S+)', output.decode(), re.IGNORECASE)
                    if res:
                        pn = res.group(1)
                    # search for the occurrence of "Description:", one or more whitespace characters (\s+),
                    # and then capture one or more non-whitespace characters (\S.*), follow by zero or more whitespace characters (\s+),
                    # finally end of string ('$').
                    # This will result in getting the Part Number.
                    res = re.search(r'Description:\s+(\S.*)(\s*)$', output.decode(), re.MULTILINE | re.IGNORECASE)
                    if res:
                        description = res.group(1)
                    cntr += 1
                    sys.stdout.write("\rFiles queried: %s" % str(cntr))
                    sys.stdout.flush()
            else:
                if debug:
                    print('flint dc command succeed')
                # search for the occurrence of "Name", one or more whitespace characters (\s+),
                # followed by '=', one or more white space characters,
                # and then capture one or more non-whitespace characters (\S+).
                # This will result in getting the Part Number.
                res = re.search(r'Name\s+=\s+(\S+)', output.decode(), re.IGNORECASE)
                if res:
                    pn = res.group(1)
                # search for the occurrence of "Description", one or more whitespace characters (\s+),
                # followed by '=', one or more white space characters,
                # and then capture one or more non-whitespace characters (\S.*),
                # followed by zero or more whitespace characters.
                # This will result in getting the Part Number.
                res = re.search(r'Description\s+\=\s+(\S.*)(\s*)$', output.decode(), re.MULTILINE | re.IGNORECASE)
                if res:
                    description = res.group(1)
                cntr += 1
                sys.stdout.write("\rFiles queried: %s" % str(cntr))
                sys.stdout.flush()
        if debug:
            print(psid + ' ' + version + ' ' + pn + ' ' + description)
        binsize = (os.stat(os.path.normpath(tmpdir + "/" + bins[i]))).st_size
        data_offset = struct.pack("!I", offset)  # 32 bits
        data_size = struct.pack(BIG_ENDIAN_4_BYTE, binsize)
        subimage_type = struct.pack(BIG_ENDIAN_2_BYTE, 1)  # 1: FW Image
        # pack to a little endian 1 byte
        resvd = struct.pack("B".encode(), 0)
        resvd1 = struct.pack(BIG_ENDIAN_2_BYTE, offset >> 32)  # 16 MSB bits added to data_offset when parsing the MFA
        num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 3)
        ver0 = ver1 = ver2 = ver3 = None
        # check if string is of type number.number.number (having a '.' between the numbers)
        res = re.match(r'(\d+)\.(\d+)\.(\d+)', version)
        # check if string ends with either _XXXX or '-YYY' with relevant length
        res2 = re.match(r'.+((_\d{4})|(-\d{3}))', version)
        if res:
            ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(1)))
            ver1 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(2)))
            ver2 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(3)))
            ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
        elif res2:
            ver0 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            ver1 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            ver2 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            version_length = len(version)
            branch = struct.pack(str(version_length) + 's', version.encode('utf-8'))
        else:
            raise RuntimeError("Failed to extract version " + str(version) + "number " + str(i) + ":" + bins[i] + "\n")
        metadatasz = 0
        metadata_size = struct.pack(BIG_ENDIAN_2_BYTE, metadatasz)
        OUTocF.write(data_offset)
        OUTocF.write(data_size)
        OUTocF.write(subimage_type)
        OUTocF.write(resvd)
        OUTocF.write(num_ver_fields)
        OUTocF.write(ver0)
        OUTocF.write(ver1)
        OUTocF.write(ver2)
        OUTocF.write(ver3)
        OUTocF.write(resvd1)
        OUTocF.write(metadata_size)
        metad_size = struct.pack(BIG_ENDIAN_2_BYTE, 0)
        if pxe_version:
            data_offset = struct.pack(BIG_ENDIAN_4_BYTE, 0xffffffff)
            data_size = struct.pack(BIG_ENDIAN_4_BYTE, 0)
            subimage_type = struct.pack(BIG_ENDIAN_2_BYTE, 0x110)
            resvd = struct.pack(BIG_ENDIAN_1_BYTE, 0)
            num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 3)
            ver0 = ver1 = ver2 = ver3 = None
            # check if string is of type number.number.number (having a '.' between the numbers)
            res = re.match(r'(\d+)\.(\d+)\.(\d+)', pxe_version)
            if res:
                ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(1)))
                ver1 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(2)))
                ver2 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(3)))
                ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            else:
                raise RuntimeError("Failed to extract pxe version " + pxe_version + "number " + i + ":" + bins[i] + "\n")
            OUTocF.write(data_offset)
            OUTocF.write(data_size)
            OUTocF.write(subimage_type)
            OUTocF.write(resvd)
            OUTocF.write(num_ver_fields)
            OUTocF.write(ver0)
            OUTocF.write(ver1)
            OUTocF.write(ver2)
            OUTocF.write(ver3)
            OUTocF.write(resvd)
            OUTocF.write(resvd)
            OUTocF.write(metad_size)
        if clp_version:
            data_offset = struct.pack(BIG_ENDIAN_4_BYTE, 0xffffffff)
            data_size = struct.pack(BIG_ENDIAN_4_BYTE, 0)
            subimage_type = struct.pack(BIG_ENDIAN_2_BYTE, 0x101)
            resvd = struct.pack(BIG_ENDIAN_1_BYTE, 0)
            num_ver_fields = None
            ver0 = ver1 = ver2 = ver3 = None
            # check if string is of type number.number.number (having a '.' between the numbers)
            res = re.match(r'(\d+)\.(\d+)\.(\d+)', clp_version)
            # check if string is of type numbers
            res2 = re.match(r'(\d+)', clp_version)
            if res:
                num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 3)
                ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(1)))
                ver1 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(2)))
                ver2 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(3)))
                ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            elif res2:
                num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 1)
                ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res2.group(1)))
                ver1 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
                ver2 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
                ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            else:
                raise RuntimeError("Failed to extract clp version " + clp_version + "number" + i + ": " + bins[i] + "\n")
            OUTocF.write(data_offset)
            OUTocF.write(data_size)
            OUTocF.write(subimage_type)
            OUTocF.write(resvd)
            OUTocF.write(num_ver_fields)
            OUTocF.write(ver0)
            OUTocF.write(ver1)
            OUTocF.write(ver2)
            OUTocF.write(ver3)
            OUTocF.write(resvd)
            OUTocF.write(resvd)
            OUTocF.write(metad_size)
        metad_size = struct.pack(BIG_ENDIAN_2_BYTE, 0)
        if uefi_version:
            data_offset = struct.pack(BIG_ENDIAN_4_BYTE, 0xffffffff)
            data_size = struct.pack(BIG_ENDIAN_4_BYTE, 0)
            subimage_type = struct.pack(BIG_ENDIAN_2_BYTE, 0x111)
            resvd = struct.pack(BIG_ENDIAN_1_BYTE, 0)
            num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 3)
            ver0 = ver1 = ver2 = ver3 = None
            # check if string is of type number.number.number (having a '.' between the numbers)
            res = re.match(r'(\d+)\.(\d+)\.(\d+)', uefi_version)
            if res:
                ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(1)))
                ver1 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(2)))
                ver2 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(3)))
                ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            else:
                raise RuntimeError("Failed to extract uefi version " + uefi_version + "number" + i + ":" + bins[i] + "\n")
            OUTocF.write(data_offset)
            OUTocF.write(data_size)
            OUTocF.write(subimage_type)
            OUTocF.write(resvd)
            OUTocF.write(num_ver_fields)
            OUTocF.write(ver0)
            OUTocF.write(ver1)
            OUTocF.write(ver2)
            OUTocF.write(ver3)
            OUTocF.write(resvd)
            OUTocF.write(resvd)
            OUTocF.write(metad_size)
        metad_size = struct.pack(BIG_ENDIAN_2_BYTE, 0)
        if fcode_version:
            data_offset = struct.pack(BIG_ENDIAN_4_BYTE, 0xffffffff)
            data_size = struct.pack(BIG_ENDIAN_4_BYTE, 0)
            subimage_type = struct.pack(BIG_ENDIAN_2_BYTE, 0x121)
            resvd = struct.pack(BIG_ENDIAN_1_BYTE, 0)
            num_ver_fields = struct.pack(BIG_ENDIAN_1_BYTE, 3)
            ver0 = ver1 = ver2 = ver3 = None
            # check if string is of type number.number.number (having a '.' between the numbers)
            res = re.match(r'(\d+)\.(\d+)\.(\d+)', fcode_version)
            if res:
                ver0 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(1)))
                ver1 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(2)))
                ver2 = struct.pack(BIG_ENDIAN_2_BYTE, int(res.group(3)))
                ver3 = struct.pack(BIG_ENDIAN_2_BYTE, 0)
            else:
                raise RuntimeError("Failed to extract fcode version " + fcode_version + "number" + i + ":" + bins[i] + "\n")
            OUTocF.write(data_offset)
            OUTocF.write(data_size)
            OUTocF.write(subimage_type)
            OUTocF.write(resvd)
            OUTocF.write(num_ver_fields)
            OUTocF.write(ver0)
            OUTocF.write(ver1)
            OUTocF.write(ver2)
            OUTocF.write(ver3)
            OUTocF.write(resvd)
            OUTocF.write(resvd)
            OUTocF.write(metad_size)
        try:
            INF = open(os.path.normpath(tmpdir + "/" + bins[i]), "rb")
        except BaseException:
            raise RuntimeError("Couldn't open " + bins[i] + "for reading\n")
        buffer = None
        buffer_size = 65536
        while True:
            buffer = INF.read(buffer_size)
            if not buffer:
                break
            OUTF.write(buffer)
        INF.close()
        offset += binsize
        psid2bin_entry = {'psid': psid,
                          'arch': 'noarch',
                          'fwfile': rmpath(bins[i]),
                          'exprom_file': '',
                          'opns': pn,
                          'version': version,
                          'size': (os.stat(os.path.normpath(tmpdir + "/" + bins[i]))).st_size,
                          'offset': binsize,
                          'metadata_size': metadatasz,
                          'description': description,
                          'clp_version': clp_version,
                          'pxe_version': pxe_version,
                          'uefi_version': uefi_version,
                          'fcode_version': fcode_version,
                          'branch': branch,
                          'image_type': image_type
                          }
        psid2bin.append(psid2bin_entry)
        if psid not in psid2fl:
            psid2fl[psid] = rmpath(bins[i])
        else:
            raise Exception("PSID=" + psid + " appears more than once, found in binary files: " + orig_file[psid2fl[psid]] + " and " + orig_file[rmpath(bins[i])] + '\n')
    print("")  # New line following "Files queried:" refreshed counter print
    OUTocF.close()
    OUTF.close()


def create_map_file(tmp_dir):
    global psid2bin
    try:
        OUT = open(os.path.normpath(tmp_dir + "/MAP_FILE"), "wb")
    except BaseException:
        raise RuntimeError("Failed to create map file\n")
    sz = len(psid2bin)
    sum = 0
    for i in range(0, sz):
        # pack to a string of length 32
        board_type_id = struct.pack("32s", psid2bin[i]["psid"].encode())
        numfiles = 1
        if psid2bin[i]["pxe_version"]:
            numfiles += 1
        if psid2bin[i]["clp_version"]:
            numfiles += 1
        if psid2bin[i]["uefi_version"]:
            numfiles += 1
        if psid2bin[i]["fcode_version"]:
            numfiles += 1
        files_per_psid = struct.pack(BIG_ENDIAN_1_BYTE, numfiles)
        resvd = struct.pack(BIG_ENDIAN_1_BYTE, 0)
        metadata_sz = struct.pack(BIG_ENDIAN_2_BYTE, 152)
        OUT.write(board_type_id)
        OUT.write(files_per_psid)
        OUT.write(resvd)
        OUT.write(metadata_sz)
        metadata_hdr0 = struct.pack(BIG_ENDIAN_1_BYTE, 0x1)
        metadata_hdr1 = struct.pack(BIG_ENDIAN_1_BYTE, 0x0)
        metadata_hdr2 = None  # number of keys in metadata
        # key0 is PN
        # key1 is description
        # key2 is branch - optional
        md_length2 = len(psid2bin[i]["branch"]) + 1
        if md_length2 > 1:
            metadata_hdr2 = struct.pack(BIG_ENDIAN_2_BYTE, 3)
        else:
            metadata_hdr2 = struct.pack(BIG_ENDIAN_2_BYTE, 2)
        OUT.write(metadata_hdr0)
        OUT.write(metadata_hdr1)
        OUT.write(metadata_hdr2)
        # pack to a string of length 3
        md_key0 = struct.pack('3s', "PN".encode())
        md_length0 = len(psid2bin[i]["opns"]) + 1
        # pack to a string of length 12
        md_key1 = struct.pack('12s', "DESCRIPTION".encode())
        md_length1 = len(psid2bin[i]["description"]) + 1
        md_length_left = 152 - md_length0 - md_length1 - 3 - 12 - 4
        md_value1 = None
        if md_length2 > 1:  # check if branch is presented
            md_length_left = md_length_left - md_length2 - 7
        if md_length_left < 0:
            # check if size of the keys combined exceeds total size
            md_length1 += md_length_left  # subtracts from description
            md_length_left = 0
            tmp_md_length1 = md_length1 - 4  # place for "..." and null
            # pack to a string of length tmp_md_length1 + 3 length for '...'
            tmp_md_value1 = struct.pack(str(tmp_md_length1) + 's3s', psid2bin[i]["description"].encode(), "...".encode())
            # pack to a string of length md_length1
            md_value1 = struct.pack(str(md_length1) + 's', tmp_md_value1)
        else:
            # pack to a string of length md_length1
            md_value1 = struct.pack(str(md_length1) + 's', psid2bin[i]["description"].encode())
        # pack to a string of length md_length0
        md_value0 = struct.pack(str(md_length0) + 's', psid2bin[i]["opns"].encode())
        # pack to a string of length md_length_left
        md_value_left = struct.pack(str(md_length_left) + 's', ''.encode())
        OUT.write(md_key0)
        OUT.write(md_value0)
        OUT.write(md_key1)
        OUT.write(md_value1)
        if md_length2 > 1:
            # adding branch if presented
            md_key2 = struct.pack('7s', "BRANCH".encode())
            md_value2 = struct.pack(str(md_length2) + 's', psid2bin[i]["branch"])
            OUT.write(md_key2)
            OUT.write(md_value2)
        OUT.write(md_value_left)
        file_cnt = 1
        toc_offset = struct.pack(BIG_ENDIAN_4_BYTE, sum)
        image_type = struct.pack(BIG_ENDIAN_2_BYTE, psid2bin[i]["image_type"])
        # pack to a string of length 32
        select_tag = struct.pack('32s', ''.encode())
        OUT.write(toc_offset)
        OUT.write(image_type)
        OUT.write(resvd)
        group_id = struct.pack(BIG_ENDIAN_1_BYTE, file_cnt)
        OUT.write(group_id)
        OUT.write(select_tag)
        sum += psid2bin[i]["metadata_size"] + 24
        if psid2bin[i]["pxe_version"]:
            file_cnt += 1
            toc_offset = struct.pack(BIG_ENDIAN_4_BYTE, sum)
            image_type = struct.pack(BIG_ENDIAN_2_BYTE, 1)
            # pack to a string of length 32
            select_tag = struct.pack('32s', ''.encode())
            OUT.write(toc_offset)
            OUT.write(image_type)
            OUT.write(resvd)
            group_id = struct.pack(BIG_ENDIAN_1_BYTE, file_cnt)
            OUT.write(group_id)
            OUT.write(select_tag)
            sum += 24
        if psid2bin[i]["clp_version"]:
            file_cnt += 1
            toc_offset = struct.pack(BIG_ENDIAN_4_BYTE, sum)
            image_type = struct.pack(BIG_ENDIAN_2_BYTE, 1)
            # pack to a string of length 32
            select_tag = struct.pack('32s', ''.encode())
            OUT.write(toc_offset)
            OUT.write(image_type)
            OUT.write(resvd)
            group_id = struct.pack(BIG_ENDIAN_1_BYTE, file_cnt)
            OUT.write(group_id)
            OUT.write(select_tag)
            sum += 24
        if psid2bin[i]["uefi_version"]:
            file_cnt += 1
            toc_offset = struct.pack(BIG_ENDIAN_4_BYTE, sum)
            image_type = struct.pack(BIG_ENDIAN_2_BYTE, 1)
            # pack to a string of length 32
            select_tag = struct.pack('32s', ''.encode())
            OUT.write(toc_offset)
            OUT.write(image_type)
            OUT.write(resvd)
            group_id = struct.pack(BIG_ENDIAN_1_BYTE, file_cnt)
            OUT.write(group_id)
            OUT.write(select_tag)
            sum += 24
        if psid2bin[i]["fcode_version"]:
            file_cnt += 1
            toc_offset = struct.pack(BIG_ENDIAN_4_BYTE, sum)
            image_type = struct.pack(BIG_ENDIAN_2_BYTE, 1)
            # pack to a string of length 32
            select_tag = struct.pack('32s', ''.encode())
            OUT.write(toc_offset)
            OUT.write(image_type)
            OUT.write(resvd)
            group_id = struct.pack(BIG_ENDIAN_1_BYTE, file_cnt)
            OUT.write(group_id)
            OUT.write(select_tag)
            sum += 24
    OUT.close()
    return 0


def swap_be32(x):
    x = (((x << 24) & 0xFF000000) |
         ((x << 8) & 0x00FF0000) |
         ((x >> 8) & 0x0000FF00) |
         ((x >> 24) & 0x000000FF))
    return x


def create_mfa(mlxarchive, mapfile, tocfile, datafile):
    global use_xz
    # pack to a string of length 4
    header_signature = struct.pack("4s", "MFAR".encode())
    header_version = struct.pack(LITTLE_ENDIAN_4_BYTE, swap_be32(mfa_version))
    reserved = struct.pack(LITTLE_ENDIAN_4_BYTE, 0)
    try:
        OUTF = open(os.path.normpath(str(mlxarchive)), "wb")
    except BaseException:
        raise RuntimeError("Couldn't open " + mlxarchive + " for writing!\n")
    OUTF.write(header_signature)
    OUTF.write(header_version)
    OUTF.write(reserved)
    OUTF.write(reserved)
    # MAP section
    map_tlv_type = struct.pack(BIG_ENDIAN_1_BYTE, 0x1)
    map_tlv_flag = struct.pack(BIG_ENDIAN_1_BYTE, use_xz)
    reserved = struct.pack(BIG_ENDIAN_1_BYTE, 0x0)
    map_tlv_size = struct.pack(BIG_ENDIAN_4_BYTE, os.stat(mapfile).st_size)
    OUTF.write(map_tlv_type)
    OUTF.write(reserved)
    OUTF.write(reserved)
    OUTF.write(map_tlv_flag)
    OUTF.write(map_tlv_size)
    buffer = ''
    buffer_size = 65536
    try:
        INF = open(os.path.normpath(mapfile), "rb")
    except BaseException:
        raise RuntimeError("Couldn't open " + mapfile + " for reading!\n")
    while True:
        buffer = INF.read(buffer_size)
        if not buffer:
            break
        OUTF.write(buffer)
    INF.close()
    # TOC section
    toc_tlv_type = struct.pack(BIG_ENDIAN_1_BYTE, 0x2)
    toc_tlv_flag = struct.pack(BIG_ENDIAN_1_BYTE, use_xz)
    toc_tlv_size = struct.pack(BIG_ENDIAN_4_BYTE, os.stat(tocfile).st_size)
    OUTF.write(toc_tlv_type)
    OUTF.write(reserved)
    OUTF.write(reserved)
    OUTF.write(toc_tlv_flag)
    OUTF.write(toc_tlv_size)
    buffer = None
    buffer_size = 65536
    try:
        INF = open(os.path.normpath(tocfile), "rb")
    except BaseException:
        raise RuntimeError("Couldn't open " + tocfile + " for reading!\n")
    while True:
        buffer = INF.read(buffer_size)
        if not buffer:
            break
        OUTF.write(buffer)
    INF.close()
    # Data section
    data_tlv_type = struct.pack(BIG_ENDIAN_1_BYTE, 0x3)
    data_tlv_flag = struct.pack(BIG_ENDIAN_1_BYTE, use_xz)
    data_tlv_size = struct.pack(BIG_ENDIAN_4_BYTE, os.stat(datafile).st_size)
    OUTF.write(data_tlv_type)
    OUTF.write(reserved)
    OUTF.write(reserved)
    OUTF.write(data_tlv_flag)
    OUTF.write(data_tlv_size)
    buffer = None
    buffer_size = 65536
    try:
        INF = open(os.path.normpath(datafile), "rb")
    except BaseException:
        raise RuntimeError("Couldn't open " + os.path.normpath(datafile) + " for reading!\n")
    while True:
        buffer = INF.read(buffer_size)
        if not buffer:
            break
        OUTF.write(buffer)
    INF.close()
    OUTF.close()
    return 0


def get_crc(file):
    global debug
    crc = 0
    cmd = ''
    mlxfwmanager = None
    if platform.system() == 'Linux':
        mlxfwmanager = os.path.normpath("mlxfwmanager")
    else:
        mlxfwmanager = os.path.normpath(exedir + "/mlxfwmanager.exe")
    cmd = mlxfwmanager + ' --crc -i ' + os.path.normpath(file)
    if debug:
        print(cmd)
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    output, exit_code = process.communicate()
    if exit_code:
        raise RuntimeError("Failed to run ", cmd)
    # check if string has CRC32 followed by one whitespace or more, followed by '=',
    # followed by one whitespace or more and then captures one or more non-whitespace characters.
    res = re.match('CRC32\\s+=\\s+(\\S+)', output.decode())
    if res:
        crc = res.group(1)
    return crc


def add_crc(mlxarchive, argcrc):
    mlxarchive_crc = struct.pack(BIG_ENDIAN_4_BYTE, int(argcrc, 16))
    try:
        OUTF = open(os.path.normpath(mlxarchive), "ab")
    except BaseException:
        raise RuntimeError("Couldn't open " + mlxarchive + " for writing!\n")
    OUTF.write(mlxarchive_crc)
    OUTF.close()


def checkRefineArgs(arg, option):
    fixedPath = None
    if not arg:
        raise RuntimeError("Missing argument for option " + option + "\n")
    # match string that start with zero or more whitespace characters followed by a hyphen (-), followed by zero or more of any character (.).
    if re.match('^\\s*-.*', arg):
        raise RuntimeError("Missing argument for option " + option + "\n")
    fixedPath = os.path.abspath(arg)
    if not fixedPath:
        raise RuntimeError("Could not evaluate path " + arg + ", for option " + option + " check if path exist\n")
    if not os.path.isdir(fixedPath):
        raise RuntimeError("Directory " + arg + " does not exist\n")
    return fixedPath
# -- Command-Line Arguments  ---------------------------------------------------


def main():
    global debug
    global orig_file
    global use_xz
    global mxz
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", help="source dir of bin files")
    parser.add_argument("-p", help="package name of resulting mfa file")
    parser.add_argument("-noxz", help="do not use xz for compression", action="store_true")
    parser.add_argument("-debug", help="print debug info", action="store_true")
    parser.add_argument("-mxz", help="use xz multithreading with X threads. (provide a number). Default is set to 1.")
    parser.add_argument("--json_metadata", help="json file specifying information on unknown binaries format")
    parser.add_argument("--keep_temp_files", action="store_true", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.s:
        srcdirs.append(checkRefineArgs(args.s, "-s"))
    if args.p:
        if re.match('\\.mfa$', args.p):
            args.p = args.p + '.mfa'
        mf_archive = args.p
        txz_archive = args.p
        # replace .mfa with .txz
        txz_archive = re.sub('\\.mfa', '\\.txz', txz_archive)
    else:
        sys.stderr.write("-E- Must specify package name\n")
        usage()
        sys.exit(1)
    if args.noxz:
        use_xz = 0
    if args.debug:
        debug = 1
    if args.mxz:
        mxz = args.mxz
    metadata = None
    if args.json_metadata:
        with open(args.json_metadata) as f:
            metadata = json.load(f)
    if srcdirs == 0:
        sys.stderr.write("-E- Must specify source dir")
        usage()
        raise RuntimeError(1)
    if not mf_archive:
        raise Exception("-E- Must specify package name")

    tmpdir = mktempDir()
    if debug:
        print("Temp dir created: %s" % tmpdir)
    orig_file = {}
    cps = 0
    total_cps = 0
    for srcdir in srcdirs:
        print("Adding bins from %s" % srcdir)
        cps = copy_bins(srcdir, tmpdir, total_cps, metadata)
        total_cps += len(cps)
    if not total_cps:
        raise RuntimeError('Failed to find any binary file, Please specify a directory with bins!\n')

    print("\nQuerying images ...")
    query_images(tmpdir, metadata)

    rc = create_map_file(tmpdir)
    if rc:
        raise RuntimeError('-E- Failed to create map file\n')
    ext = ''
    if use_xz:
        print("Compressing ... (this may take a minute)\n")
        ext = '.xz'
        cmd = xzcmd() + ' -q --check=crc32 ' + os.path.normpath(tmpdir + "/MAP_FILE")
        if debug:
            print(cmd)
        result = os.system(cmd)
        if result:
            raise RuntimeError('Failed to run xz on ' + os.path.normpath(tmpdir + "/MAP_FILE"))
        cmd = xzcmd() + ' -q --check=crc32 ' + os.path.normpath(tmpdir + "/TOC_FILE")
        if debug:
            print(cmd)
        result = os.system(cmd)
        if result:
            raise RuntimeError('Failed to run xz on ' + os.path.normpath(tmpdir + "/TOC_FILE"))
        cmd = xzcmd() + ' -q --check=crc32 ' + os.path.normpath(tmpdir + "/DATA_FILE")
        if debug:
            print(cmd)
        result = os.system(cmd)
        if result:
            raise RuntimeError('Failed to run xz on ' + os.path.normpath(tmpdir + "/DATA_FILE"))
    argcrc = 0x0
    rc = create_mfa(mf_archive, os.path.normpath(tmpdir + "/MAP_FILE" + ext), os.path.normpath(tmpdir + "/TOC_FILE" + ext), os.path.normpath(tmpdir + "/DATA_FILE" + ext))
    argcrc = get_crc(mf_archive)
    if not argcrc:
        raise RuntimeError('-E- Failed to calculate crc\n')
    add_crc(mf_archive, argcrc)
    print("\nArchive: " + mf_archive)
    if not args.keep_temp_files:
        try:
            shutil.rmtree(tmpdir)
        except Exception as e:
            print("-W- Failed to remove temp dir - %s. Error: %s" % (str(tmpdir), str(e)))
    else:
        print("-I- Temp files can be found at - %s" % tmpdir)
    end_time = time.time()
    mins = int((end_time - start_time) / 60)
    seconds = int((end_time - start_time) - mins * 60)
    print("Total time: %d.m %d" % (mins, seconds))
    sys.exit(rc)


# --------------------------------------------------------------------------------------------
if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        import traceback
        traceback.print_exc()
        print("-E- %s." % str(e))
        sys.exit(1)
