| ## @file | |
| # This file is used to define the BIOS Tree Node. | |
| # | |
| # Copyright (c) 2021-, Intel Corporation. All rights reserved.<BR> | |
| # SPDX-License-Identifier: BSD-2-Clause-Patent | |
| ## | |
| from FirmwareStorageFormat.UPLHeader import * | |
| from FirmwareStorageFormat.FvHeader import * | |
| from FirmwareStorageFormat.FfsFileHeader import * | |
| from FirmwareStorageFormat.SectionHeader import * | |
| from FirmwareStorageFormat.PECOFFHeader import * | |
| from FirmwareStorageFormat.Common import * | |
| from utils.FmmtLogger import FmmtLogger as logger | |
| import uuid | |
| SectionHeaderType = { | |
| 0x01:'EFI_COMPRESSION_SECTION', | |
| 0x02:'EFI_GUID_DEFINED_SECTION', | |
| 0x03:'EFI_SECTION_DISPOSABLE', | |
| 0x10:'EFI_SECTION_PE32', | |
| 0x11:'EFI_SECTION_PIC', | |
| 0x12:'EFI_SECTION_TE', | |
| 0x13:'EFI_SECTION_DXE_DEPEX', | |
| 0x14:'EFI_SECTION_VERSION', | |
| 0x15:'EFI_SECTION_USER_INTERFACE', | |
| 0x16:'EFI_SECTION_COMPATIBILITY16', | |
| 0x17:'EFI_SECTION_FIRMWARE_VOLUME_IMAGE', | |
| 0x18:'EFI_FREEFORM_SUBTYPE_GUID_SECTION', | |
| 0x19:'EFI_SECTION_RAW', | |
| 0x1B:'EFI_SECTION_PEI_DEPEX', | |
| 0x1C:'EFI_SECTION_MM_DEPEX' | |
| } | |
| HeaderType = [0x01, 0x02, 0x14, 0x15, 0x18] | |
| class BinaryNode: | |
| def __init__(self, name: str) -> None: | |
| self.Size = 0 | |
| self.Name = "BINARY" + str(name) | |
| self.HOffset = 0 | |
| self.Data = b'' | |
| class ElfNode: | |
| def __init__(self, buffer: bytes) -> None: | |
| self.Header = ELF_HEADER32.from_buffer_copy(buffer) | |
| if self.Header.ELF_Identification[0:4] != b'\x7fELF': | |
| logger.error('Invalid Elf Header! Elf Identification {} is not ".ELF".'.format(self.Header.ELF_Identification)) | |
| raise Exception("Process Failed: Invalid ELF Header Identification!") | |
| self.Class = self.Header.ELF_Identification[4] | |
| if self.Class == 0x02: | |
| self.Header = ELF_HEADER64.from_buffer_copy(buffer) | |
| elif self.Class != 0x01: | |
| logger.error('Invalid Elf Class! Elf Class {} is not 0x01 or 0x02.'.format(self.Class)) | |
| raise Exception("Process Failed: Invalid ELF Class!") | |
| self.ProList = [] | |
| self.SecList = [] | |
| self.UpldInfoSection = None | |
| self.UpldInfo = None | |
| self.UpldBuffer = b'' | |
| self.Name = "ELF" | |
| self.HeaderLength = len(struct2stream(self.Header)) | |
| self.HOffset = 0 | |
| self.DOffset = 0 | |
| self.ROffset = 0 | |
| self.Data = b'' | |
| self.PadData = b'' | |
| self.Upld_Info_Align = False | |
| def GetProgramList(self, buffer: bytes) -> None: | |
| for i in range(self.Header.ELF_PHNum): | |
| if self.Class == 0x01: | |
| ElfProgramHeader = ELF_PROGRAM_HEADER32.from_buffer_copy(buffer[i*self.Header.ELF_PHEntSize:]) | |
| elif self.Class == 0x02: | |
| ElfProgramHeader = ELF_PROGRAM_HEADER64.from_buffer_copy(buffer[i*self.Header.ELF_PHEntSize:]) | |
| self.ProList.append(ElfProgramHeader) | |
| def GetSectionList(self, buffer: bytes) -> None: | |
| for i in range(self.Header.ELF_SHNum): | |
| if self.Class == 0x01: | |
| ElfSectionHeader = ELF_SECTION_HEADER32.from_buffer_copy(buffer[i*self.Header.ELF_SHEntSize:]) | |
| elif self.Class == 0x02: | |
| ElfSectionHeader = ELF_SECTION_HEADER64.from_buffer_copy(buffer[i*self.Header.ELF_SHEntSize:]) | |
| self.SecList.append(ElfSectionHeader) | |
| def FindUPLDSection(self, buffer: bytes) -> None: | |
| for item in self.SecList: | |
| if buffer[item.SH_Offset:item.SH_Offset+4] == b'PLDH' or buffer[item.SH_Offset:item.SH_Offset+4] == b'UPLD': | |
| self.UpldInfoSection = item | |
| self.UpldInfo = UNIVERSAL_PAYLOAD_INFO.from_buffer_copy(buffer[item.SH_Offset:item.SH_Offset+item.SH_Size]) | |
| self.UpldBuffer = struct2stream(self.UpldInfo) | |
| if (self.UpldInfoSection.SH_Offset) % 4 == 0: | |
| # if (self.UpldInfoSection.SH_Offset - self.Header.ELF_Entry) % 4 == 0: | |
| self.Upld_Info_Align = True | |
| class FvNode: | |
| def __init__(self, name, buffer: bytes) -> None: | |
| self.Header = EFI_FIRMWARE_VOLUME_HEADER.from_buffer_copy(buffer) | |
| Map_num = (self.Header.HeaderLength - 56)//8 | |
| self.Header = Refine_FV_Header(Map_num).from_buffer_copy(buffer) | |
| self.FvId = "FV" + str(name) | |
| self.Name = "FV" + str(name) | |
| if self.Header.ExtHeaderOffset: | |
| self.ExtHeader = EFI_FIRMWARE_VOLUME_EXT_HEADER.from_buffer_copy(buffer[self.Header.ExtHeaderOffset:]) | |
| self.Name = uuid.UUID(bytes_le=struct2stream(self.ExtHeader.FvName)) | |
| self.ExtEntryOffset = self.Header.ExtHeaderOffset + 20 | |
| if self.ExtHeader.ExtHeaderSize != 20: | |
| self.ExtEntryExist = 1 | |
| self.ExtEntry = EFI_FIRMWARE_VOLUME_EXT_ENTRY.from_buffer_copy(buffer[self.ExtEntryOffset:]) | |
| self.ExtTypeExist = 1 | |
| if self.ExtEntry.ExtEntryType == 0x01: | |
| nums = (self.ExtEntry.ExtEntrySize - 8) // 16 | |
| self.ExtEntry = Refine_FV_EXT_ENTRY_OEM_TYPE_Header(nums).from_buffer_copy(buffer[self.ExtEntryOffset:]) | |
| elif self.ExtEntry.ExtEntryType == 0x02: | |
| nums = self.ExtEntry.ExtEntrySize - 20 | |
| self.ExtEntry = Refine_FV_EXT_ENTRY_GUID_TYPE_Header(nums).from_buffer_copy(buffer[self.ExtEntryOffset:]) | |
| elif self.ExtEntry.ExtEntryType == 0x03: | |
| self.ExtEntry = EFI_FIRMWARE_VOLUME_EXT_ENTRY_USED_SIZE_TYPE.from_buffer_copy(buffer[self.ExtEntryOffset:]) | |
| else: | |
| self.ExtTypeExist = 0 | |
| else: | |
| self.ExtEntryExist = 0 | |
| self.Size = self.Header.FvLength | |
| self.HeaderLength = self.Header.HeaderLength | |
| self.HOffset = 0 | |
| self.DOffset = 0 | |
| self.ROffset = 0 | |
| self.Data = b'' | |
| if self.Header.Signature != 1213613663: | |
| logger.error('Invalid Fv Header! Fv {} signature {} is not "_FVH".'.format(struct2stream(self.Header), self.Header.Signature)) | |
| raise Exception("Process Failed: Fv Header Signature!") | |
| self.PadData = b'' | |
| self.Free_Space = 0 | |
| self.ModCheckSum() | |
| def ModCheckSum(self) -> None: | |
| # Fv Header Sums to 0. | |
| Header = struct2stream(self.Header)[::-1] | |
| Size = self.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) -> None: | |
| # 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) -> None: | |
| # 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) -> None: | |
| if self.Header.ExtHeaderOffset: | |
| ExtHeaderData = struct2stream(self.ExtHeader) | |
| ExtHeaderDataOffset = self.Header.ExtHeaderOffset - self.HeaderLength | |
| self.Data = self.Data[:ExtHeaderDataOffset] + ExtHeaderData + self.Data[ExtHeaderDataOffset+20:] | |
| if self.Header.ExtHeaderOffset and self.ExtEntryExist: | |
| ExtHeaderEntryData = 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, buffer: bytes) -> None: | |
| self.Header = EFI_FFS_FILE_HEADER.from_buffer_copy(buffer) | |
| # self.Attributes = unpack("<B", buffer[21:22])[0] | |
| if self.Header.FFS_FILE_SIZE != 0 and self.Header.Attributes != 0xff and self.Header.Attributes & 0x01 == 1: | |
| logger.error('Error Ffs Header! Ffs {} Header Size and Attributes is not matched!'.format(uuid.UUID(bytes_le=struct2stream(self.Header.Name)))) | |
| raise Exception("Process Failed: Error Ffs Header!") | |
| if self.Header.FFS_FILE_SIZE == 0 and self.Header.Attributes & 0x01 == 1: | |
| self.Header = EFI_FFS_FILE_HEADER2.from_buffer_copy(buffer) | |
| self.Name = uuid.UUID(bytes_le=struct2stream(self.Header.Name)) | |
| self.UiName = b'' | |
| self.Version = b'' | |
| self.Size = self.Header.FFS_FILE_SIZE | |
| self.HeaderLength = self.Header.HeaderLength | |
| self.HOffset = 0 | |
| self.DOffset = 0 | |
| self.ROffset = 0 | |
| self.Data = b'' | |
| self.PadData = b'' | |
| self.SectionMaxAlignment = SECTION_COMMON_ALIGNMENT # 4-align | |
| self.PeCoffSecIndex = None | |
| self.IsFsp = False | |
| self.IsVtf = False | |
| self.IfFsp() | |
| self.IfVtf() | |
| def ModCheckSum(self) -> None: | |
| HeaderData = 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 | |
| def IfFsp(self) -> None: | |
| if self.Name == EFI_FSP_GUID: | |
| self.IsFsp = True | |
| def IfVtf(self) -> None: | |
| if self.Name == EFI_FFS_VOLUME_TOP_FILE_GUID: | |
| self.IsVtf = True | |
| class SectionNode: | |
| def __init__(self, buffer: bytes) -> None: | |
| if buffer[0:3] != b'\xff\xff\xff': | |
| self.Header = EFI_COMMON_SECTION_HEADER.from_buffer_copy(buffer) | |
| else: | |
| self.Header = EFI_COMMON_SECTION_HEADER2.from_buffer_copy(buffer) | |
| if self.Header.Type in SectionHeaderType: | |
| self.Name = SectionHeaderType[self.Header.Type] | |
| elif self.Header.Type == 0: | |
| self.Name = "EFI_SECTION_ALL" | |
| else: | |
| self.Name = "SECTION" | |
| self.IsPadSection = False | |
| self.IsUiSection = False | |
| self.IsVerSection = False | |
| if self.Header.Type in HeaderType: | |
| self.ExtHeader = self.GetExtHeader(self.Header.Type, buffer[self.Header.Common_Header_Size():], (self.Header.SECTION_SIZE-self.Header.Common_Header_Size())) | |
| self.HeaderLength = self.Header.Common_Header_Size() + self.ExtHeader.ExtHeaderSize() | |
| else: | |
| self.ExtHeader = None | |
| self.HeaderLength = self.Header.Common_Header_Size() | |
| self.Size = self.Header.SECTION_SIZE | |
| self.Type = self.Header.Type | |
| self.HOffset = 0 | |
| self.DOffset = 0 | |
| self.ROffset = 0 | |
| self.Data = b'' | |
| self.OriData = b'' | |
| self.OriHeader = b'' | |
| self.PadData = b'' | |
| self.IsPadSection = False | |
| self.SectionMaxAlignment = SECTION_COMMON_ALIGNMENT # 4-align | |
| def GetExtHeader(self, Type: int, buffer: bytes, nums: int=0) -> None: | |
| if Type == 0x01: | |
| return EFI_COMPRESSION_SECTION.from_buffer_copy(buffer) | |
| elif Type == 0x02: | |
| return EFI_GUID_DEFINED_SECTION.from_buffer_copy(buffer) | |
| elif Type == 0x14: | |
| self.IsVerSection = True | |
| return Get_VERSION_Header((nums - 2)//2).from_buffer_copy(buffer) | |
| elif Type == 0x15: | |
| self.IsUiSection = True | |
| return Get_USER_INTERFACE_Header(nums//2).from_buffer_copy(buffer) | |
| elif Type == 0x18: | |
| return EFI_FREEFORM_SUBTYPE_GUID_SECTION.from_buffer_copy(buffer) | |
| class PeCoffNode: | |
| def __init__(self, buffer: bytes, offset: int, size: int = 0) -> None: | |
| self.Name = 'PeCoff' | |
| self.offset = offset | |
| self.Size = size | |
| self.OriData = buffer | |
| self.Data = buffer | |
| self.IsTeImage = True | |
| self.RelocationsStripped = True | |
| self.PeCoffHeaderOffset = 0 | |
| self.ImageAddress = 0 | |
| self.DestinationAddress = 0 | |
| self.DebugDirectoryEntryVirtualAddress = 0 | |
| self.PeHeader = None | |
| self.TeHeader = None | |
| self.Machine = None | |
| self.ImageType = None | |
| self.OptionalHeader = None | |
| self.BlkHeaderOffset = 0 | |
| self.BlkHeader = None | |
| self.RelocationsFieldSize = 0 | |
| self.RelocationsData = None | |
| self.RelocList = [] | |
| self.SizeOfImage = 0 | |
| self.HOffset = self.offset | |
| self.DOffset = 0 | |
| self.IfRebase = True | |
| self.TeHeader = EFI_TE_IMAGE_HEADER.from_buffer_copy(self.Data) | |
| if self.TeHeader.Signature == EFI_IMAGE_DOS_SIGNATURE: | |
| self.IsTeImage = False | |
| self.TeHeader = None | |
| self.DOffset = 0 | |
| self.DosHeader = EFI_IMAGE_DOS_HEADER.from_buffer_copy(self.Data) | |
| if self.DosHeader.e_magic == EFI_IMAGE_DOS_SIGNATURE: | |
| self.PeCoffHeaderOffset = self.DosHeader.e_lfanew #0xC0 | |
| self.PeHeader = EFI_IMAGE_OPTIONAL_HEADER_UNION.from_buffer_copy(self.Data[self.PeCoffHeaderOffset:]) | |
| if self.PeHeader.Pe32.Signature != EFI_IMAGE_NT_SIGNATURE: | |
| self.TeHeader = self.PeHeader.Te | |
| if self.TeHeader.Signature != EFI_TE_IMAGE_HEADER_SIGNATURE: | |
| logger.error('Invalid Te Header! Te signature {} is not "VZ".'.format(self.TeHeader.Signature)) | |
| raise Exception('Not support TeHeader which signature is not "VZ"!') | |
| self.IsTeImage = True | |
| self.PeCoffLoaderCheckImageType() | |
| self.PeCoffRebaseFlag() | |
| if self.IfRebase: | |
| self.PeCoffParseReloc() | |
| def HasRelocTable(self) -> bool: | |
| """ | |
| Check if the PE/COFF image has a valid relocation table. | |
| Returns True if the relocation table exists and is non-empty, False otherwise. | |
| """ | |
| # For TE image | |
| if self.IsTeImage: | |
| # DataDirectory[0] is Base Relocation Table | |
| if hasattr(self.TeHeader, 'DataDirectory') and self.TeHeader.DataDirectory[0].Size > 0: | |
| return True | |
| return False | |
| # For PE image | |
| if self.PeHeader: | |
| if hasattr(self.PeHeader.Pe32, 'OptionalHeader'): | |
| dir_entry = self.PeHeader.Pe32.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC] | |
| if dir_entry.Size > 0: | |
| return True | |
| if hasattr(self.PeHeader, 'Pe32Plus') and hasattr(self.PeHeader.Pe32Plus, 'OptionalHeader'): | |
| dir_entry = self.PeHeader.Pe32Plus.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC] | |
| if dir_entry.Size > 0: | |
| return True | |
| return False | |
| def FillPeReloc(self) -> bool: | |
| """ | |
| Supplement the PE/COFF relocation table if missing or incomplete. | |
| Returns True if the table was supplemented, False otherwise. | |
| """ | |
| from ctypes import sizeof, c_uint16, c_uint32 | |
| from FirmwareStorageFormat.PECOFFHeader import EFI_IMAGE_BASE_RELOCATION, EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC, EFI_IMAGE_SIZEOF_BASE_RELOCATION | |
| # If already has a reloc table, do nothing | |
| if self.HasRelocTable(): | |
| return False | |
| # For TE image | |
| if self.IsTeImage and self.TeHeader: | |
| if hasattr(self.TeHeader, 'DataDirectory') and self.TeHeader.DataDirectory[0].Size == 0: | |
| # Check if Data is valid | |
| if not isinstance(self.Data, (bytes, bytearray)): | |
| logger.error('TE image Data is not bytes/bytearray') | |
| return False | |
| # Construct the minimal valid reloc block | |
| reloc = EFI_IMAGE_BASE_RELOCATION() | |
| reloc.VirtualAddress = 0 | |
| reloc.SizeOfBlock = EFI_IMAGE_SIZEOF_BASE_RELOCATION + 2 # header+1个WORD | |
| reloc_block = bytes(reloc) + (c_uint16(0).value).to_bytes(2, 'little') | |
| # Boundary check: DataDirectory[0].VirtualAddress must not overflow | |
| va = len(self.Data) | |
| if va > 0xFFFFFFFF: | |
| logger.error('Reloc VA overflow for TE image') | |
| return False | |
| self.TeHeader.DataDirectory[0].VirtualAddress = va | |
| self.TeHeader.DataDirectory[0].Size = len(reloc_block) | |
| self.Data += reloc_block | |
| return True | |
| return False | |
| # For PE image | |
| if self.PeHeader: | |
| # 32-bit PE | |
| if hasattr(self.PeHeader.Pe32, 'OptionalHeader'): | |
| opt = self.PeHeader.Pe32.OptionalHeader | |
| # Check NumberOfRvaAndSizes | |
| if hasattr(opt, 'NumberOfRvaAndSizes') and opt.NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC: | |
| dir_entry = opt.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC] | |
| if dir_entry.Size == 0: | |
| # Check if Data is valid | |
| if not isinstance(self.Data, (bytes, bytearray)): | |
| logger.error('PE image Data is not bytes/bytearray') | |
| return False | |
| reloc = EFI_IMAGE_BASE_RELOCATION() | |
| reloc.VirtualAddress = 0 | |
| reloc.SizeOfBlock = EFI_IMAGE_SIZEOF_BASE_RELOCATION + 2 | |
| reloc_block = bytes(reloc) + (c_uint16(0).value).to_bytes(2, 'little') | |
| va = len(self.Data) | |
| if va > 0xFFFFFFFF: | |
| logger.error('Reloc VA overflow for PE32') | |
| return False | |
| dir_entry.VirtualAddress = va | |
| dir_entry.Size = len(reloc_block) | |
| self.Data += reloc_block | |
| return True | |
| # 64-bit PE | |
| if hasattr(self.PeHeader, 'Pe32Plus') and hasattr(self.PeHeader.Pe32Plus, 'OptionalHeader'): | |
| opt = self.PeHeader.Pe32Plus.OptionalHeader | |
| if hasattr(opt, 'NumberOfRvaAndSizes') and opt.NumberOfRvaAndSizes > EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC: | |
| dir_entry = opt.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC] | |
| if dir_entry.Size == 0: | |
| if not isinstance(self.Data, (bytes, bytearray)): | |
| logger.error('PE+ image Data is not bytes/bytearray') | |
| return False | |
| reloc = EFI_IMAGE_BASE_RELOCATION() | |
| reloc.VirtualAddress = 0 | |
| reloc.SizeOfBlock = EFI_IMAGE_SIZEOF_BASE_RELOCATION + 2 | |
| reloc_block = bytes(reloc) + (c_uint16(0).value).to_bytes(2, 'little') | |
| va = len(self.Data) | |
| if va > 0xFFFFFFFF: | |
| logger.error('Reloc VA overflow for PE32+') | |
| return False | |
| dir_entry.VirtualAddress = va | |
| dir_entry.Size = len(reloc_block) | |
| self.Data += reloc_block | |
| return True | |
| return False | |
| def PeCoffLoaderCheckImageType(self) -> None: | |
| MachineTypeList = [EFI_IMAGE_FILE_MACHINE_I386, EFI_IMAGE_FILE_MACHINE_EBC, EFI_IMAGE_FILE_MACHINE_X64, EFI_IMAGE_FILE_MACHINE_ARMT, EFI_IMAGE_FILE_MACHINE_ARM64, EFI_IMAGE_FILE_MACHINE_RISCV64] | |
| ImageTypeList = [EFI_IMAGE_SUBSYSTEM_EFI_APPLICATION, EFI_IMAGE_SUBSYSTEM_EFI_BOOT_SERVICE_DRIVER, EFI_IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER, EFI_IMAGE_SUBSYSTEM_SAL_RUNTIME_DRIVER] | |
| # Check Machine Type | |
| if self.IsTeImage: | |
| self.Machine = self.TeHeader.Machine | |
| else: | |
| self.Machine = self.PeHeader.Pe32.FileHeader.Machine | |
| if self.Machine not in MachineTypeList: | |
| if self.Machine == EFI_IMAGE_FILE_MACHINE_ARM: | |
| self.Machine = EFI_IMAGE_FILE_MACHINE_ARMT | |
| if self.IsTeImage: | |
| self.TeHeader.Machine = self.Machine | |
| else: | |
| self.PeHeader.Pe32.FileHeader.Machine = self.Machine | |
| else: | |
| logger.error('The Machine Type {} is not supported!'.format(self.Machine)) | |
| raise Exception('The Machine Type {} is not supported!'.format(self.Machine)) | |
| # Check Image Type | |
| if self.IsTeImage: | |
| self.ImageType = self.TeHeader.Subsystem | |
| else: | |
| self.ImageType = self.PeHeader.Pe32.OptionalHeader.Subsystem | |
| if self.ImageType not in ImageTypeList: | |
| logger.error('The Image Type {} is not supported!'.format(self.ImageType)) | |
| raise Exception('The Image Type {} is not supported!'.format(self.ImageType)) | |
| def PeCoffRebaseFlag(self): | |
| if self.IsTeImage: | |
| pass | |
| else: | |
| if self.PeHeader.Pe32.OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC: | |
| self.SizeOfImage = self.PeHeader.Pe32.OptionalHeader.SizeOfImage | |
| else: | |
| self.SizeOfImage = self.PeHeader.Pe32Plus.OptionalHeader.SizeOfImage | |
| if self.SizeOfImage != len(self.Data): | |
| self.IfRebase = False | |
| def PeCoffParseReloc(self) -> None: | |
| if self.IsTeImage: | |
| self.ImageAddress = self.TeHeader.ImageBase + self.TeHeader.StrippedSize - EFI_TE_IMAGE_HEADER_SIZE | |
| self.BlkHeaderOffset = self.offset + EFI_TE_IMAGE_HEADER_SIZE - self.TeHeader.StrippedSize +self.TeHeader.DataDirectory[0].VirtualAddress | |
| self.BlkSize = self.TeHeader.DataDirectory[0].Size | |
| else: | |
| if self.PeHeader.Pe32.OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC: | |
| self.ImageAddress = self.PeHeader.Pe32.OptionalHeader.ImageBase | |
| self.BlkHeaderOffset = self.offset + self.PeHeader.Pe32.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress | |
| self.BlkSize = self.PeHeader.Pe32.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size | |
| else: | |
| self.ImageAddress = self.PeHeader.Pe32Plus.OptionalHeader.ImageBase | |
| self.BlkHeaderOffset = self.offset + self.PeHeader.Pe32Plus.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress | |
| self.BlkSize = self.PeHeader.Pe32Plus.OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC].Size | |
| CurOff = self.BlkHeaderOffset | |
| while CurOff < self.BlkHeaderOffset + self.BlkSize: | |
| if CurOff % 4: | |
| CurOff += 4 - CurOff % 4 | |
| self.BlkHeader = EFI_BLK_HEADER.from_buffer_copy(self.Data[CurOff-self.offset:]) | |
| self.RelocationsFieldSize = self.BlkHeader.BlockSize - EFI_BLK_HEADER_SIZE | |
| self.RelocationsData = (c_uint16 * int(self.RelocationsFieldSize/2)).from_buffer_copy(self.Data[CurOff - self.offset + EFI_BLK_HEADER_SIZE:CurOff - self.offset + EFI_BLK_HEADER_SIZE + self.RelocationsFieldSize]) | |
| for EachDataField in self.RelocationsData: | |
| # Rtype [15:12] Roffset [11:0] | |
| EachRType = EachDataField >> 12 | |
| EachROff = EachDataField & 0xfff | |
| if EachRType == 0: # IMAGE_REL_BASED_ABSOLUTE | |
| continue | |
| if ((EachRType != 3) and (EachRType != 10)): # IMAGE_REL_BASED_HIGHLOW and IMAGE_REL_BASED_DIR64 | |
| logger.error("ERROR: Unsupported relocation type %d!" % EachRType) | |
| raise Exception("ERROR: Unsupported relocation type %d!" % EachRType) | |
| if self.TeHeader: | |
| TarROff = self.offset + self.BlkHeader.PageRVA + EachROff + EFI_TE_IMAGE_HEADER_SIZE - self.TeHeader.StrippedSize | |
| else: | |
| TarROff = self.offset + self.BlkHeader.PageRVA + EachROff | |
| self.RelocList.append((EachRType, TarROff)) | |
| CurOff += self.BlkHeader.BlockSize | |
| def PeCoffRebase(self, DeltaSize=0, CalcuFlag=0, base_address=None, xip_offset=0) -> None: | |
| """ | |
| Support absolute base address (base_address) and XIP offset (xip_offset) parameters, compatible with original DeltaSize/CalcuFlag logic. | |
| base_address has higher priority than DeltaSize/CalcuFlag. | |
| """ | |
| # If base_address is explicitly specified, force absolute relocation | |
| # XIP offset can be used for future extension (if needed) | |
| ## Rebase function | |
| # Architecture relocation type table driven | |
| SUPPORTED_RELOC_TYPES = { | |
| 3: 4, # IMAGE_REL_BASED_HIGHLOW (x86/x64 32bit) | |
| 10: 8, # IMAGE_REL_BASED_DIR64 (x64 64bit) | |
| # Extendable: e.g. ARM/ARM64 etc. | |
| } | |
| # If base_address is explicitly specified, force absolute relocation | |
| if base_address is not None: | |
| CalcuFlag = 1 | |
| DeltaSize = base_address | |
| # XIP offset can be used for future extension (if needed) | |
| if self.TeHeader: | |
| ImageBase = self.TeHeader.ImageBase | |
| CurOff = self.offset + EFI_TE_IMAGE_HEADER.ImageBase.offset | |
| ImageBaseSize = EFI_TE_IMAGE_HEADER.ImageBase.size | |
| else: | |
| CurOff = self.offset + self.PeCoffHeaderOffset | |
| CurOff += EFI_IMAGE_NT_HEADERS32.OptionalHeader.offset | |
| if self.PeHeader.Pe32.OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC: # PE32+ image | |
| ImageBase = EFI_IMAGE_OPTIONAL_HEADER64.ImageBase | |
| CurOff += EFI_IMAGE_OPTIONAL_HEADER64.ImageBase.offset | |
| ImageBaseSize = EFI_IMAGE_OPTIONAL_HEADER64.ImageBase.size | |
| else: | |
| ImageBase = EFI_IMAGE_OPTIONAL_HEADER32.ImageBase | |
| CurOff += EFI_IMAGE_OPTIONAL_HEADER32.ImageBase.offset | |
| ImageBaseSize = EFI_IMAGE_OPTIONAL_HEADER32.ImageBase.size | |
| if CalcuFlag: | |
| NewImageBase = DeltaSize | |
| if self.TeHeader: | |
| DeltaSize = DeltaSize - self.TeHeader.ImageBase | |
| self.TeHeader.ImageBase = NewImageBase | |
| self.ImageAddress = self.TeHeader.ImageBase + self.TeHeader.StrippedSize - EFI_TE_IMAGE_HEADER_SIZE | |
| else: | |
| DeltaSize = DeltaSize - self.ImageAddress | |
| if self.PeHeader.Pe32.OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC: | |
| self.PeHeader.Pe32.OptionalHeader.ImageBase = NewImageBase | |
| self.ImageAddress = self.PeHeader.Pe32.OptionalHeader.ImageBase | |
| else: | |
| self.PeHeader.Pe32Plus.OptionalHeader.ImageBase = NewImageBase | |
| self.ImageAddress = self.PeHeader.Pe32Plus.OptionalHeader.ImageBase | |
| CurValue = NewImageBase | |
| else: | |
| if self.TeHeader: | |
| self.TeHeader.ImageBase += DeltaSize | |
| self.ImageAddress = self.TeHeader.ImageBase + self.TeHeader.StrippedSize - EFI_TE_IMAGE_HEADER_SIZE | |
| CurValue = self.TeHeader.ImageBase | |
| else: | |
| if self.PeHeader.Pe32.OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC: | |
| self.PeHeader.Pe32.OptionalHeader.ImageBase += DeltaSize | |
| self.ImageAddress = self.PeHeader.Pe32.OptionalHeader.ImageBase | |
| else: | |
| self.PeHeader.Pe32Plus.OptionalHeader.ImageBase += DeltaSize | |
| self.ImageAddress = self.PeHeader.Pe32Plus.OptionalHeader.ImageBase | |
| CurValue = self.ImageAddress | |
| self.Data = self.Data[:CurOff-self.offset] + CurValue.to_bytes(ImageBaseSize, byteorder='little',signed=False) + self.Data[CurOff-self.offset+ImageBaseSize:] | |
| for (EachRType, TarROff) in self.RelocList: | |
| if EachRType in SUPPORTED_RELOC_TYPES: | |
| size = SUPPORTED_RELOC_TYPES[EachRType] | |
| CurValue = Bytes2Val(self.Data[TarROff-self.offset:TarROff+size-self.offset]) | |
| CurValue += DeltaSize | |
| self.Data = self.Data[:TarROff-self.offset] + CurValue.to_bytes(size, byteorder='little',signed=False) + self.Data[TarROff+size-self.offset:] | |
| else: | |
| logger.error(f"ERROR: Unsupported relocation type {EachRType}! (please extend SUPPORTED_RELOC_TYPES)") | |
| raise Exception(f"ERROR: Unsupported relocation type {EachRType}! (please extend SUPPORTED_RELOC_TYPES)") | |
| class FreeSpaceNode: | |
| def __init__(self, buffer: bytes) -> None: | |
| self.Name = 'Free_Space' | |
| self.Data = buffer | |
| self.Size = len(buffer) | |
| self.HOffset = 0 | |
| self.DOffset = 0 | |
| self.ROffset = 0 | |
| self.PadData = b'' |