1#!/usr/bin/env python3
2
3"""Script to create PLDM FW update package"""
4
5import argparse
6import binascii
7from datetime import datetime
8import json
9import os
10import struct
11import sys
12
13import math
14from bitarray import bitarray
15from bitarray.util import ba2int
16
17string_types = dict([
18    ("Unknown", 0),
19    ("ASCII", 1),
20    ("UTF8", 2),
21    ("UTF16", 3),
22    ("UTF16LE", 4),
23    ("UTF16BE", 5)])
24
25initial_descriptor_type_name_length = {
26    0x0000: ["PCI Vendor ID", 2],
27    0x0001: ["IANA Enterprise ID", 4],
28    0x0002: ["UUID", 16],
29    0x0003: ["PnP Vendor ID", 3],
30    0x0004: ["ACPI Vendor ID", 4]}
31
32descriptor_type_name_length = {
33    0x0000: ["PCI Vendor ID", 2],
34    0x0001: ["IANA Enterprise ID", 4],
35    0x0002: ["UUID", 16],
36    0x0003: ["PnP Vendor ID", 3],
37    0x0004: ["ACPI Vendor ID", 4],
38    0x0100: ["PCI Device ID", 2],
39    0x0101: ["PCI Subsystem Vendor ID", 2],
40    0x0102: ["PCI Subsystem ID", 2],
41    0x0103: ["PCI Revision ID", 1],
42    0x0104: ["PnP Product Identifier", 4],
43    0x0105: ["ACPI Product Identifier", 4]}
44
45
46def check_string_length(string):
47    """Check if the length of the string is not greater than 255."""
48    if len(string) > 255:
49        sys.exit("ERROR: Max permitted string length is 255")
50
51
52def write_pkg_release_date_time(pldm_fw_up_pkg, release_date_time):
53    '''
54    Write the timestamp into the package header. The timestamp is formatted as
55    series of 13 bytes defined in DSP0240 specification.
56
57        Parameters:
58            pldm_fw_up_pkg: PLDM FW update package
59            release_date_time: Package Release Date Time
60    '''
61    time = release_date_time.time()
62    date = release_date_time.date()
63    us_bytes = time.microsecond.to_bytes(3, byteorder='little')
64    pldm_fw_up_pkg.write(
65        struct.pack(
66            '<hBBBBBBBBHB',
67            0,
68            us_bytes[0],
69            us_bytes[1],
70            us_bytes[2],
71            time.second,
72            time.minute,
73            time.hour,
74            date.day,
75            date.month,
76            date.year,
77            0))
78
79
80def write_package_version_string(pldm_fw_up_pkg, metadata):
81    '''
82    Write PackageVersionStringType, PackageVersionStringLength and
83    PackageVersionString to the package header.
84
85        Parameters:
86            pldm_fw_up_pkg: PLDM FW update package
87            metadata: metadata about PLDM FW update package
88    '''
89    # Hardcoded string type to ASCII
90    string_type = string_types["ASCII"]
91    package_version_string = \
92        metadata["PackageHeaderInformation"]["PackageVersionString"]
93    check_string_length(package_version_string)
94    format_string = '<BB' + str(len(package_version_string)) + 's'
95    pldm_fw_up_pkg.write(
96        struct.pack(
97            format_string,
98            string_type,
99            len(package_version_string),
100            package_version_string.encode('ascii')))
101
102
103def write_component_bitmap_bit_length(pldm_fw_up_pkg, metadata):
104    '''
105    ComponentBitmapBitLength in the package header indicates the number of bits
106    that will be used represent the bitmap in the ApplicableComponents field
107    for a matching device. The value shall be a multiple of 8 and be large
108    enough to contain a bit for each component in the package. The number of
109    components in the JSON file is used to populate the bitmap length.
110
111        Parameters:
112            pldm_fw_up_pkg: PLDM FW update package
113            metadata: metadata about PLDM FW update package
114
115        Returns:
116            ComponentBitmapBitLength: number of bits that will be used
117            represent the bitmap in the ApplicableComponents field for a
118            matching device
119    '''
120    # The script supports upto 32 components now
121    max_components = 32
122    bitmap_multiple = 8
123
124    num_components = len(metadata["ComponentImageInformationArea"])
125    if num_components > max_components:
126        sys.exit("ERROR: only upto 32 components supported now")
127    component_bitmap_bit_length = bitmap_multiple * \
128        math.ceil(num_components/bitmap_multiple)
129    pldm_fw_up_pkg.write(struct.pack('<H', int(component_bitmap_bit_length)))
130    return component_bitmap_bit_length
131
132
133def write_pkg_header_info(pldm_fw_up_pkg, metadata):
134    '''
135    ComponentBitmapBitLength in the package header indicates the number of bits
136    that will be used represent the bitmap in the ApplicableComponents field
137    for a matching device. The value shall be a multiple of 8 and be large
138    enough to contain a bit for each component in the package. The number of
139    components in the JSON file is used to populate the bitmap length.
140
141        Parameters:
142            pldm_fw_up_pkg: PLDM FW update package
143            metadata: metadata about PLDM FW update package
144
145        Returns:
146            ComponentBitmapBitLength: number of bits that will be used
147            represent the bitmap in the ApplicableComponents field for a
148            matching device
149    '''
150    uuid = metadata["PackageHeaderInformation"]["PackageHeaderIdentifier"]
151    package_header_identifier = bytearray.fromhex(uuid)
152    pldm_fw_up_pkg.write(package_header_identifier)
153
154    package_header_format_revision = \
155        metadata["PackageHeaderInformation"]["PackageHeaderFormatVersion"]
156    # Size will be computed and updated subsequently
157    package_header_size = 0
158    pldm_fw_up_pkg.write(
159        struct.pack(
160            '<BH',
161            package_header_format_revision,
162            package_header_size))
163
164    try:
165        release_date_time = datetime.strptime(
166            metadata["PackageHeaderInformation"]["PackageReleaseDateTime"],
167            "%d/%m/%Y %H:%M:%S")
168        write_pkg_release_date_time(pldm_fw_up_pkg, release_date_time)
169    except KeyError:
170        write_pkg_release_date_time(pldm_fw_up_pkg, datetime.now())
171
172    component_bitmap_bit_length = write_component_bitmap_bit_length(
173        pldm_fw_up_pkg, metadata)
174    write_package_version_string(pldm_fw_up_pkg, metadata)
175    return component_bitmap_bit_length
176
177
178def get_applicable_components(device, components, component_bitmap_bit_length):
179    '''
180    This function figures out the components applicable for the device and sets
181    the ApplicableComponents bitfield accordingly.
182
183        Parameters:
184            device: device information
185            components: list of components in the package
186            component_bitmap_bit_length: length of the ComponentBitmapBitLength
187
188        Returns:
189            The ApplicableComponents bitfield
190    '''
191    applicable_components_list = device["ApplicableComponents"]
192    applicable_components = bitarray(component_bitmap_bit_length,
193                                     endian='little')
194    applicable_components.setall(0)
195    for component in components:
196        if component["ComponentIdentifier"] in applicable_components_list:
197            applicable_components[components.index(component)] = 1
198    return applicable_components
199
200
201def prepare_record_descriptors(descriptors):
202    '''
203    This function processes the Descriptors and prepares the RecordDescriptors
204    section of the the firmware device ID record.
205
206        Parameters:
207            descriptors: Descriptors entry
208
209        Returns:
210            RecordDescriptors, DescriptorCount
211    '''
212    record_descriptors = bytearray()
213    vendor_defined_desc_type = 65535
214    vendor_desc_title_str_type_len = 1
215    vendor_desc_title_str_len_len = 1
216    descriptor_count = 0
217
218    for descriptor in descriptors:
219
220        descriptor_type = descriptor["DescriptorType"]
221        if descriptor_count == 0:
222            if initial_descriptor_type_name_length.get(descriptor_type) \
223                    is None:
224                sys.exit("ERROR: Initial descriptor type not supported")
225        else:
226            if descriptor_type_name_length.get(descriptor_type) is None and \
227                    descriptor_type != vendor_defined_desc_type:
228                sys.exit("ERROR: Descriptor type not supported")
229
230        if descriptor_type == vendor_defined_desc_type:
231            vendor_desc_title_str = \
232                descriptor["VendorDefinedDescriptorTitleString"]
233            vendor_desc_data = descriptor["VendorDefinedDescriptorData"]
234            check_string_length(vendor_desc_title_str)
235            vendor_desc_title_str_type = string_types["ASCII"]
236            descriptor_length = vendor_desc_title_str_type_len + \
237                vendor_desc_title_str_len_len + len(vendor_desc_title_str) + \
238                len(bytearray.fromhex(vendor_desc_data))
239            format_string = '<HHBB' + str(len(vendor_desc_title_str)) + 's'
240            record_descriptors.extend(struct.pack(
241                format_string,
242                descriptor_type,
243                descriptor_length,
244                vendor_desc_title_str_type,
245                len(vendor_desc_title_str),
246                vendor_desc_title_str.encode('ascii')))
247            record_descriptors.extend(bytearray.fromhex(vendor_desc_data))
248            descriptor_count += 1
249        else:
250            descriptor_type = descriptor["DescriptorType"]
251            descriptor_data = descriptor["DescriptorData"]
252            descriptor_length = len(bytearray.fromhex(descriptor_data))
253            if descriptor_length != \
254                    descriptor_type_name_length.get(descriptor_type)[1]:
255                err_string = "ERROR: Descriptor type - " + \
256                    descriptor_type_name_length.get(descriptor_type)[0] + \
257                    " length is incorrect"
258                sys.exit(err_string)
259            format_string = '<HH'
260            record_descriptors.extend(struct.pack(
261                format_string,
262                descriptor_type,
263                descriptor_length))
264            record_descriptors.extend(bytearray.fromhex(descriptor_data))
265            descriptor_count += 1
266    return record_descriptors, descriptor_count
267
268
269def write_fw_device_identification_area(pldm_fw_up_pkg, metadata,
270                                        component_bitmap_bit_length):
271    '''
272    Write firmware device ID records into the PLDM package header
273
274    This function writes the DeviceIDRecordCount and the
275    FirmwareDeviceIDRecords into the firmware update package by processing the
276    metadata JSON. Currently there is no support for optional
277    FirmwareDevicePackageData.
278
279        Parameters:
280            pldm_fw_up_pkg: PLDM FW update package
281            metadata: metadata about PLDM FW update package
282            component_bitmap_bit_length: length of the ComponentBitmapBitLength
283    '''
284    # The spec limits the number of firmware device ID records to 255
285    max_device_id_record_count = 255
286    devices = metadata["FirmwareDeviceIdentificationArea"]
287    device_id_record_count = len(devices)
288    if device_id_record_count > max_device_id_record_count:
289        sys.exit(
290            "ERROR: there can be only upto 255 entries in the \
291                FirmwareDeviceIdentificationArea section")
292
293    # DeviceIDRecordCount
294    pldm_fw_up_pkg.write(struct.pack('<B', device_id_record_count))
295
296    for device in devices:
297        # RecordLength size
298        record_length = 2
299
300        # DescriptorCount
301        record_length += 1
302
303        # DeviceUpdateOptionFlags
304        device_update_option_flags = bitarray(32, endian='little')
305        device_update_option_flags.setall(0)
306        # Continue component updates after failure
307        supported_device_update_option_flags = [0]
308        for option in device["DeviceUpdateOptionFlags"]:
309            if option not in supported_device_update_option_flags:
310                sys.exit("ERROR: unsupported DeviceUpdateOptionFlag entry")
311            device_update_option_flags[option] = 1
312        record_length += 4
313
314        # ComponentImageSetVersionStringType supports only ASCII for now
315        component_image_set_version_string_type = string_types["ASCII"]
316        record_length += 1
317
318        # ComponentImageSetVersionStringLength
319        component_image_set_version_string = \
320            device["ComponentImageSetVersionString"]
321        check_string_length(component_image_set_version_string)
322        record_length += len(component_image_set_version_string)
323        record_length += 1
324
325        # Optional FirmwareDevicePackageData not supported now,
326        # FirmwareDevicePackageDataLength is set to 0x0000
327        fw_device_pkg_data_length = 0
328        record_length += 2
329
330        # ApplicableComponents
331        components = metadata["ComponentImageInformationArea"]
332        applicable_components = \
333            get_applicable_components(device,
334                                      components,
335                                      component_bitmap_bit_length)
336        applicable_components_bitfield_length = \
337            round(len(applicable_components)/8)
338        record_length += applicable_components_bitfield_length
339
340        # RecordDescriptors
341        descriptors = device["Descriptors"]
342        record_descriptors, descriptor_count = \
343            prepare_record_descriptors(descriptors)
344        record_length += len(record_descriptors)
345
346        format_string = '<HBIBBH' + \
347            str(applicable_components_bitfield_length) + 's' + \
348            str(len(component_image_set_version_string)) + 's'
349        pldm_fw_up_pkg.write(
350            struct.pack(
351                format_string,
352                record_length,
353                descriptor_count,
354                ba2int(device_update_option_flags),
355                component_image_set_version_string_type,
356                len(component_image_set_version_string),
357                fw_device_pkg_data_length,
358                applicable_components.tobytes(),
359                component_image_set_version_string.encode('ascii')))
360        pldm_fw_up_pkg.write(record_descriptors)
361
362
363def write_component_image_info_area(pldm_fw_up_pkg, metadata, image_files):
364    '''
365    Write component image information area into the PLDM package header
366
367    This function writes the ComponentImageCount and the
368    ComponentImageInformation into the firmware update package by processing
369    the metadata JSON. Currently there is no support for
370    ComponentComparisonStamp field and the component option use component
371    comparison stamp.
372
373    Parameters:
374        pldm_fw_up_pkg: PLDM FW update package
375        metadata: metadata about PLDM FW update package
376        image_files: component images
377    '''
378    components = metadata["ComponentImageInformationArea"]
379    # ComponentImageCount
380    pldm_fw_up_pkg.write(struct.pack('<H', len(components)))
381    component_location_offsets = []
382    # ComponentLocationOffset position in individual component image
383    # information
384    component_location_offset_pos = 12
385
386    for component in components:
387        # Record the location of the ComponentLocationOffset to be updated
388        # after appending images to the firmware update package
389        component_location_offsets.append(pldm_fw_up_pkg.tell() +
390                                          component_location_offset_pos)
391
392        # ComponentClassification
393        component_classification = component["ComponentClassification"]
394        if component_classification < 0 or component_classification > 0xFFFF:
395            sys.exit(
396                "ERROR: ComponentClassification should be [0x0000 - 0xFFFF]")
397
398        # ComponentIdentifier
399        component_identifier = component["ComponentIdentifier"]
400        if component_identifier < 0 or component_identifier > 0xFFFF:
401            sys.exit(
402                "ERROR: ComponentIdentifier should be [0x0000 - 0xFFFF]")
403
404        # ComponentComparisonStamp not supported
405        component_comparison_stamp = 0xFFFFFFFF
406
407        # ComponentOptions
408        component_options = bitarray(16, endian='little')
409        component_options.setall(0)
410        supported_component_options = [0]
411        for option in component["ComponentOptions"]:
412            if option not in supported_component_options:
413                sys.exit(
414                    "ERROR: unsupported ComponentOption in\
415                    ComponentImageInformationArea section")
416            component_options[option] = 1
417
418        # RequestedComponentActivationMethod
419        requested_component_activation_method = bitarray(16, endian='little')
420        requested_component_activation_method.setall(0)
421        supported_requested_component_activation_method = [0, 1, 2, 3, 4, 5]
422        for option in component["RequestedComponentActivationMethod"]:
423            if option not in supported_requested_component_activation_method:
424                sys.exit(
425                    "ERROR: unsupported RequestedComponent\
426                        ActivationMethod entry")
427            requested_component_activation_method[option] = 1
428
429        # ComponentLocationOffset
430        component_location_offset = 0
431        # ComponentSize
432        component_size = 0
433        # ComponentVersionStringType
434        component_version_string_type = string_types["ASCII"]
435        # ComponentVersionStringlength
436        # ComponentVersionString
437        component_version_string = component["ComponentVersionString"]
438        check_string_length(component_version_string)
439
440        format_string = '<HHIHHIIBB' + str(len(component_version_string)) + 's'
441        pldm_fw_up_pkg.write(
442            struct.pack(
443                format_string,
444                component_classification,
445                component_identifier,
446                component_comparison_stamp,
447                ba2int(component_options),
448                ba2int(requested_component_activation_method),
449                component_location_offset,
450                component_size,
451                component_version_string_type,
452                len(component_version_string),
453                component_version_string.encode('ascii')))
454
455    index = 0
456    pkg_header_checksum_size = 4
457    start_offset = pldm_fw_up_pkg.tell() + pkg_header_checksum_size
458    # Update ComponentLocationOffset and ComponentSize for all the components
459    for offset in component_location_offsets:
460        file_size = os.stat(image_files[index]).st_size
461        pldm_fw_up_pkg.seek(offset)
462        pldm_fw_up_pkg.write(
463            struct.pack(
464                '<II', start_offset, file_size))
465        start_offset += file_size
466        index += 1
467    pldm_fw_up_pkg.seek(0, os.SEEK_END)
468
469
470def write_pkg_header_checksum(pldm_fw_up_pkg):
471    '''
472    Write PackageHeaderChecksum into the PLDM package header.
473
474        Parameters:
475            pldm_fw_up_pkg: PLDM FW update package
476    '''
477    pldm_fw_up_pkg.seek(0)
478    package_header_checksum = binascii.crc32(pldm_fw_up_pkg.read())
479    pldm_fw_up_pkg.seek(0, os.SEEK_END)
480    pldm_fw_up_pkg.write(struct.pack('<I', package_header_checksum))
481
482
483def update_pkg_header_size(pldm_fw_up_pkg):
484    '''
485    Update PackageHeader in the PLDM package header. The package header size
486    which is the count of all bytes in the PLDM package header structure is
487    calculated once the package header contents is complete.
488
489        Parameters:
490            pldm_fw_up_pkg: PLDM FW update package
491    '''
492    pkg_header_checksum_size = 4
493    file_size = pldm_fw_up_pkg.tell() + pkg_header_checksum_size
494    pkg_header_size_offset = 17
495    # Seek past PackageHeaderIdentifier and PackageHeaderFormatRevision
496    pldm_fw_up_pkg.seek(pkg_header_size_offset)
497    pldm_fw_up_pkg.write(struct.pack('<H', file_size))
498    pldm_fw_up_pkg.seek(0, os.SEEK_END)
499
500
501def append_component_images(pldm_fw_up_pkg, image_files):
502    '''
503    Append the component images to the firmware update package.
504
505        Parameters:
506            pldm_fw_up_pkg: PLDM FW update package
507            image_files: component images
508    '''
509    for image in image_files:
510        with open(image, 'rb') as file:
511            for line in file:
512                pldm_fw_up_pkg.write(line)
513
514
515def main():
516    """Create PLDM FW update (DSP0267) package based on a JSON metadata file"""
517    parser = argparse.ArgumentParser()
518    parser.add_argument("pldmfwuppkgname",
519                        help="Name of the PLDM FW update package")
520    parser.add_argument("metadatafile", help="Path of metadata JSON file")
521    parser.add_argument(
522        "images", nargs='+',
523        help="One or more firmware image paths, in the same order as\
524            ComponentImageInformationArea entries")
525
526    args = parser.parse_args()
527    image_files = args.images
528    with open(args.metadatafile) as file:
529        try:
530            metadata = json.load(file)
531        except ValueError:
532            sys.exit("ERROR: Invalid metadata JSON file")
533
534    # Validate the number of component images
535    if len(image_files) != len(metadata["ComponentImageInformationArea"]):
536        sys.exit("ERROR: number of images passed != number of entries \
537            in ComponentImageInformationArea")
538
539    try:
540        with open(args.pldmfwuppkgname, 'w+b') as pldm_fw_up_pkg:
541            component_bitmap_bit_length = write_pkg_header_info(pldm_fw_up_pkg,
542                                                                metadata)
543            write_fw_device_identification_area(pldm_fw_up_pkg,
544                                                metadata,
545                                                component_bitmap_bit_length)
546            write_component_image_info_area(pldm_fw_up_pkg, metadata,
547                                            image_files)
548            update_pkg_header_size(pldm_fw_up_pkg)
549            write_pkg_header_checksum(pldm_fw_up_pkg)
550            append_component_images(pldm_fw_up_pkg, image_files)
551            pldm_fw_up_pkg.close()
552    except BaseException:
553        pldm_fw_up_pkg.close()
554        os.remove(args.pldmfwuppkgname)
555        raise
556
557
558if __name__ == "__main__":
559    main()
560