Fix replace

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')
-```
+```

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:

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}')

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('<I', self.count)
         self.update_checksum()
 
-    # 00b: x86 Physical address
-    # 01b: Offset from start of the BIOS (flash offset)
-    # 10b: Offset from start of directory header
-    # 11b: Offset from start of partition
     @property
-    def address_mode(self):
+    def address_mode(self) -> AddressMode:
         info = struct.unpack('<L', self.additional_info)[0]
         version = (info >> 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()

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,23 +84,23 @@ def type(self):
 
     @type.setter
     def type(self, value):
-        self.set_bytes(0, 1, struct.pack('<B', value))
+        self.set_bytes(0, struct.pack('<B', value))
 
     @property
     def subprogram(self):
         return struct.unpack('<B', self[1:2])[0]
 
     @subprogram.setter
     def subprogram(self, value):
-        self.set_bytes(1, 1, struct.pack('<B', value))
+        self.set_bytes(1, struct.pack('<B', value))
 
     @property
     def flags(self):
         return struct.unpack('<H', self[2:4])[0]
 
     @flags.setter
     def flags(self, value):
-        self.set_bytes(2, 2, struct.pack('<H', value))
+        self.set_bytes(2, struct.pack('<H', value))
 
     @property
     def instance(self):
@@ -112,31 +116,27 @@ def size(self):
 
     @size.setter
     def size(self, value):
-        self.set_bytes(4, 4, struct.pack('<I', value))
+        self.set_bytes(4, struct.pack('<I', value))
 
     @property
     def offset(self):
         return struct.unpack('<I', self[8:12])[0]
 
     @offset.setter
     def offset(self, value):
-        self.set_bytes(8, 4, struct.pack('<I', value))
+        self.set_bytes(8, struct.pack('<I', value))
 
     @property
     def rsv0(self):
         return struct.unpack('<I', self[12:16])[0]
 
     @rsv0.setter
     def rsv0(self, value):
-        self.set_bytes(12, 4, struct.pack('<I', value))
+        self.set_bytes(12, struct.pack('<I', value))
 
-    # 00b: x86 Physical address
-    # 01b: Offset from start of the BIOS (flash offset)
-    # 10b: Offset from start of directory header
-    # 11b: Offset from start of partition
     @property
-    def address_mode(self):
-        return (self.rsv0 >> 30) & 3
+    def address_mode(self) -> AddressMode:
+        return AddressMode((self.rsv0 >> 30) & 3)
 
 class BiosDirectoryEntry(DirectoryEntry):
     ENTRY_SIZE = 4 * 6
@@ -150,23 +150,23 @@ def destination(self):
 
     @destination.setter
     def destination(self, value):
-        self.set_bytes(16, 8, struct.pack('<Q', value))
+        self.set_bytes(16, struct.pack('<Q', value))
 
     @property
     def region_type(self):
         return struct.unpack('<B', self[1:2])[0]
 
     @region_type.setter
     def region_type(self, value):
-        self.set_bytes(1, 1, struct.pack('<B', value))
+        self.set_bytes(1, struct.pack('<B', value))
 
     @property
     def flags(self):
         return struct.unpack('<H', self[2:4])[0]
 
     @flags.setter
     def flags(self, value):
-        self.set_bytes(2, 2, struct.pack('<H', value))
+        self.set_bytes(2, struct.pack('<H', value))
 
     @property
     def subprogram(self):

psptool/file.py

@@ -386,7 +386,7 @@ def move_buffer(self, new_address, size):
         current_address = self.get_address()
         move_offset = new_address - current_address
         self.buffer_offset += move_offset
-        self.buffer_size = int(ceil(size / self.ENTRY_ALIGNMENT)) * self.ENTRY_ALIGNMENT
+        self.buffer_size = size
 
         # update all directories' headers that point to this entry
         for directory in self.references: