From 78c14e7c97801fc516287e9eb608781789ce704c Mon Sep 17 00:00:00 2001
From: Elliot Nunn <elliotnunn@me.com>
Date: Sat, 25 May 2019 10:02:41 +0800
Subject: [PATCH] Add bootinfo rebuilding

---
 tbxi/bootinfo_build.py | 66 +++++++++++++++++++++++++++++++++++++++++-
 tbxi/bootinfo_dump.py  |  5 ++--
 tbxi/dispatcher.py     |  2 +-
 3 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/tbxi/bootinfo_build.py b/tbxi/bootinfo_build.py
index 92c5275..689e00a 100644
--- a/tbxi/bootinfo_build.py
+++ b/tbxi/bootinfo_build.py
@@ -1,5 +1,69 @@
+from os import path
+import re
+import zlib
+
 from . import dispatcher
 
 
+def append_checksum(binary):
+    cksum = ('\r\\ h# %08X' % zlib.adler32(binary)).encode('ascii')
+    binary.extend(cksum)
+
+
 def build(src):
-    raise dispatcher.WrongFormat
+    try:
+        with open(path.join(src, 'Bootscript'), 'rb') as f:
+            booter = bytearray(f.read().replace(b'\n', b'\r'))
+    except (NotADirectoryError, FileNotFoundError):
+        raise dispatcher.WrongFormat
+
+    constants = dict()
+    constant_spans = dict()
+    for m in re.finditer(rb'h#\s+([A-Fa-f0-9]+)\s+constant\s+([-\w]+)', booter):
+        key = m.group(2).decode('ascii')
+        val = int(m.group(1), 16)
+        constants[key] = val
+        constant_spans[key] = m.span(1)
+
+    booter.append(4) # EOT
+
+    for f in ['elf-offset', 'elf-size']:
+        assert f in constants
+
+    if 'elf-offset' in constants:
+        # special case: pad according to residual info in script
+        booter.extend(b'\0' * (constants['elf-offset'] - len(booter)))
+
+        constants['elf-offset'] = len(booter)
+        booter.extend(dispatcher.build_path(path.join(src, 'MacOS.elf')))
+        constants['elf-size'] = len(booter) - constants['elf-offset']
+
+    if 'lzss-offset' in constants:
+        base = 'lzss'
+    else:
+        base = 'parcels'
+
+    if base + '-offset' in constants:
+        constants[base + '-offset'] = len(booter)
+        for attempt in ['MacROM', 'Parcels']:
+            try:
+                booter.extend(dispatcher.build_path(path.join(src, attempt)))
+            except:
+                pass
+            else:
+                break
+        else:
+            raise FileNotFoundError
+
+        constants[base + '-size'] = len(booter) - constants[base + '-offset']
+
+    constants['info-size'] = len(booter)
+
+    for key, (start, stop) in reversed(sorted(constant_spans.items())):
+        insert = ('%X' % constants[key]).zfill(stop - start).encode('ascii')
+        assert start + len(insert) == stop
+        booter[start:stop] = insert
+
+    append_checksum(booter)
+
+    return bytes(booter)
diff --git a/tbxi/bootinfo_dump.py b/tbxi/bootinfo_dump.py
index 0bbdd67..f3a4feb 100644
--- a/tbxi/bootinfo_dump.py
+++ b/tbxi/bootinfo_dump.py
@@ -24,8 +24,9 @@ def dump(binary, dest_dir):
         val = int(m.group(1), 16)
         constants[key] = val
 
-        for i in range(*m.span(1)):
-            chrp_boot_zeroed[i:i+1] = b'0'
+        if key != 'elf-offset':
+            for i in range(*m.span(1)):
+                chrp_boot_zeroed[i:i+1] = b'0'
 
     with open(path.join(dest_dir, 'Bootscript'), 'wb') as f:
         f.write(chrp_boot_zeroed)
diff --git a/tbxi/dispatcher.py b/tbxi/dispatcher.py
index b6f6fce..d36be82 100644
--- a/tbxi/dispatcher.py
+++ b/tbxi/dispatcher.py
@@ -51,7 +51,7 @@ def build_path(p):
     # try building the file from a directory
     for np in [p + '.src', p]:
         try:
-            data = build_dir(p)
+            data = build_dir(np)
         except WrongFormat:
             pass
         else: