Add support for multi-node capsules to capsule_helper script (#269)

Add routines to capsule_helper script to support generation of mult-node capsules.
Added save_multinode_capsule which generates a capsule package for a multi-node capsule, and create_multinode_inf_file which generates a multi-node inf.

Updates edk2-pytool-library dependency.

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,13 +187,38 @@ 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
     to the provided options
 
     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
 

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()

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={