mirror of
https://github.com/BrianPugh/game-and-watch-patch.git
synced 2025-12-16 07:16:26 +01:00
300 lines
9.6 KiB
Python
300 lines
9.6 KiB
Python
from .compression import lzma_compress
|
|
|
|
def twos_compliment(value, bits):
|
|
if value >= 0:
|
|
return value
|
|
else:
|
|
# Two's Compliment
|
|
return (1 << bits) + value
|
|
|
|
|
|
class FirmwarePatchMixin:
|
|
""" Patch commands that apply to a single firmware instance.
|
|
"""
|
|
|
|
def replace(self, offset : int, data, size=None) -> int:
|
|
"""
|
|
"""
|
|
|
|
if offset >= len(self):
|
|
raise IndexError(f"Patch offset {offset} exceeds firmware length {len(self)}")
|
|
|
|
if offset >= self.STOCK_ROM_END:
|
|
raise IndexError(f"Patch offset {offset} exceeds stock firmware region {self.STOCK_ROM_END}")
|
|
|
|
n_bytes_patched = 0
|
|
|
|
if isinstance(data, bytes):
|
|
# Write the bytes at that address as is.
|
|
self[offset:offset + len(data)] = data
|
|
n_bytes_patched = len(data)
|
|
elif isinstance(data, str):
|
|
if size:
|
|
raise ValueError("Don't specify size when providing a symbol name.")
|
|
data = self.address(data)
|
|
self[offset:offset+4] = data.to_bytes(4, 'little')
|
|
n_bytes_patched = 4
|
|
elif isinstance(data, int):
|
|
# must be 1, 2, or 4 bytes
|
|
if size is None:
|
|
raise ValueError("Must specify \"size\" when providing int data")
|
|
if size not in (1, 2, 4):
|
|
raise ValueError(f"Size must be one of {1, 2, 4}; got {size}")
|
|
self[offset:offset+size] = data.to_bytes(size, 'little')
|
|
n_bytes_patched = size
|
|
else:
|
|
raise ValueError(f"Don't know how to parse data type \"{data}\"")
|
|
|
|
return n_bytes_patched
|
|
|
|
def relative(self, offset, data, size=None) -> int:
|
|
"""
|
|
data
|
|
If str, looks up a function
|
|
"""
|
|
src = self.FLASH_BASE + offset
|
|
|
|
if isinstance(data, str):
|
|
if size:
|
|
raise ValueError("Don't specify size when providing a symbol name.")
|
|
dst = self.address(data)
|
|
elif isinstance(data, int):
|
|
# must be 1, 2, or 4 bytes
|
|
if size is None:
|
|
raise ValueError("Must specify \"size\" when providing int data.")
|
|
if data < self.FLASH_BASE:
|
|
data += self.FLASH_BASE
|
|
dst = data
|
|
rel_distance = dst - src
|
|
if rel_distance < 0:
|
|
rel_distance += 0x1_0000_0000
|
|
|
|
print(f"Computed relative distance 0x{rel_distance:08X}")
|
|
|
|
return self.replace(offset, rel_distance, size=4)
|
|
|
|
|
|
def b(self, offset : int, data : int) -> int:
|
|
""" Unconditional branch inserted at ``offset`` to offset ``data``
|
|
|
|
data should be the aboolute offset into the firmware (i.e. NOT in the
|
|
form 0x08XX_XXXX or 0x9XXX_XXXX)
|
|
|
|
2 byte command
|
|
"""
|
|
|
|
pc = offset + 4
|
|
jump = data - pc
|
|
|
|
if abs(jump) > (2 * (1<<10)):
|
|
# Max +-2KB jump
|
|
raise ValueError(f"Too large of a jump {jump} specified.")
|
|
|
|
jump >>= 1
|
|
jump = twos_compliment(jump, 11)
|
|
|
|
byte_0 = 0b1110_0000 | ((jump >> 8) & 0x7)
|
|
byte_1 = jump & 0xFF
|
|
|
|
self[offset + 0] = byte_1
|
|
self[offset + 1] = byte_0
|
|
|
|
return 2
|
|
|
|
|
|
def bl(self, offset : int, data : str) -> int:
|
|
""" Replace a branch-link statement to a branch to one of our functions
|
|
|
|
4 byte command.
|
|
"""
|
|
|
|
dst_address = self.address(data)
|
|
|
|
pc = self.FLASH_BASE + offset + 4
|
|
|
|
jump = dst_address - pc
|
|
|
|
if abs(jump) > (4 * (1 << 20)):
|
|
# Max +-4MB jump
|
|
raise ValueError(f"Too large of a jump {jump} specified.")
|
|
|
|
# Where H=0
|
|
offset_stage_1 = jump >> 12
|
|
offset_stage_1 = twos_compliment(offset_stage_1, 11)
|
|
|
|
stage_1_byte_0 = 0b1111_0000 | ((offset_stage_1 >> 8) & 0x7)
|
|
stage_1_byte_1 = offset_stage_1 & 0xFF
|
|
|
|
# Where H=1
|
|
offset_stage_2 = (jump - (offset_stage_1 << 12)) >> 1
|
|
if offset_stage_2 >> 11:
|
|
raise ValueError(f"bl jump 0x{jump:08X} too large!")
|
|
|
|
stage_2_byte_0 = 0b1111_1000 | ((offset_stage_2 >> 8) & 0x7)
|
|
stage_2_byte_1 = offset_stage_2 & 0xFF
|
|
|
|
# Store the instructions in little endian order
|
|
self[offset + 0] = stage_1_byte_1
|
|
self[offset + 1] = stage_1_byte_0
|
|
self[offset + 2] = stage_2_byte_1
|
|
self[offset + 3] = stage_2_byte_0
|
|
|
|
return 4
|
|
|
|
def bkpt(self, offset, size=2):
|
|
""" Insert software breakpoint(s) """
|
|
if size % 2:
|
|
raise ValueError("Size of breakpoints must be even.")
|
|
|
|
n_bkpts = size // 2
|
|
self[offset:offset+size] = b"\x00\xbe" * n_bkpts
|
|
return size
|
|
|
|
@property
|
|
def _ks(self):
|
|
try:
|
|
return self._ks_inst
|
|
except AttributeError:
|
|
from keystone import Ks, KS_ARCH_ARM, KS_MODE_THUMB
|
|
self._ks_inst = Ks(KS_ARCH_ARM, KS_MODE_THUMB)
|
|
return self._ks_inst
|
|
|
|
def asm(self, offset : int, data : str, size=None) -> int:
|
|
"""
|
|
Parameters
|
|
----------
|
|
data : str
|
|
Assembly instructions
|
|
"""
|
|
encoding, _ = self._ks.asm(data)
|
|
print(f" \"{data}\" -> {[hex(x) for x in encoding]}")
|
|
if size:
|
|
assert len(encoding) == size
|
|
for i, val in enumerate(encoding):
|
|
self[offset + i] = val
|
|
return len(encoding)
|
|
|
|
def nop(self, offset : int, data : int) -> int:
|
|
"""Insert N NOP operations (each 2 bytes long)."""
|
|
size = data * 2
|
|
self[offset:offset+size] = b"\x00\xbf" * data
|
|
return size
|
|
|
|
def _move_copy(self, offset : int, data : int, size : int, delete : bool) -> int:
|
|
""" Move from offset -> data """
|
|
|
|
if not isinstance(data, int):
|
|
raise ValueError(f"Data must be int, got {type(data)}")
|
|
|
|
old_start = offset
|
|
old_end = old_start + size
|
|
new_start = offset + data
|
|
new_end = new_start + size
|
|
print(f" moving {size} bytes from 0x{old_start:08X} to 0x{new_start:08X}")
|
|
self[new_start:new_end] = self[old_start:old_end]
|
|
|
|
# Erase old copy
|
|
if delete:
|
|
if data < 0:
|
|
if new_end > offset:
|
|
self.clear_range(new_end, old_end)
|
|
else:
|
|
self.clear_range(old_start, old_end)
|
|
else:
|
|
if new_start < old_end:
|
|
self.clear_range(old_start, new_start)
|
|
else:
|
|
self.clear_range(old_start, old_end)
|
|
|
|
for i in range(size):
|
|
self._lookup[self.FLASH_BASE + old_start + i] = self.FLASH_BASE + new_start + i
|
|
|
|
return size
|
|
|
|
def move(self, offset : int, data : int, size : int) -> int:
|
|
return self._move_copy(offset, data, size, True)
|
|
|
|
def copy(self, offset : int, data : int, size : int) -> int:
|
|
return self._move_copy(offset, data, size, False)
|
|
|
|
def add(self, offset : int, data : int, size : int = 4) -> int:
|
|
val = self.int(offset, size)
|
|
val += data
|
|
self[offset:offset+size] = val.to_bytes(size, "little")
|
|
|
|
return size
|
|
|
|
def shorten(self, data : int) -> int:
|
|
data = abs(data)
|
|
if data == 0:
|
|
return
|
|
|
|
self.ENC_LEN -= data
|
|
if self.ENC_LEN < 0:
|
|
self.ENC_LEN = 0
|
|
if len(self) == data:
|
|
self[:] = bytearray()
|
|
else:
|
|
self[:] = self[:-data]
|
|
|
|
return data
|
|
|
|
def compress(self, offset : int, size : int) -> int:
|
|
""" Apply in-place LZMA compression. """
|
|
data = self[offset:offset+size]
|
|
|
|
compressed_data = lzma_compress(data)
|
|
|
|
# Clear the original data
|
|
self.clear_range(offset, offset+size)
|
|
# Insert the compressed data
|
|
self[offset:offset+len(compressed_data)] = compressed_data
|
|
|
|
print(f" compressed {len(data)}->{len(compressed_data)} bytes (saves {len(data)-len(compressed_data)})")
|
|
|
|
return len(compressed_data)
|
|
|
|
def lookup(self, offsets):
|
|
size = 4
|
|
|
|
if not isinstance(offsets, list):
|
|
offsets = [offsets]
|
|
|
|
for offset in offsets:
|
|
val = self.int(offset, size)
|
|
try:
|
|
new_val = self._lookup[val]
|
|
except KeyError:
|
|
raise KeyError(f"0x{val:08X} at offset 0x{offset:08X}")
|
|
self[offset:offset+size] = new_val.to_bytes(size, "little")
|
|
#print(f" lookup 0x{val:08X}->0x{new_val:08x}")
|
|
|
|
|
|
class DevicePatchMixin:
|
|
def _move_copy(self, dst, dst_offset : int, src, src_offset : int, size : int, delete : bool) -> int:
|
|
dst[dst_offset:dst_offset+size] = src[src_offset:src_offset+size]
|
|
if delete:
|
|
src.clear_range(src_offset, src_offset + size)
|
|
|
|
for i in range(size):
|
|
self.lookup[src.FLASH_BASE + src_offset + i] = dst.FLASH_BASE + dst_offset + i
|
|
|
|
return size
|
|
|
|
def move(self, dst, dst_offset : int, src, src_offset : int, size : int) -> int:
|
|
return self._move_copy(dst, dst_offset, src, src_offset, size, True)
|
|
|
|
def copy(self, dst, dst_offset : int, src, src_offset : int, size : int) -> int:
|
|
return self._move_copy(dst, dst_offset, src, src_offset, size, False)
|
|
|
|
# Convenience methods for move and copy
|
|
def move_to_int(self, ext_offset:int, int_offset:int, size:int) -> int:
|
|
return self.move(self.internal, int_offset, self.external, ext_offset, size)
|
|
|
|
def copy_to_int(self, ext_offset:int, int_offset:int, size:int) -> int:
|
|
return self.copy(self.internal, int_offset, self.external, ext_offset, size)
|
|
|
|
def move_to_sram3(self, ext_offset:int, sram_offset:int, size:int) -> int:
|
|
return self.move(self.sram3, sram_offset, self.external, ext_offset, size)
|
|
|