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