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