Update Driver to existing Binary of specified GUID

docs/user_guide/uefi_binary_parsing.md

@@ -5,6 +5,7 @@ The parser is able to generate the json representation from BIOS or IFWI image.
 Key features:
 -   JSON representation, lightweight database with keys and values with ease of readability
 -   Works with both SUT and offline image
+-   Replace new driver in a bios image.
 
 Working with SUT:
 
@@ -44,3 +45,43 @@ if uefi_parser.guid_to_store:
     # Store guid stored result to json file
     uefi_parser.write_result_to_file(user_guid_out_file, output_dict=uefi_parser.stored_guids)
 ```
+
+
+
+Replace the Driver Ffs in Binary with below commands:
+```python
+from xmlcli.common import bios_fw_parser
+
+bios_image = "absolute-path/to/bios-image.rom"
+driver_image = "absolute-path/to/driver_image.ffs"
+output_image = "absolute-path/to/replaced_bios_image.rom"
+
+uefi_parser = bios_fw_parser.UefiParser(bin_file=bios_image,  # binary file to parse
+                                parsing_level=0,  # parsing level to manage number of parsing features
+                                base_address=0,  # (optional) provide base address of bios FV region to start the parsing (default 0x0)
+                                guid_to_store=[]  # if provided the guid for parsing then parser will look for every GUID in the bios image
+                                )
+
+newffs_parser = bios_fw_parser.UefiParser(bin_file=driver_image,  # binary file to parse
+                                parsing_level=0,  # parsing level to manage number of parsing features
+                                base_address=0,  # (optional) provide base address of bios FV region to start the parsing (default 0x0)
+                                guid_to_store=[]  # if provided the guid for parsing then parser will look for every GUID in the bios image
+                                )
+
+# parse bios image into a binary_tree
+bios_output_dict = uefi_parser.parse_binary()
+
+# parse driver ffs image into a binary tree node
+ffs_output_dict = newffs_parser.parse_binary()
+# get the target ffs guid through ffs file, extract the target tree node
+TargetFfsGuid = newffs_parser.binary_tree.Position.ChildNodeList[0].Data.Name
+newffsnode = newffs_parser.binary_tree.Position.ChildNodeList[0]
+
+# replace the target ffs with new one
+uefi_parser.find_ffs_node(TargetFfsGuid)
+uefi_parser.ReplaceFfs(newffsnode, uefi_parser.TargetFfsList[0])
+uefi_parser.binary_tree.WholeTreeData = b''
+uefi_parser.Encapsulate_binary(uefi_parser.binary_tree)
+# dump the bios image with replaced ffs info
+uefi_parser.dump_binary(replaced_image)
+```

src/xmlcli/common/compress.py

@@ -8,9 +8,9 @@
 from datetime import datetime
 
 # Custom imports
-from . import utils
-from . import logger
-from . import configurations
+import utils
+import logger
+import configurations
 
 __author__ = "Gahan Saraiya"
 

src/xmlcli/common/logger.py

@@ -9,7 +9,7 @@
 from datetime import datetime
 
 # Custom imports
-from .configurations import XMLCLI_CONFIG, XMLCLI_DIR, OUT_DIR, PLATFORM, PY_VERSION, ENCODING, STATUS_CODE_RECORD_FILE
+from configurations import XMLCLI_CONFIG, XMLCLI_DIR, OUT_DIR, PLATFORM, PY_VERSION, ENCODING, STATUS_CODE_RECORD_FILE
 
 ###############################################################################
 # START: LOG Settings #########################################################

src/xmlcli/common/structure.py

@@ -12,7 +12,6 @@
 except ModuleNotFoundError as e:
   DescriptorRegion = None
 
-
 __version__ = "0.0.1"
 __author__ = "Gahan Saraiya"
 
@@ -53,6 +52,16 @@ def is_bios(bin_file):
   bios_header = EfiFirmwareVolumeHeader.read_from(bin_file)
   return FV_SIGNATURE == bios_header.Signature
 
+def struct2stream(target_struct):
+    length = ctypes.sizeof(target_struct)
+    p = ctypes.cast(ctypes.pointer(target_struct), ctypes.POINTER(ctypes.c_char * length))
+    return p.contents.raw
+
+def get_pad_size(size: int, alignment: int):
+    if size % alignment == 0:
+        return 0
+    pad_Size = alignment - size % alignment
+    return pad_Size
 
 def read_structure(method, base_structure, buffer, buffer_pointer):
   """Read Valid structure
@@ -170,6 +179,7 @@ def process_efi_cert_type_rsa2048_sha256_guid(buffer, buffer_pointer, section):
 # i.e. EFI_FVB2_ALIGNMENT_128 will be 0x00070000
 EFI_FVB2_ALIGNMENT = lambda size=None: "0x001F0000" if not size else f"0x{size:0>4x}0000"
 EFI_FVB2_WEAK_ALIGNMENT = 0x80000000
+EFI_FVB2_ERASE_POLARITY = 0x00000800
 
 
 class EfiTime(utils.StructureHelper):  # 16 bytes
@@ -526,6 +536,33 @@ def dump_dict(self):
     result.pop("Checksum")
     return result
 
+def Refine_EfiFirmwareVolumeHeader(nums):
+  class EfiFirmwareVolumeHeader(utils.StructureHelper):
+    # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h
+    _fields_ = [
+      ("ZeroVector", ctypes.ARRAY(ctypes.c_uint8, 16)),
+      ("FileSystemGuid", utils.Guid),  # EFI_GUID
+      ("FvLength", ctypes.c_uint64),
+      ("Signature", ctypes.ARRAY(ctypes.c_char, 4)),  # actually it's a signature of 4 characters... (UINT32 Signature)
+      ("Attributes", EFI_FVB_ATTRIBUTES_2),  # UINT32 EFI_FVB_ATTRIBUTES_2
+      ("HeaderLength", ctypes.c_uint16),
+      ("Checksum", ctypes.c_uint16),
+      ("ExtHeaderOffset", ctypes.c_uint16),
+      ("Reserved", ctypes.ARRAY(ctypes.c_uint8, 1)),
+      ("Revision", ctypes.c_uint8),
+      ("BlockMap", ctypes.ARRAY(EfiFvBlockMapEntry, nums))  # EFI_FV_BLOCK_MAP_ENTRY
+    ]
+
+    @property
+    def get_guid(self):
+      return self.FileSystemGuid
+
+    def dump_dict(self):
+      result = super(EfiFirmwareVolumeHeader, self).dump_dict()
+      result.pop("BlockMap")
+      result.pop("Checksum")
+      return result
+  return EfiFirmwareVolumeHeader
 
 class EfiFirmwareVolumeExtHeader(utils.StructureHelper):
   # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h
@@ -554,6 +591,15 @@ class EfiFirmwareVolumeExtEntryOemType(utils.StructureHelper):
     ("TypeMask", ctypes.c_uint32)
   ]
 
+def Refine_FV_EXT_ENTRY_OEM_TYPE_Header(nums: int):
+  class EfiFirmwareVolumeExtEntryOemType(utils.StructureHelper):
+    # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h
+    _fields_ = [
+      ("Hdr", EfiFirmwareVolumeExtEntry),
+      ("TypeMask", ctypes.c_uint32),
+      ('Types',    ctypes.ARRAY(utils.Guid, nums))
+      ]
+  return EfiFirmwareVolumeExtEntryOemType()
 
 class EfiFirmwareVolumeExtEntryGuidType(utils.StructureHelper):
   # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h
@@ -566,6 +612,20 @@ class EfiFirmwareVolumeExtEntryGuidType(utils.StructureHelper):
   def get_guid(self):
     return self.FormatType
 
+def Refine_FV_EXT_ENTRY_GUID_TYPE_Header(nums: int):
+  class EfiFirmwareVolumeExtEntryGuidType(utils.StructureHelper):
+    # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h
+    _fields_ = [
+      ("Hdr", EfiFirmwareVolumeExtEntry),
+      ("FormatType", utils.Guid),
+      ('Data',       ctypes.ARRAY(ctypes.c_uint8, nums))
+    ]
+    @property
+
+    def get_guid(self):
+      return self.FormatType
+
+  return EfiFirmwareVolumeExtEntryGuidType()
 
 class EfiFirmwareVolumeExtEntryUsedSizeType(utils.StructureHelper):
   # source: Edk2/BaseTools/Source/C/Include/Common/PiFirmwareVolume.h

src/xmlcli/common/tree_and_node.py

@@ -0,0 +1,131 @@
+# -*- coding: utf-8 -*-
+
+# Built-in imports
+from collections import namedtuple
+
+# custom imports
+import structure
+from logger import log
+
+
+__version__ = "0.0.1"
+__author__ = "Christine Chen & Yuting2 Yang"
+
+class TreeStructure:
+  def __init__(self) -> None:
+    self.ParentNode = None
+    self.ChildNodeList = []
+
+class TreeNode:
+  def __init__(self, Key) -> None:
+    self.Key = Key
+    self.Type = None
+    self.Data = b''
+    self.WholeTreeData = b''
+    self.Position = TreeStructure()
+
+  def HasChild(self):
+    if self.Position.ChildNodeList != []:
+      return True
+    else:
+      return False
+
+  def InsertChildNode(self, NewChild, order=None):
+    if order is not None:
+      self.Position.ChildNodeList.insert(order, NewChild)
+      NewChild.Position.ParentNode = self
+    else:
+      self.Position.ChildNodeList.append(NewChild)
+      NewChild.Position.ParentNode = self
+
+  def RemoveChildNode(self, RemoveChild):
+    self.Position.ChildNodeList.remove(RemoveChild)
+
+  def FindNode(self, key, findlist) -> None:
+    if self.Type == 'FFS' and self.Data.Name.guid==key.guid:
+      findlist.append(self)
+    for item in self.Position.ChildNodeList:
+      item.FindNode(key, findlist)
+
+class FvNode:
+  def __init__(self):
+    self.Name = ''
+    self.Header = None
+    self.ExtHeader = None
+    self.Data = b''
+    self.Info = None
+    self.IsValid = None
+    self.Free_Space = 0
+
+  def InitExtEntry(self, buffer, buffer_pointer):
+    self.ExtEntryOffset = self.Header.ExtHeaderOffset + 20
+    buffer.seek(buffer_pointer+self.ExtEntryOffset)
+    if self.ExtHeader.ExtHeaderSize != 20:
+      self.ExtEntryExist = 1
+      self.ExtEntry = structure.EfiFirmwareVolumeExtEntry.read_from(buffer)
+      self.ExtTypeExist = 1
+      if self.ExtEntry.ExtEntryType == 0x01:
+        nums = (self.ExtEntry.ExtEntrySize - 8) // 16
+        self.ExtEntry = structure.Refine_FV_EXT_ENTRY_OEM_TYPE_Header(nums).read_from(buffer)
+      elif self.ExtEntry.ExtEntryType == 0x02:
+        nums = self.ExtEntry.ExtEntrySize - 20
+        self.ExtEntry = structure.Refine_FV_EXT_ENTRY_GUID_TYPE_Header(nums).read_from(buffer)
+      elif self.ExtEntry.ExtEntryType == 0x03:
+        self.ExtEntry = structure.EfiFirmwareVolumeExtEntryUsedSizeType.read_from(buffer)
+      else:
+        self.ExtTypeExist = 0
+    else:
+      self.ExtEntryExist = 0
+
+  def ModCheckSum(self):
+    # Fv Header Sums to 0.
+    Header = structure.struct2stream(self.Header)[::-1]
+    Size = self.Header.HeaderLength // 2
+    Sum = 0
+    for i in range(Size):
+      Sum += int(Header[i*2: i*2 + 2].hex(), 16)
+    if Sum & 0xffff:
+      self.Header.Checksum = 0x10000 - (Sum - self.Header.Checksum) % 0x10000
+
+  def ModFvExt(self):
+    # If used space changes and self.ExtEntry.UsedSize exists, self.ExtEntry.UsedSize need to be changed.
+    if self.Header.ExtHeaderOffset and self.ExtEntryExist and self.ExtTypeExist and self.ExtEntry.Hdr.ExtEntryType == 0x03:
+      self.ExtEntry.UsedSize = self.Header.FvLength - self.Free_Space
+
+  def ModFvSize(self):
+    # If Fv Size changed, self.Header.FvLength and self.Header.BlockMap[i].NumBlocks need to be changed.
+    BlockMapNum = len(self.Header.BlockMap)
+    for i in range(BlockMapNum):
+      if self.Header.BlockMap[i].Length:
+        self.Header.BlockMap[i].NumBlocks = self.Header.FvLength // self.Header.BlockMap[i].Length
+
+  def ModExtHeaderData(self):
+    if self.Header.ExtHeaderOffset:
+      ExtHeaderData = structure.struct2stream(self.ExtHeader)
+      ExtHeaderDataOffset = self.Header.ExtHeaderOffset - self.Header.HeaderLength
+      self.Data = self.Data[:ExtHeaderDataOffset] + ExtHeaderData + self.Data[ExtHeaderDataOffset+20:]
+    if self.Header.ExtHeaderOffset and self.ExtEntryExist:
+      ExtHeaderEntryData = structure.struct2stream(self.ExtEntry)
+      ExtHeaderEntryDataOffset = self.Header.ExtHeaderOffset + 20 - self.HeaderLength
+      self.Data = self.Data[:ExtHeaderEntryDataOffset] + ExtHeaderEntryData + self.Data[ExtHeaderEntryDataOffset+len(ExtHeaderEntryData):]
+
+
+class FfsNode:
+  def __init__(self):
+    self.Name = ''
+    self.Header = None
+    self.Data = b''
+    self.PadData = b''
+    self.Info = None
+    self.FfsType = None
+
+  def ModCheckSum(self) -> None:
+    HeaderData = structure.struct2stream(self.Header)
+    HeaderSum = 0
+    for item in HeaderData:
+      HeaderSum += item
+    HeaderSum -= self.Header.State
+    HeaderSum -= self.Header.IntegrityCheck.Checksum.File
+    if HeaderSum & 0xff:
+      Header = self.Header.IntegrityCheck.Checksum.Header + 0x100 - HeaderSum % 0x100
+      self.Header.IntegrityCheck.Checksum.Header = Header % 0x100

src/xmlcli/common/utils.py

@@ -20,8 +20,8 @@
 from collections import OrderedDict
 
 # Custom imports
-from .logger import log
-from .configurations import XMLCLI_CONFIG, ENCODING, XMLCLI_DIR, OUT_DIR, PY3, STATUS_CODE_RECORD_FILE
+from logger import log
+from configurations import XMLCLI_CONFIG, ENCODING, XMLCLI_DIR, OUT_DIR, PY3, STATUS_CODE_RECORD_FILE
 
 try:
   from defusedxml import ElementTree as ET

tests/UefiParserTest.py

@@ -103,6 +103,9 @@
   def test_write_to_json(self):
     self.parse_image(self.bios_image)
 
+  def test_replace_driver(self):
+    self.replace_ffs(self.bios_image, self.new_driver_file, self.replaced_image_file)
+
   def parse_image(self, bios_image):
     self.log.info(f"{'=' * 50}\n>>>>>>>>> PROCESSING IMAGE: {bios_image} <<<<<<<<<\n{'=' * 50}")
     binary_file_name = os.path.splitext(bios_image)[0]  # get filename without extension
@@ -134,6 +137,36 @@
       # Validate whether content written in json or not
       self.assertGreater(os.path.getsize(user_guid_out_file), 0)
 
+  def replace_ffs(self, bios_image, driver_image, replaced_image):
+    self.log.info(f"{'=' * 50}\n>>>>>>>>> REPLACING DRIVER: {driver_image} <<<<<<<<<\n{'=' * 50}")
+    uefi_parser = bios_fw_parser.UefiParser(bin_file=bios_image,  # binary file to parse
+                                parsing_level=0,  # parsing level to manage number of parsing features
+                                base_address=0,  # (optional) provide base address of bios FV region to start the parsing (default 0x0)
+                                guid_to_store=[]  # if provided the guid for parsing then parser will look for every GUID in the bios image
+                                )
+
+    newffs_parser = bios_fw_parser.UefiParser(bin_file=driver_image,  # binary file to parse
+                                    parsing_level=0,  # parsing level to manage number of parsing features
+                                    base_address=0,  # (optional) provide base address of bios FV region to start the parsing (default 0x0)
+                                    guid_to_store=[]  # if provided the guid for parsing then parser will look for every GUID in the bios image
+                                    )
+
+    # parse bios image into a binary_tree
+    bios_output_dict = uefi_parser.parse_binary()
+
+    # parse driver ffs image into a binary tree node
+    ffs_output_dict = newffs_parser.parse_binary()
+    # get the target ffs guid through ffs file, extract the target tree node
+    TargetFfsGuid = newffs_parser.binary_tree.Position.ChildNodeList[0].Data.Name
+    newffsnode = newffs_parser.binary_tree.Position.ChildNodeList[0]
+
+    # replace the target ffs with new one
+    uefi_parser.find_ffs_node(TargetFfsGuid)
+    uefi_parser.ReplaceFfs(newffsnode, uefi_parser.TargetFfsList[0])
+    uefi_parser.binary_tree.WholeTreeData = b''
+    uefi_parser.Encapsulate_binary(uefi_parser.binary_tree)
+    uefi_parser.dump_binary(replaced_image)
+
   def test_compare_with_old(self):
     for bios_image in self.bios_roms:
       with open(bios_image, 'rb') as BiosBinFile:

tests/UnitTestHelper.py

@@ -33,6 +33,8 @@
 ACCESS_METHOD = TEST_SUITE_CONFIG.get("TEST_SETTINGS", "ACCESS_METHOD")
 RUN_OPTIONAL_TEST = TEST_SUITE_CONFIG.getboolean("TEST_SETTINGS", "RUN_OPTIONAL_TEST")
 BIOS_IMAGES_DIR = os.path.abspath(TEST_SUITE_CONFIG.get("TEST_SETTINGS", "BIOS_IMAGES_DIR"))
+NEW_DRIVER_FILE = os.path.abspath(TEST_SUITE_CONFIG.get("TEST_SETTINGS", "NEW_DRIVER_FILE"))
+REPLACED_IMAGE_FILE = os.path.abspath(TEST_SUITE_CONFIG.get("TEST_SETTINGS", "REPLACED_IMAGE_FILE"))
 LITE_FEATURE_TESTING = TEST_SUITE_CONFIG.getboolean("TEST_SETTINGS", "LITE_FEATURE_TESTING")
 
 LOG_TITLE = TEST_SUITE_CONFIG.get("LOG_SETTINGS", "LOGGER_TITLE")
@@ -89,6 +91,14 @@ def access_method(self):
   def bios_image_dir(self):
     return BIOS_IMAGES_DIR
 
+  @property
+  def new_driver_file(self):
+    return NEW_DRIVER_FILE
+
+  @property
+  def replaced_image_file(self):
+    return REPLACED_IMAGE_FILE
+
   @property
   def bios_image(self):
     return self.bios_roms[0]