diff --git a/README.md b/README.md index 1a89473..42d5c64 100644 --- a/README.md +++ b/README.md @@ -856,6 +856,6 @@ b'\x01\x00\x00\x00\x1b\xb9\x87\xc3YIF\x06\xb1t\x94V\x01\xc9\xea[\x1b\xb9\x87\xc3 [...] > my_stuff = [...] > psp.blob.roms[0].directories[0].files[1].move_buffer(0x60000, 0x1000) -> psp.blob.roms[0].set_bytes(0x60000, 0x1000, my_stuff) +> psp.blob.roms[0].set_bytes(0x60000, my_stuff) > psp.to_file('my_modified_bios.bin') -``` \ No newline at end of file +``` diff --git a/psptool/__main__.py b/psptool/__main__.py index bd970aa..1a29294 100755 --- a/psptool/__main__.py +++ b/psptool/__main__.py @@ -123,7 +123,8 @@ def main(): '-r idx: specifies rom_index (default: 0)', '-d idx: specifies directory_index', '-e idx: specifies file_index', - '-s file: specifies subfile (i.e. the new file contents)', + '-s file: specifies substitution file (i.e. the new file contents)', + '(optional to allow plain re-signs)', '-o file: specifies outfile', '-p file: specifies file-stub (e.g. \'keys/id\') for the re-signing keys', '-a pass: specifies password for the re-signing keys' @@ -239,13 +240,13 @@ def main(): if args.directory_index is not None and args.file_index is not None and args.outfile is not None: file = psp.blob.roms[args.rom_index].directories[args.directory_index].files[args.file_index] - # Substituting an file is actually optional to allow plain re-signs + # Substituting a file is actually optional to allow plain re-signs if args.subfile is not None: with open(args.subfile, 'rb') as f: sub_binary = f.read() # Keep the existing file's address, but adapt its size - file.move_buffer(file.get_address(), len(sub_binary)) - file.set_bytes(0, len(sub_binary), sub_binary) + file.move_buffer(file.get_address(), len(sub_binary) + file.signature.len) + file.set_bytes(0, sub_binary) privkeys = None if args.privkeystub: diff --git a/psptool/cert_tree.py b/psptool/cert_tree.py index 940b15b..d151377 100644 --- a/psptool/cert_tree.py +++ b/psptool/cert_tree.py @@ -114,7 +114,7 @@ def resign_only(self, privkey: PrivateKey): assert len(signature) == self.signature.buffer_size, f'Could not resign {self} with {privkey}: ' \ f'The new signature has the wrong length ' \ f'{len(signature)} != {self.signature.buffer_size}' - self.signature.set_bytes(0, len(signature), signature) + self.signature.set_bytes(0, signature) def resign_and_replace(self, privkeys: PrivateKeyDict = None, recursive: bool = False): # this resigns self (multiple times!) @@ -215,7 +215,7 @@ def get_public_key(self) -> PublicKey: def replace_crypto_material(self, crypto_material: bytes): assert len(crypto_material) == self._crypto_material.buffer_size, f'Crypto material has wrong size: {len(crypto_material)} != {self._crypto_material.buffer_size}' self._public_key = None - self._crypto_material.set_bytes(0, len(crypto_material), crypto_material) + self._crypto_material.set_bytes(0, crypto_material) def replace_only(self, pubkey: PublicKey): print(f'Replacing {self}') diff --git a/psptool/directory.py b/psptool/directory.py index eb29816..2668005 100644 --- a/psptool/directory.py +++ b/psptool/directory.py @@ -16,7 +16,7 @@ import struct -from .entry import DirectoryEntry, BiosDirectoryEntry +from .entry import DirectoryEntry, BiosDirectoryEntry, AddressMode from .utils import NestedBuffer, fletcher32 from .file import File, BiosFile, SECONDARY_DIRECTORY_ENTRY_TYPES, TERTIARY_DIRECTORY_ENTRY_TYPES @@ -174,6 +174,7 @@ def __init__(self, parent_rom, offset: int, psptool, zen_generation=None): # create/link files for entry in self.entries: file = self.FILE_CLASS.create_file_if_not_exists(self, entry) + entry.file = file if file is not None: self.files.append(file) @@ -204,18 +205,14 @@ def count(self, value): self.header[8:12] = struct.pack(' AddressMode: info = struct.unpack('> 31) & 1 if version == 1: - return (info >> 24) & 3 + return AddressMode((info >> 24) & 3) else: - return (info >> 29) & 3 + return AddressMode((info >> 29) & 3) def verify_checksum(self): data = self[0x8:] @@ -227,40 +224,28 @@ def verify_checksum(self): def update_checksum(self): data = self[0x8:] # checksum is calculated from after the checksum field in the header - self.checksum.set_bytes(0, 4, fletcher32(data)) + self.checksum.set_bytes(0, fletcher32(data)) def update_entry_fields(self, file: File, type_, size, offset): - # 1. Find respective Entry for a given File + # 1. Find respective entry for a given file entry = None - for index, my_entry in enumerate(self.entries): - # We assume that each directory has at most one entry of a given type - if my_entry.type == file.type: - entry = my_entry + for e in self.entries: + if e.file is file: + entry = e break - assert(entry is not None) + assert entry is not None, f'No entry for {file} in {self}' # 2. Update fields entry.type = type_ if entry.type not in File.NO_SIZE_ENTRY_TYPES: entry.size = size - # Convert the ROM buffer offset back to the value the entry expects, preserving address mode. - # This mirrors the inverse of file_offset() in entry.py. - addr_mode = self.address_mode - if addr_mode == 2 or addr_mode == 3: - addr_mode = entry.address_mode - - if addr_mode == 0: - # x86 physical: preserve upper bits (e.g. 0xFF000000), replace lower bits - addr_mask = self.rom.addr_mask - upper_mask = 0xFFFFFFFF ^ addr_mask - entry.offset = (entry.offset & upper_mask) | (offset & addr_mask) - elif addr_mode == 1: - entry.offset = offset - elif addr_mode == 2 or addr_mode == 3: - entry.offset = offset - self.buffer_offset - else: - entry.offset = offset + # file_offset() is affine in entry.offset (within the + # address mask), so shifting entry.offset by the file's + # displacement keeps the address mode and upper bits + # intact, and leaves in-place replacements untouched. + entry.offset += offset - entry.file_offset() + assert entry.file_offset() == offset, f'Offset {offset:#x} is not representable by {entry}' # 3. Update checksum self.update_checksum() diff --git a/psptool/entry.py b/psptool/entry.py index 010343c..c8b2aec 100644 --- a/psptool/entry.py +++ b/psptool/entry.py @@ -17,11 +17,17 @@ import struct from typing import TYPE_CHECKING +from enum import IntEnum if TYPE_CHECKING: from directory import Directory from .utils import NestedBuffer +class AddressMode(IntEnum): + PHYSICAL = 0 # x86 Physical address + REL_BIOS = 1 # Offset from start of the BIOS (flash offset) most common on modern systems + REL_DIR = 2 # Offset from start of directory header + REL_SLOT = 3 # Offset from start of partition class DirectoryEntry(NestedBuffer): ENTRY_SIZE = 4 * 4 @@ -30,6 +36,7 @@ def __init__(self, parent_directory: 'Directory', entry_offset): super().__init__(parent_directory.body, self.ENTRY_SIZE, entry_offset) self.parent_directory = parent_directory self.entry_offset = entry_offset + self.file = None def __repr__(self): return f'{self.__class__.__name__}({self.type=:#x}, {self.flags=:#x}, {self.size=:#x}, ' \ @@ -42,17 +49,17 @@ def file_offset(self): return dir_start + self.entry_offset addr_mode = self.parent_directory.address_mode - # If directory address mode is 2 or 3 (relative to dir or slot), the + # If directory address mode is relative to dir or slot, the # entry address mode must be taken into account, otherwise ignored. - if addr_mode == 2 or addr_mode == 3: + if addr_mode == AddressMode.REL_DIR or addr_mode == AddressMode.REL_SLOT: addr_mode = self.address_mode # coreboot's amdfwtool writes APOB NV entries on older SoCs as an x86 - # physical address (entry address mode 0) in a mode 1 directory. No + # physical address in a directory with relative-to-bios mode. No # flash offset reaches 0xff000000. - elif addr_mode == 1 and self.address_mode == 0 and self.offset >= 0xFF000000: - addr_mode = 0 + elif addr_mode == AddressMode.REL_BIOS and self.address_mode == AddressMode.PHYSICAL and self.offset >= 0xFF000000: + addr_mode = AddressMode.PHYSICAL - if addr_mode == 0: + if addr_mode == AddressMode.PHYSICAL: # x86 physical address, should be in range 0xff000000 - 0xffffffff # But some images use flash offset in x86 physical address mode. # If ROM is bigger than 16MB and the entry address is in the @@ -63,14 +70,11 @@ def file_offset(self): return self.offset & 0x00FFFFFF else: return self.offset & self.parent_directory.rom.addr_mask - elif addr_mode == 1: - # Flash offset from start of BIOS, most common on modern systems + elif addr_mode == AddressMode.REL_BIOS: return self.offset - elif addr_mode == 2: - # Flash offset from start of directory header + elif addr_mode == AddressMode.REL_DIR: return self.parent_directory.buffer_offset + self.offset - elif addr_mode == 3: - # Flash offset from start of the slot + elif addr_mode == AddressMode.REL_SLOT: # TODO: How to calculate the offset from slot? Is this correct? return self.parent_directory.buffer_offset + self.offset @@ -80,7 +84,7 @@ def type(self): @type.setter def type(self, value): - self.set_bytes(0, 1, struct.pack('> 30) & 3 + def address_mode(self) -> AddressMode: + return AddressMode((self.rsv0 >> 30) & 3) class BiosDirectoryEntry(DirectoryEntry): ENTRY_SIZE = 4 * 6 @@ -150,7 +150,7 @@ def destination(self): @destination.setter def destination(self, value): - self.set_bytes(16, 8, struct.pack(' bool: + return self is not SignatureType.NONE + + @classmethod + def from_header(cls, signed: NestedBuffer, type: int) -> "SignatureType": + signed = int.from_bytes(signed.get_bytes(), 'little') + if signed not in {0, 1, 0xffff0000}: + raise File.ParseError(f'Did not expect signed to be 0x{signed:x}') + if signed == 0: + return SignatureType.NONE + for st in cls: + if st.type == type: + return st + return File.ParseError(f"Unknown signature type 0x{type:x}") + +class Signature(NestedBuffer): + def __init__(self, file: 'HeaderFile', type: SignatureType): + self.type = type + super().__init__(file, type.len, file.rom_size - type.len) + + @property + def len(self) -> int: + return self.type.len + class HeaderFile(File): @@ -44,8 +79,9 @@ def _parse(self): self.encrypted = struct.unpack(' bool: @property def is_signed(self) -> bool: - signed = int.from_bytes(self._signed.get_bytes(), 'little') - if signed not in {0, 1, 0xffff0000}: - raise self.ParseError(f'Did not expect signed to be 0x{signed:x}') - return signed != 0 + return self.signature.type.is_signed def get_readable_version(self): return '.'.join([hex(b)[2:].upper() for b in self.version]) @@ -215,4 +241,4 @@ def md5(self): m.update(self.body.get_bytes()) except: self.psptool.ph.print_warning(f"Get bytes failed at file: 0x{self.get_address():x} type: {self.get_readable_type()} size: 0x{self.buffer_size:x}") - return m.hexdigest() \ No newline at end of file + return m.hexdigest() diff --git a/psptool/utils.py b/psptool/utils.py index 963245f..5ff2339 100644 --- a/psptool/utils.py +++ b/psptool/utils.py @@ -106,8 +106,8 @@ def get_bytes(self, offset: int = 0x0, size: int = None) -> bytes: size = self.buffer_size if size is None else size return bytes(self[offset:offset + size]) - def set_bytes(self, address: int, size: int, value): - assert len(value) == size, "set_bytes: value length does not match size argument" + def set_bytes(self, address: int, value: bytes): + size = len(value) self[address:address + size] = value def get_chunks(self, size: int, offset: int = 0): diff --git a/tests/integration/test_replace.py b/tests/integration/test_replace.py new file mode 100644 index 0000000..f2e1c2c --- /dev/null +++ b/tests/integration/test_replace.py @@ -0,0 +1,169 @@ +# PSPTool - Display, extract and manipulate PSP firmware inside UEFI images +# Copyright (C) 2026 contributors +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import contextlib +import io +import os +import subprocess +import sys +import tempfile +import unittest + +from psptool import PSPTool +from psptool.header_file import HeaderFile + +dirname = os.path.dirname(__file__) +rom_fixtures_path = os.path.join(dirname, "fixtures/roms") + + +class TestReplace(unittest.TestCase): + # Cache PSPTool objects across tests in a class member + cached_pts = {} + + def fixture_roms(self): + for subdir, dirs, files in os.walk(rom_fixtures_path): + for file in files: + if file[0] == ".": + continue + filename = os.path.join(subdir, file) + yield filename + + def pt_from_file(self, filename) -> PSPTool: + if filename not in self.__class__.cached_pts.keys(): + with io.StringIO() as stderr_buf: + with contextlib.redirect_stderr(stderr_buf): + self.cached_pts[filename] = PSPTool.from_file(filename) + warnings = stderr_buf.getvalue().split("\n") + + return self.__class__.cached_pts[filename] + + def find_substitutable_entry(self, pt: PSPTool): + """Find the first unsigned HeaderFile entry that can be substituted without re-signing.""" + for r_i, rom in enumerate(pt.blob.roms): + for d_i, d in enumerate(rom.directories): + for f_i, file in enumerate(d.files): + if type(file) is HeaderFile and not getattr( + file, "is_signed", False + ): + return r_i, d_i, f_i, file + return None + + def test_extract_and_substitute_cli(self): + """Extract a file via CLI (-X) and reinsert/substitute it back (-R). + The newly created ROM should be byte-for-byte identical to the original fixture. + """ + repo_root = os.path.abspath(os.path.join(dirname, "../..")) + env = dict(os.environ) + env["PYTHONPATH"] = repo_root + ( + os.pathsep + env["PYTHONPATH"] if "PYTHONPATH" in env else "" + ) + + for filename in self.fixture_roms(): + with self.subTest(filename=filename): + try: + pt = self.pt_from_file(filename) + except Exception as e: + self.skipTest(f"Failed to parse ROM: {e}") + + target = self.find_substitutable_entry(pt) + if not target: + self.skipTest("No substitutable unsigned HeaderFile found in ROM") + + r_i, d_i, f_i, file = target + + with tempfile.TemporaryDirectory() as td: + extracted_path = os.path.join(td, "extracted.bin") + new_rom_path = os.path.join(td, "new_rom.bin") + + # 1. Extract file via CLI (-X) + cmd_extract = [ + sys.executable, + "-m", + "psptool", + filename, + "-X", + "-r", + str(r_i), + "-d", + str(d_i), + "-e", + str(f_i), + "-o", + extracted_path, + ] + res_extract = subprocess.run( + cmd_extract, capture_output=True, text=True, env=env + ) + self.assertEqual( + res_extract.returncode, + 0, + f"CLI extract (-X) failed on {os.path.basename(filename)}:\n" + f"stdout: {res_extract.stdout}\nstderr: {res_extract.stderr}", + ) + self.assertTrue(os.path.isfile(extracted_path)) + self.assertGreater(os.path.getsize(extracted_path), 0) + + # 2. Reinsert / substitute file back via CLI (-R) + cmd_replace = [ + sys.executable, + "-m", + "psptool", + filename, + "-R", + "-r", + str(r_i), + "-d", + str(d_i), + "-e", + str(f_i), + "-s", + extracted_path, + "-o", + new_rom_path, + ] + res_replace = subprocess.run( + cmd_replace, capture_output=True, text=True, env=env + ) + self.assertEqual( + res_replace.returncode, + 0, + f"CLI replace (-R) failed on {os.path.basename(filename)}:\n" + f"stdout: {res_replace.stdout}\nstderr: {res_replace.stderr}", + ) + self.assertTrue(os.path.isfile(new_rom_path)) + + # 3. Verify original fixture and newly created ROM are identical + with open(filename, "rb") as f_orig, open( + new_rom_path, "rb" + ) as f_new: + orig_bytes = f_orig.read() + new_bytes = f_new.read() + + self.assertEqual( + len(orig_bytes), + len(new_bytes), + f"File size mismatch for {os.path.basename(filename)}: " + f"original is {len(orig_bytes)} bytes, new is {len(new_bytes)} bytes", + ) + self.assertEqual( + orig_bytes, + new_bytes, + f"Content mismatch: newly created ROM differs from original fixture {os.path.basename(filename)}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/synthetic_rom.py b/tests/unit/synthetic_rom.py index f9a1feb..f5cb679 100644 --- a/tests/unit/synthetic_rom.py +++ b/tests/unit/synthetic_rom.py @@ -63,7 +63,7 @@ def _build_psp_directory(entry_offset: int, entry_size: int) -> bytes: 0x00, # subprogram 0x0000, # flags entry_size, # size - entry_offset, # offset (raw — directory addr_mode 1 returns this as-is) + entry_offset, # offset (raw — directory relative to BIOS address_mode returns this as-is) 0x00000000, # rsv0 )