1import fnmatch
2import struct
3
4import KickRom
5import RemusFile
6import amitools.util.DataDir as DataDir
7from amitools.binfmt.BinImage import *
8
9
10class RomSplitter:
11  def __init__(self, split_data_path=None):
12    # get data file path
13    if split_data_path is None:
14      split_data_path = DataDir.ensure_data_sub_dir("splitdata")
15    # setup remus file set
16    self.rfs = RemusFile.RemusFileSet()
17    self.rfs.load(split_data_path)
18    # state
19    self.chk_sum = None
20    self.rom_data = None
21    self.remus_rom = None
22
23  def list_roms(self, out, query=None, show_entries=False):
24    roms = self.rfs.get_roms()
25    for rom in roms:
26      if query is None or fnmatch.fnmatch(rom.name, query):
27        self.remus_rom = rom
28        self.print_rom(out, show_entries)
29
30  def find_rom(self, rom_path):
31    """load ROM and try to find a matching dat file.
32       Returns True if ROM was matched"""
33    self.rom_data = KickRom.Loader.load(rom_path)
34    # get check sum
35    kh = KickRom.KickRomAccess(self.rom_data)
36    if kh.is_kick_rom():
37      self.chk_sum = kh.read_check_sum()
38    else:
39      self.chk_sum = kh.calc_check_sum()
40    # search rom in Remus data base
41    self.remus_rom = self.rfs.find_rom(self.rom_data, self.chk_sum)
42    return self.remus_rom
43
44  def print_rom(self, out, show_entries=False):
45    rom = self.remus_rom
46    out("@%08x  +%08x  sum=%08x  sum_off=%08x  %s" % \
47        (rom.base_addr, rom.size, rom.chk_sum, rom.sum_off, rom.name))
48    if show_entries:
49      for e in rom.entries:
50        self.print_entry(out, e)
51
52  def print_entry(self, out, entry):
53    out("  @%06x  +%06x  =%06x  relocs=#%5d  %s" % \
54        (entry.offset, entry.size, entry.offset+entry.size,
55         len(entry.relocs), entry.name))
56
57  def print_entries(self, out, entries):
58    for e in entries:
59      self.print_entry(out, e)
60
61  def get_all_entries(self):
62    return self.remus_rom.entries
63
64  def query_entries(self, query_str):
65    res = []
66    for e in self.remus_rom.entries:
67      if fnmatch.fnmatch(e.name, query_str):
68        res.append(e)
69    return res
70
71  def extract_entry(self, entry):
72    """return data, relocs"""
73    data = self.rom_data[entry.offset:entry.offset+entry.size]
74    relocs = entry.relocs
75    entry_addr = self.remus_rom.base_addr + entry.offset
76    data = self._clean_relocs(data, relocs, entry_addr)
77    return data, relocs
78
79  def _clean_relocs(self, data, relocs, base_addr):
80    if type(data) is not bytearray:
81      data = bytearray(data)
82    for off in relocs:
83      addr = struct.unpack_from(">I", data, off)[0]
84      if addr < base_addr:
85        raise ValueError("Invalid relocatable address: %08x" % addr)
86      addr -= base_addr
87      struct.pack_into(">I", data, off, addr)
88    return data
89
90  def extract_bin_img(self, entry):
91    data, relocs = self.extract_entry(entry)
92    # create a bin image
93    bi = BinImage(BIN_IMAGE_TYPE_HUNK)
94    seg = Segment(SEGMENT_TYPE_CODE, len(data), data)
95    bi.add_segment(seg)
96    # create reloc for target segment
97    rl = Relocations(seg)
98    # add offsets
99    for o in relocs:
100      r = Reloc(o)
101      rl.add_reloc(r)
102    seg.add_reloc(seg, rl)
103    # return final binary image
104    return bi
105
106  def write_index_file(self, idx_path):
107    with open(idx_path, "w") as fh:
108      for e in self.remus_rom.entries:
109        fh.write(e.name + "\n")
110