BIOSUtilities/biosutilities/apple_efi_pbzx.py
Plato Mavropoulos 35455f735c BIOSUtilities v24.10.18
Removed all Python built-in library keyword arguments (#55)
2024-10-18 13:37:52 +03:00

133 lines
4.8 KiB
Python

#!/usr/bin/env python3 -B
# coding=utf-8
"""
Apple PBZX Extract
Apple EFI PBZX Extractor
Copyright (C) 2021-2024 Plato Mavropoulos
"""
import ctypes
import logging
import lzma
import os
from typing import Any, Final
from biosutilities.common.compression import is_szip_supported, szip_decompress
from biosutilities.common.paths import make_dirs, path_stem
from biosutilities.common.patterns import PAT_APPLE_PBZX
from biosutilities.common.structs import ctypes_struct, UINT32
from biosutilities.common.system import printer
from biosutilities.common.templates import BIOSUtility
from biosutilities.common.texts import file_to_bytes
class PbzxChunk(ctypes.BigEndianStructure):
""" PBZX Chunk Header """
_pack_ = 1
_fields_ = [
('Reserved0', UINT32), # 0x00
('InitSize', UINT32), # 0x04
('Reserved1', UINT32), # 0x08
('CompSize', UINT32) # 0x0C
# 0x10
]
def struct_print(self, padding: int = 0) -> None:
""" Display structure information """
printer(message=['Reserved 0 :', f'0x{self.Reserved0:X}'], padding=padding, new_line=False)
printer(message=['Initial Size :', f'0x{self.InitSize:X}'], padding=padding, new_line=False)
printer(message=['Reserved 1 :', f'0x{self.Reserved1:X}'], padding=padding, new_line=False)
printer(message=['Compressed Size:', f'0x{self.CompSize:X}'], padding=padding, new_line=False)
class AppleEfiPbzxExtract(BIOSUtility):
""" Apple EFI PBZX Extractor """
TITLE: str = 'Apple EFI PBZX Extractor'
PBZX_CHUNK_HDR_LEN: Final[int] = ctypes.sizeof(PbzxChunk)
def check_format(self, input_object: str | bytes | bytearray) -> bool:
""" Check if input is Apple PBZX image """
input_buffer: bytes = file_to_bytes(in_object=input_object)
return bool(PAT_APPLE_PBZX.search(input_buffer, 0, 4))
def parse_format(self, input_object: str | bytes | bytearray, extract_path: str, padding: int = 0) -> bool:
""" Parse & Extract Apple PBZX image """
input_buffer: bytes = file_to_bytes(in_object=input_object)
make_dirs(in_path=extract_path, delete=True)
cpio_bin: bytes = b'' # Initialize PBZX > CPIO Buffer
cpio_len: int = 0x0 # Initialize PBZX > CPIO Length
chunk_off: int = 0xC # First PBZX Chunk starts at 0xC
while chunk_off < len(input_buffer):
chunk_hdr: Any = ctypes_struct(buffer=input_buffer, start_offset=chunk_off, class_object=PbzxChunk)
printer(message=f'PBZX Chunk at 0x{chunk_off:08X}\n', padding=padding)
chunk_hdr.struct_print(padding=padding + 4)
# PBZX Chunk data starts after its Header
comp_bgn: int = chunk_off + self.PBZX_CHUNK_HDR_LEN
# To avoid a potential infinite loop, double-check Compressed Size
comp_end: int = comp_bgn + max(chunk_hdr.CompSize, self.PBZX_CHUNK_HDR_LEN)
comp_bin: bytes = input_buffer[comp_bgn:comp_end]
try:
# Attempt XZ decompression, if applicable to Chunk data
cpio_bin += lzma.LZMADecompressor().decompress(comp_bin)
printer(message='Successful LZMA decompression!', padding=padding + 8)
except Exception as error: # pylint: disable=broad-except
logging.debug('Error: Failed to LZMA decompress PBZX Chunk 0x%X: %s', chunk_off, error)
# Otherwise, Chunk data is not compressed
cpio_bin += comp_bin
# Final CPIO size should match the sum of all Chunks > Initial Size
cpio_len += chunk_hdr.InitSize
# Next Chunk starts at the end of current Chunk's data
chunk_off = comp_end
# Check that CPIO size is valid based on all Chunks > Initial Size
if cpio_len != len(cpio_bin):
printer(message='Error: Unexpected CPIO archive size!', padding=padding)
return False
cpio_name: str = path_stem(in_path=input_object) if isinstance(input_object, str) else 'Payload'
cpio_path: str = os.path.join(extract_path, f'{cpio_name}.cpio')
with open(cpio_path, 'wb') as cpio_object:
cpio_object.write(cpio_bin)
# Decompress PBZX > CPIO archive with 7-Zip
if is_szip_supported(in_path=cpio_path, padding=padding, args=['-tCPIO'], silent=False):
if szip_decompress(in_path=cpio_path, out_path=extract_path, in_name='CPIO',
padding=padding, args=['-tCPIO']):
os.remove(cpio_path) # Successful extraction, delete PBZX > CPIO archive
else:
return False
else:
return False
return True
if __name__ == '__main__':
AppleEfiPbzxExtract().run_utility()