"""fwf.py - Firmware File collector:

Runs on Python 2.7.x with following module(s) installed:
    pip install uefi_firmware
        https://github.com/theopolis/uefi-firmware-parser

Copyright (C) 2016-2019 Phoenix Technologies Ltd. All Rights Reserved. """

import struct
import collections
import uefi_firmware

def sguid(b, big=False):
    '''RFC4122 binary GUID as string.'''
    if b is None or len(b) != 16:
        return ""
    a, b, c, d = struct.unpack("%sIHH8s" % (">" if big else "<"), b)
    d = ''.join('%02X' % ord(c) for c in d)
    return "%08X-%04X-%04X-%s-%s" % (a, b, c, d[:4], d[4:])

#EFI_SECTION_TYPES = {
#    0x01: ("Compression",               "compressed",   None),
#    0x02: ("Guid Defined",              "guid",         None),
#    0x03: ("Disposable",                "disposable",   None),
#    0x10: ("PE32 image",                "pe",           "PE32"),
#    0x11: ("PE32+ PIC image",           "pic.pe",       "PIC"),
#    0x12: ("Terse executable (TE)",     "te",           "TE"),
#    0x13: ("DXE dependency expression", "dxe.depex",    "DXE_DEPEX"),
#    # Added from previous code (not in Phoenix spec
#    0x14: ("Version section",           "version",      "VERSION"),
#    0x15: ("User interface name",       "ui",           "UI"),
#
#    0x16: ("IA-32 16-bit image",        "ia32.16bit",   "COMPAT16"),
#    0x17: ("Firmware volume image",     "fv",           "FV_IMAGE"),
#    # See FdfParser.py in EDKII's GenFds
#    0x18: ("Free-form GUID",            "freeform.guid", "SUBTYPE_GUID"),
#    0x19: ("Raw",                       "raw",          "RAW"),
#    0x1b: ("PEI dependency expression", "pie.depex",    "PEI_DEPEX"),
#    0x1c: ("SMM dependency expression", "smm.depex",    "SMM_DEPEX")
executable_sction_types = {0x10, 0x11, 0x12}

complain_duplicate_file = False
found_files = collections.OrderedDict()

raw_files = []

def FwFiles(FirmwareImagePath, search_guid=''):
    """Collect firmware-files with their file-section-ids

    Input:
        1. bios file path
        2. search_guid
    Output: raw_files[]
    """

    if search_guid:
        search_guid = search_guid.lower()

    def traverse_MultiVolumeContainer(self):
        for volume in self.volumes:
            volume.traverse()

    def traverse_MultiObject(self):
        for obj in self.objects:
            obj.traverse()

    def traverse_FirmwareVolume(self):
        if self.valid_header and len(self.data):
            for _ffs in self.firmware_filesystems:
                _ffs.traverse()

    def traverse_FirmwareFileSystem(self):
        for firmware_file in self.files:
            firmware_file.traverse()

    def traverse_FirmwareFile(self):
        global raw_files
        if search_guid:
            if self.guid_label == search_guid:
                raw_files += [(self.guid_label, self.data)]
        else:
            raw_files += [(self.guid_label, self.data)]
        for section in self.sections:
            section.traverse()

    def traverse_FirmwareFileSystemSection(self):
        if self.parsed_object:
            self.parsed_object.traverse()

    def traverse_GuidDefinedSection(self):
        for section in self.subsections:
            section.traverse()

    def traverse_RawObject(self):
        return

    def traverse_FirmwareCapsule(self):
        if self.valid_header and len(self.data):
            if self.capsule_body:
                self.capsule_body.traverse()

    def traverse_CompressedSection(self):
        for object in self.subsections:
            pass #object.traverse()

    def traverse_FlashDescriptor(self):
        for region in self.regions:
            region.traverse()

    def traverse_FlashRegion(self):
        for section in self.sections:
            section.traverse()

    def traverse_MeContainer(self):
        for partition in self.partitions:
            partition.traverse()

    def traverse_MeModule(self):
        pass

    def traverse_MeVariableModule(self):
        pass

    def traverse_MeManifestHeader(self):
        for module in self.modules:
            module.traverse()
        for module in self.variable_modules:
            module.traverse()
        if self.huffman_llut is not None:
            self.huffman_llut.traverse()

    def traverse_PartitionEntry(self):
        if self.manifest:
            self.manifest.traverse()

    def traverse_CPDManifestHeader(self):
        for entry in self.modules:
            entry.traverse()

    def traverse_CPDEntry(self):
        pass

    def traverse_MeLLUT(self):
        pass

    # monkey-patching
    uefi_firmware.MultiVolumeContainer.traverse             = traverse_MultiVolumeContainer
    uefi_firmware.MultiObject.traverse                      = traverse_MultiObject
    uefi_firmware.uefi.FirmwareVolume.traverse              = traverse_FirmwareVolume
    uefi_firmware.uefi.FirmwareFileSystem.traverse          = traverse_FirmwareFileSystem
    uefi_firmware.uefi.FirmwareFile.traverse                = traverse_FirmwareFile
    uefi_firmware.uefi.FirmwareFileSystemSection.traverse   = traverse_FirmwareFileSystemSection
    uefi_firmware.uefi.GuidDefinedSection.traverse          = traverse_GuidDefinedSection
    uefi_firmware.uefi.RawObject.traverse                   = traverse_RawObject
    uefi_firmware.uefi.FirmwareCapsule.traverse             = traverse_FirmwareCapsule
    uefi_firmware.uefi.CompressedSection.traverse           = traverse_CompressedSection
    uefi_firmware.flash.FlashDescriptor.traverse            = traverse_FlashDescriptor
    uefi_firmware.flash.FlashRegion.traverse                = traverse_FlashRegion
    uefi_firmware.me.MeContainer.traverse                   = traverse_MeContainer
    uefi_firmware.me.MeManifestHeader.traverse              = traverse_MeManifestHeader
    uefi_firmware.me.MeModule.traverse                      = traverse_MeModule
    uefi_firmware.me.MeVariableModule.traverse              = traverse_MeVariableModule
    uefi_firmware.me.MeLLUT.traverse                        = traverse_MeLLUT
    uefi_firmware.me.PartitionEntry.traverse                = traverse_PartitionEntry
    uefi_firmware.me.CPDManifestHeader.traverse             = traverse_CPDManifestHeader
    uefi_firmware.me.CPDEntry.traverse                      = traverse_CPDEntry

    with open(FirmwareImagePath, 'rb') as fh:
        file_content = fh.read()

    if not len(file_content):
        raise Exception("Invalid file size: %d" % len(file_content))

    parser = uefi_firmware.AutoParser(file_content)

    if parser.type() != 'unknown':
        firmware = parser.parse()
        firmware.traverse()

    return raw_files

def int32(data):
    ret = 0
    for i, v in enumerate(data[::-1]):
        ret *= 0x100
        ret += ord(v)
        if i == 3:
            break
    return ret

def mcu_header(blob):
    inc = 0
    count = 1
    while inc < len(blob):
        blobx = blob[inc:]
        print('#%d - %8.8X : rev%X (@%8X)' % (
            count, int32(blobx[3*4:4*4]), int32(blobx[1*4:2*4]), int32(blobx[2*4:3*4])
        ))
        inc += int32(blobx[8*4:9*4])
        count += 1


def mcu_header_v(blob):
    inc = 0
    count = 1
    hd = [
        'Header Version',
        'Patch ID',
        'DATE',
        'CPUID',
        'Checksum',
        'Loader Version',
        'Processor Flags',
        'Data Size',
        'Total Size'
    ]
    while inc < len(blob):
        print('\n[MCU %d]' % count)
        for i, n in enumerate(hd):
            print('%-16s : %8X' % (n, int32(blob[4*i:4*(i+1)])))
        count += 1
        inc += int32(blob[8*4: 9*4])


if __name__ == '__main__':
    bios = "CML_RVP_X64.bin"
    fwfdb0 = FwFiles(bios, '197DB236-F856-4924-90F8-CDF12FB875F3')
    for fn, data in fwfdb0:
        print('Write: %s' % fn)
        with open(fn, 'wb') as fout:
            fout.write(data)
        mcu_header(data)