1from __future__ import print_function
2
3import os
4import struct
5
6from RomAccess import RomAccess
7
8class KickRomAccess(RomAccess):
9
10  EXT_HEADER_SIZE = 0x10
11  FOOTER_SIZE = 0x18
12  ROMHDR_SIZE = 8
13  ROMHDR_256K = 0x11114ef9
14  ROMHDR_512K = 0x11144ef9
15
16  def __init__(self, rom_data):
17    RomAccess.__init__(self, rom_data)
18
19  def is_kick_rom(self):
20    return self.detect_kick_rom() == "ok"
21
22  def detect_kick_rom(self):
23    if not self.check_size():
24      return "size check failed"
25    if not self.check_header():
26      return "header check failed"
27    if not self.check_footer():
28      return "footer check failed"
29    if not self.check_rom_size_field():
30      return "rom size field mismatch"
31    if not self.verify_check_sum():
32      return "check sum mismatch"
33    return "ok"
34
35  def check_header(self):
36    # expect 0x1114 0x4ef9
37    val = self.read_long(0)
38    if self.kib == 512:
39      return val == self.ROMHDR_512K
40    elif self.kib == 256:
41      return val == self.ROMHDR_256K
42    else:
43      return False
44
45  def check_kickety_split(self):
46    # expect 0x1111 0x4ef9
47    if self.kib == 512:
48      val = self.read_long(0x40000)
49      return val == self.ROMHDR_256K
50    else:
51      return False
52
53  def check_footer(self):
54    # expect 0x0019 ... 0x001f
55    off = self.size - 14
56    num = 0x19
57    for i in xrange(7):
58      val = self.read_word(off)
59      if val != num:
60        return False
61      num += 1
62      off += 2
63    return True
64
65  def check_size(self):
66    # expect 256k or 512k ROM
67    if self.size % 1024 != 0:
68      return False
69    if self.kib not in (256, 512):
70      return False
71    return True
72
73  def check_rom_size_field(self):
74    return self.read_rom_size_field() == self.size
75
76  def check_magic_reset(self):
77    return self.read_word(0xd0) == 0x4e70
78
79  def calc_check_sum(self, skip_off=None):
80    """Check internal kickstart checksum and return True if is correct"""
81    chk_sum = 0
82    num_longs = self.size / 4
83    off = 0
84    max_u32 = 0xffffffff
85    for i in xrange(num_longs):
86      val = struct.unpack_from(">I", self.rom_data, off)[0]
87      if off != skip_off:
88        chk_sum += val
89      off += 4
90      if chk_sum > max_u32:
91        chk_sum = chk_sum & max_u32
92        chk_sum += 1
93    return max_u32 - chk_sum
94
95  def verify_check_sum(self):
96    chk_sum = self.calc_check_sum()
97    return chk_sum == 0
98
99  def read_check_sum(self):
100    sum_off = self.size - 0x18
101    return self.read_long(sum_off)
102
103  def recalc_check_sum(self):
104    sum_off = self.size - 0x18
105    return self.calc_check_sum(sum_off)
106
107  def write_check_sum(self):
108    cs = self.recalc_check_sum()
109    sum_off = self.size - 0x18
110    self.write_long(sum_off, cs)
111    return cs
112
113  def write_rom_size_field(self):
114    off = self.size - 0x14
115    self.write_long(off, self.size)
116
117  def write_header(self, jump_addr, kickety_split=False):
118    if kickety_split:
119      offset = 0x40000
120      hdr = self.ROMHDR_256K
121    elif self.kib == 256:
122      offset = 0
123      hdr = self.ROMHDR_256K
124    else:
125      offset = 0
126      hdr = self.ROMHDR_512K
127    self.write_long(offset, hdr)
128    self.write_long(offset+4, jump_addr)
129
130  def write_ext_header(self, jump_addr, rom_rev):
131    self.write_header(jump_addr)
132    self.write_word(8,0)
133    self.write_word(10,0xffff)
134    self.write_word(12,rom_rev[0])
135    self.write_word(14,rom_rev[1])
136
137  def write_ext_footer(self):
138    self.write_footer()
139    self.write_rom_size_field()
140    self.write_check_sum()
141
142  def write_footer(self):
143    off = self.size - 0x10
144    num = 0x18
145    for i in xrange(8):
146      self.write_word(off, num)
147      num += 1
148      off += 2
149
150  def write_rom_ver_rev(self, rom_rev):
151    """get (ver, rev) version info from ROM"""
152    return struct.pack_into(">HH", self.rom_data, 12, rom_rev[0], rom_rev[1])
153
154  def read_boot_pc(self):
155    """return PC for booting the ROM"""
156    return self.read_long(4)
157
158  def read_rom_ver_rev(self):
159    """get (ver, rev) version info from ROM"""
160    return struct.unpack_from(">HH", self.rom_data, 12)
161
162  def read_exec_ver_rev(self):
163    """get (ver, rev) version info from ROM"""
164    return struct.unpack_from(">HH", self.rom_data, 16)
165
166  def read_rom_size_field(self):
167    """return size of ROM stored in ROM itself"""
168    off = self.size - 0x14
169    return self.read_long(off)
170
171  def get_base_addr(self):
172    return self.read_boot_pc() & ~0xffff
173
174
175class Loader(object):
176  """Load kick rom images in different formats"""
177  @classmethod
178  def load(cls, kick_file, rom_key_file=None):
179    raw_img = None
180    rom_key = None
181    # read rom image
182    with open(kick_file, "rb") as fh:
183      raw_img = fh.read()
184    # coded rom?
185    need_key = False
186    if raw_img[:11] == 'AMIROMTYPE1':
187      rom_img = raw_img[11:]
188      need_key = True
189    else:
190      rom_img = raw_img
191    # decode rom
192    if need_key:
193      # read key file
194      if rom_key_file is None:
195        path = os.path.dirname(kick_file)
196        rom_key_file = os.path.join(path, "rom.key")
197      with open(rom_key_file, "rb") as fh:
198        rom_key = fh.read()
199      rom_img = cls._decode(rom_img, rom_key)
200    return rom_img
201
202  @classmethod
203  def _decode(cls, img, rom_key):
204    data = bytearray(img)
205    n = len(rom_key)
206    for i in xrange(len(data)):
207      off = i % n
208      data[i] = data[i] ^ ord(rom_key[off])
209    return bytes(data)
210
211
212# tiny test
213if __name__ == '__main__':
214  import sys
215  args = sys.argv
216  n = len(args)
217  if n > 1:
218    ks_file = args[1]
219  else:
220    ks_file = 'amiga-os-310-a500.rom'
221  print(ks_file)
222  ks = Loader.load(ks_file,'rom.key')
223  kh = KickRomAccess(ks)
224  print("is_kick_rom", kh.is_kick_rom())
225  print("detect_kick_rom", kh.detect_kick_rom())
226  print("pc=%08x" % kh.read_boot_pc())
227  print("ver,rev=", kh.read_rom_ver_rev())
228  print("ver,rev=", kh.read_exec_ver_rev())
229  print("size %08x == %08x" % (kh.read_rom_size_field(), len(ks)))
230  print("base %08x" % kh.get_base_addr())
231  print("get chk_sum=%08x" % kh.read_check_sum())
232  print("calc chk_sum=%08x" % kh.recalc_check_sum())
233#  with open("out.rom", "wb") as fh:
234#    fh.write(ks.get_data())
235
236