diff --git a/edk2toolext/capsule/capsule_helper.py b/edk2toolext/capsule/capsule_helper.py index 617a5609..c9928272 100644 --- a/edk2toolext/capsule/capsule_helper.py +++ b/edk2toolext/capsule/capsule_helper.py @@ -13,8 +13,13 @@ import uuid import os import struct +import datetime -from edk2toollib.windows.capsule import inf_generator, cat_generator +from typing import List +from dataclasses import dataclass +from dataclasses import field + +from edk2toollib.windows.capsule import inf_generator2, cat_generator from edk2toollib.uefi.uefi_capsule_header import UefiCapsuleHeaderClass from edk2toollib.uefi.fmp_capsule_header import FmpCapsuleHeaderClass, FmpCapsuleImageHeaderClass from edk2toollib.uefi.fmp_auth_header import FmpAuthHeaderClass @@ -24,6 +29,57 @@ PKCS7_SIGNED_DATA_OID = '1.2.840.113549.1.7.2' +@dataclass +class CapsulePayload: + '''Stores information about a specific capsule payload. + + CapsulePayload instances have the following attributes: + payload - an instance of UefiCapsuleHeaderClass that represents the payload data. + payload_filename - the payload filename as a string + esrt_guid - the payload ESRT guid as a uuid.UUID instance. + version - the 32-bit ESRT version for the payload. + firmware_description - the firmware payload description. + tag - a string uniquely identifying the payload. optional, if not present, will be auto-generated. + rollback - indicates whether this is a rollback payload. optional, defaults to false. + integrity_data - integrity data for this payload. optional. + integrity_filename - integrity filename. optional if integrity_data is None, required otherwise. + ''' + payload: UefiCapsuleHeaderClass + payload_filename: str + esrt_guid: uuid.UUID + version: int + firmware_description: str + tag: str = None + rollback: bool = False + integrity_data: bytes = None + integrity_filename: str = None + + +@dataclass +class Capsule: + '''Stores information about a capsule (potentially with multiple payloads) + + Capsule instances have the following attributes: + version_string - the version of the entire capsule driver package as a string (e.g. 1.0.0.1) + name - the name of the capsule package + provider_name - the name of the capsule provider + arch - the architecture targeted by the capsule + os - the OS targeted by the capsule. + manufacturer_name - name of the capsule manufacturer. optional, defaults to provider_name if None. + date - a datetime.date object indicating when the capsule was built. optional, defaults to + datetime.date.today(). + payloads - a list of capsule payloads. optional, defaults to empty list + ''' + version_string: str + name: str + provider_name: str + arch: str = None + os: str = None + manufacturer_name: str = None + date: datetime.date = datetime.date.today() + payloads: List[CapsulePayload] = field(default_factory=list) + + def get_capsule_file_name(capsule_options: dict) -> str: '''from the shared capsule_options dictionary, returns the formatted capsule file name''' return f"{capsule_options['fw_name']}_{capsule_options['fw_version_string']}.bin" @@ -131,6 +187,30 @@ def save_capsule(uefi_capsule_header: UefiCapsuleHeaderClass, capsule_options: d return capsule_file_path +def save_multinode_capsule(capsule: Capsule, save_path: str) -> str: + ''' + takes in a Capsule object and a filesystem directory path and generates the capsule files at that path. + + capsule - a Capsule object containing the capsule details. + save_path - directory path to save the capsule contents into + + returns save_path + ''' + os.makedirs(save_path, exist_ok=True) + for capsule_payload in capsule.payloads: + payload_file_path = os.path.join(save_path, capsule_payload.payload_filename) + with open(payload_file_path, 'wb') as payload_file: + payload_file.write(capsule_payload.payload.Encode()) + + if capsule_payload.integrity_data is not None: + if (capsule_payload.integrity_filename is None): + raise ValueError("Integrity data specified, but no integrity filename specified.") + integrity_file_path = os.path.join(save_path, capsule_payload.integrity_filename) + with open(integrity_file_path, 'wb') as integrity_file: + integrity_file.write(capsule_payload.integrity_data) + return save_path + + def create_inf_file(capsule_options: dict, save_path: str) -> str: ''' takes in a dictionary of capsule_options and creates the Windows INF file for the UEFI capsule according @@ -138,6 +218,7 @@ def create_inf_file(capsule_options: dict, save_path: str) -> str: will save the final file to the save_path with a name determined from the capsule_options ''' + # Expand the version string prior to creating INF file. capsule_options['fw_version_string'] = get_normalized_version_string(capsule_options['fw_version_string']) @@ -146,23 +227,77 @@ def create_inf_file(capsule_options: dict, save_path: str) -> str: capsule_options['arch'] = capsule_options.get('arch', get_default_arch()) capsule_options['mfg_name'] = capsule_options.get('mfg_name', capsule_options['provider_name']) - # Create the INF. - infgenerator = inf_generator.InfGenerator( + inf_file = inf_generator2.InfFile( capsule_options['fw_name'], + capsule_options['fw_version_string'], + datetime.date.today().strftime("%m/%d/%Y"), capsule_options['provider_name'], - capsule_options['esrt_guid'], - capsule_options['arch'], + capsule_options['mfg_name'], + capsule_options['arch'] + ) + + inf_file.AddFirmware( + "Firmware", capsule_options['fw_description'], - capsule_options['fw_version_string'], - capsule_options['fw_version'] + capsule_options['esrt_guid'], + capsule_options['fw_version'], + get_capsule_file_name(capsule_options), + Rollback=capsule_options['is_rollback'], + IntegrityFile=capsule_options.get('fw_integrity_file', None) ) - infgenerator.Manufacturer = capsule_options['mfg_name'] - if 'fw_integrity_file' in capsule_options: - infgenerator.IntegrityFilename = os.path.basename(capsule_options['fw_integrity_file']) + inf_file_path = os.path.join(save_path, f"{capsule_options['fw_name']}.inf") - ret = infgenerator.MakeInf(inf_file_path, get_capsule_file_name(capsule_options), capsule_options['is_rollback']) - if(ret != 0): - raise RuntimeError("MakeInf Failed with errorcode %d!" % ret) + with open(inf_file_path, "w") as fp: + fp.write(str(inf_file)) + + return inf_file_path + + +def create_multinode_inf_file(capsule: Capsule, save_path: str) -> str: + ''' + Takes in a capsule object containing payload information and creates the Windows INF file in save_path + + capsule - capsule object containing payload information + save_path - path to directory where inf file will be created. + + returns the name of the created inf file. + ''' + # Expand the version string prior to creating INF file. + capsule.version_string = get_normalized_version_string(capsule.version_string) + + # set defaults for non-specified fields + if (capsule.arch is None): + capsule.arch = get_default_arch() + if (capsule.manufacturer_name is None): + capsule.manufacturer_name = capsule.provider_name + + inf_file = inf_generator2.InfFile( + capsule.name, + capsule.version_string, + capsule.date.strftime("%m/%d/%Y"), + capsule.provider_name, + capsule.manufacturer_name, + capsule.arch + ) + + idx = 0 + for payload in capsule.payloads: + if payload.tag is None: + payload.tag = f"Firmware{idx}" + idx += 1 + inf_file.AddFirmware( + payload.tag, + payload.firmware_description, + str(payload.esrt_guid), + str(payload.version), + payload.payload_filename, + Rollback=payload.rollback, + IntegrityFile=payload.integrity_filename + ) + + inf_file_path = os.path.join(save_path, f"{capsule.name}.inf") + with open(inf_file_path, "w") as fp: + fp.write(str(inf_file)) return inf_file_path diff --git a/edk2toolext/tests/capsule/capsule_helper_test.py b/edk2toolext/tests/capsule/capsule_helper_test.py index ea89b83a..c6057a63 100644 --- a/edk2toolext/tests/capsule/capsule_helper_test.py +++ b/edk2toolext/tests/capsule/capsule_helper_test.py @@ -163,5 +163,101 @@ def test_should_be_able_to_generate_cat(self): self.assertTrue(os.path.isfile(cat_file_path)) +class MultiNodeFileGenerationTest(unittest.TestCase): + + @staticmethod + def buildPayload(esrt): + fmp_capsule_image_header = FmpCapsuleImageHeaderClass() + fmp_capsule_image_header.UpdateImageTypeId = uuid.UUID(esrt) + fmp_capsule_image_header.UpdateImageIndex = 1 + + fmp_capsule_header = FmpCapsuleHeaderClass() + fmp_capsule_header.AddFmpCapsuleImageHeader(fmp_capsule_image_header) + + uefi_capsule_header = UefiCapsuleHeaderClass() + uefi_capsule_header.FmpCapsuleHeader = fmp_capsule_header + uefi_capsule_header.PersistAcrossReset = True + uefi_capsule_header.InitiateReset = True + + return uefi_capsule_header + + @classmethod + def setUpClass(cls): + cls.temp_dir = tempfile.mkdtemp() + cls.temp_output_dir = tempfile.mkdtemp() + + cls.capsule = capsule_helper.Capsule( + version_string="1.2.3", + name="TEST_FW", + provider_name="Tester", + ) + + cls.capsule.payloads.append( + capsule_helper.CapsulePayload( + cls.buildPayload("ea5c13fe-cac9-4fd7-ac30-37709bd668f2"), + "test1.bin", + uuid.UUID("ea5c13fe-cac9-4fd7-ac30-37709bd668f2"), + 0xDEADBEEF, + "TEST FW" + ) + ) + + cls.capsule.payloads.append( + capsule_helper.CapsulePayload( + cls.buildPayload("43e67b4e-b2f1-4891-9ff2-a6acd9c74cbd"), + "test2.bin", + uuid.UUID("43e67b4e-b2f1-4891-9ff2-a6acd9c74cbd"), + 0xDEADBEEF, + "TEST FW" + ) + ) + + def test_should_be_able_to_save_a_multi_node_capsule(self): + + capsule_file_path = capsule_helper.save_multinode_capsule(self.capsule, self.temp_output_dir) + + # make sure all the files we expect got created + for payload in self.capsule.payloads: + payload_file = os.path.join(capsule_file_path, payload.payload_filename) + self.assertTrue(os.path.isfile(payload_file)) + with open(payload_file, 'rb') as fh: + capsule_bytes = fh.read() + self.assertIn(payload.esrt_guid.bytes_le, capsule_bytes) + + def test_should_be_able_to_save_a_multi_node_capsule_with_integrity(self): + + self.capsule.payloads[0].integrity_data = uuid.UUID("ea5c13fe-cac9-4fd7-ac30-37709bd668f2").bytes + self.capsule.payloads[0].integrity_filename = "integrity1.bin" + + self.capsule.payloads[1].integrity_data = uuid.UUID("43e67b4e-b2f1-4891-9ff2-a6acd9c74cbd").bytes + self.capsule.payloads[1].integrity_filename = "integrity2.bin" + + capsule_file_path = capsule_helper.save_multinode_capsule(self.capsule, self.temp_output_dir) + + for payload in self.capsule.payloads: + payload_file = os.path.join(capsule_file_path, payload.payload_filename) + self.assertTrue(os.path.isfile(payload_file)) + with open(payload_file, 'rb') as fh: + capsule_bytes = fh.read() + self.assertIn(payload.esrt_guid.bytes_le, capsule_bytes) + + integrityFile = os.path.join(capsule_file_path, payload.integrity_filename) + self.assertTrue(os.path.isfile(integrityFile)) + with open(integrityFile, 'rb') as fh: + integrity_bytes = fh.read() + self.assertIn(payload.integrity_data, integrity_bytes) + + self.capsule.payloads[0].integrity_data = None + self.capsule.payloads[0].integrity_filename = None + + self.capsule.payloads[1].integrity_data = None + self.capsule.payloads[1].integrity_filename = None + + def test_should_be_able_to_generate_multi_node_inf_file(self): + + inf_file_path = capsule_helper.create_multinode_inf_file(self.capsule, self.temp_output_dir) + self.assertTrue(os.path.isfile(inf_file_path)) + + if __name__ == '__main__': unittest.main() diff --git a/setup.py b/setup.py index 821875b5..2c5e3e87 100644 --- a/setup.py +++ b/setup.py @@ -78,7 +78,7 @@ def run(self): }, install_requires=[ 'pyyaml>=5.3.1', - 'edk2-pytool-library>=0.10.15', + 'edk2-pytool-library>=0.11.0', 'pefile>=2019.4.18' ], extras_require={