QB String Insert v2

#v2: +0x20 to EBOOT offsets to accomodate new program header
import struct
import zlib
import os
import shutil

INPUT_COL = 4
if os.path.isfile('shared.orig'):
    pass
elif os.path.isfile('shared.bin'):
    shutil.copy('shared.bin','shared.orig')
else:
    print('Missing shared.bin.')
    quit()

s1 = " .,'!?abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890~-&>():;/{%"
s2 = "　．，’！？ａｂｃｄｅｆｇｈｉｊｋｌｍｎｏｐｑｒｓｔｕｖｗｘｙｚＡＢＣＤＥＦＧＨＩＪＫＬＭＮＯＰＱＲＳＴＵＶＷＸＹＺ１２３４５６７８９０～‐＆＞（）：；／\x1F％"
translate_table = str.maketrans(s1, s2)    #For converting between ASCII and fullwidth SHIFT-JIS characters

inputdata = []
with open('s2_new.tsv', 'r', encoding='utf-8') as f:
    for line in f:
        line = line.rstrip('\r\n').split('\t')
        if line[0] != '':
            Expl_index = int(line[0])
            linenum = 0
        else:
            linenum += 1
        while '"' in line[INPUT_COL]:
            line[INPUT_COL] = line[INPUT_COL].replace('"', '“', 1)
            line[INPUT_COL] = line[INPUT_COL].replace('"', '”', 1)
        s = bytes(line[INPUT_COL].translate(translate_table), 'cp932') + b'\x00'
        inputdata.append([Expl_index,
                          linenum,
                          s])
inputdata = sorted(inputdata, key=lambda x: len(x[2]), reverse=True)
strings = []
count = 0
num_p = 0
ptrs = []
null_ptr = 0
ptrs.append(0)
Sent = bytearray()
for i, (Expl_index, linenum, new_s) in enumerate(inputdata):
    if new_s == b'\x00':                #Special Case
        inputdata[i].append(null_ptr)
        continue
    if new_s in Sent:
        ptr_tgt = Sent.find(new_s)
        if ptr_tgt in ptrs:             #string+ptr exist
            inputdata[i].append(ptrs.index(ptr_tgt))
        else:                           #string exist, no ptr
            inputdata[i].append(len(ptrs))
            ptrs.append(ptr_tgt)
    else:                               #neither string nor ptr exist
        inputdata[i].append(len(ptrs))
        ptrs.append(len(Sent))
        Sent += new_s
ptrs = [x + 0xC + 4 * len(ptrs) for x in ptrs]
ptrs[0] = 7     #For null pointer

Sent_len = 0xC + 4 * len(ptrs) + len(Sent)
Sent = b''.join((b'Sent',
                 struct.pack('<I', Sent_len),
                 struct.pack('<I', len(ptrs)),
                 b''.join([struct.pack('<I', x) for x in ptrs]),
                 Sent))
##with open('Sent.bin', 'wb') as f:
##    f.write(Sent)

Expl_index_old = 0
Expl = bytearray()
Expl_data = []
Expl_data_sub = []
inputdata = sorted(sorted(inputdata, key=lambda x: x[1]), key=lambda x: x[0])
for Expl_index, linenum, news, Sent_index in inputdata:
    if Expl_index == Expl_index_old:
        Expl_data_sub.append(Sent_index)
    else:
        Expl_data.append([len(Expl_data_sub)] + Expl_data_sub)
        Expl_data_sub = []
        Expl_index_old = Expl_index
        Expl_data_sub.append(Sent_index)
Expl_data.append([len(Expl_data_sub)] + Expl_data_sub)
Expl = bytearray()
ptrs = []
for nums in Expl_data:
    ptrs.append(len(Expl))
    Expl += b''.join([struct.pack('<H', x) for x in nums])
Expl = b''.join((b'Expl',
                 struct.pack('<I', 0xC + 4 * len(ptrs) + len(Expl)),
                 struct.pack('<I', len(Expl_data)),
                 b''.join([struct.pack('<I', x + 0xC + 4 * len(ptrs)) for x in ptrs]),
                 Expl))
##with open('Expl.bin', 'wb') as f:
##    f.write(Expl)
#-----End of Sent/Expl
inputdata = []
linenum = 0
INPUT_COL = 2
with open('s1.tsv', 'r', encoding='utf-8') as f:
    for line in f:
        line = line.rstrip('\r\n').split('\t')
        while '"' in line[INPUT_COL]:
            line[INPUT_COL] = line[INPUT_COL].replace('"', '“', 1)
            line[INPUT_COL] = line[INPUT_COL].replace('"', '”', 1)
        if line[INPUT_COL] != '':
            s = bytes(line[INPUT_COL].translate(translate_table), 'cp932') + b'\x00'
            inputdata.append([linenum, s])
            linenum += 1
if linenum != 772:
    print('Input error in s1. Must be exactly 772 strings.')
    quit()
inputdata = sorted(inputdata, key=lambda x: len(x[1]), reverse=True)
Strg = bytearray()
for i, (index, new_s) in enumerate(inputdata):
    if new_s in Strg:
        ptr_tgt = Strg.find(new_s)
        inputdata[i].append(ptr_tgt)
    else:
        inputdata[i].append(len(Strg))
        Strg += new_s
inputdata = sorted(inputdata, key=lambda x: x[0])
ptrs = [x[2] for x in inputdata]
Strg = b''.join((b'Strg',
                 struct.pack('<I', 0xC + 4 * linenum + len(Strg)),
                 struct.pack('<I', linenum),
                 b''.join([struct.pack('<I', x + 0xC + 4 * len(ptrs)) for x in ptrs]),
                 Strg))
##with open('Strg.bin', 'wb') as f:
##    f.write(Strg)

with open('shared.orig','rb') as f:
    f.seek(0x40800)
    filedata = bytearray(zlib.decompress(f.read(0x87F1)))
filedata = bytearray(filedata[:0xB520])
for file, ptr_pos in ((Strg, 0x50),
                      (Expl, 0x54),
                      (Sent, 0x58)):
    filedata[ptr_pos:ptr_pos+4] = struct.pack('<I', len(filedata))
    filedata += file
    if len(filedata) % 4 != 0:
        filedata += b'\x00' * (4 - (len(filedata) % 4))
##with open('40800.uncompressed', 'wb') as f:
##    f.write(filedata)
size = len(filedata)
filedata = zlib.compress(filedata, 9)
compressed_size = len(filedata)
if compressed_size > 0x800 * 26:
    print('Too big')
    quit()
tmpfilename = 'shared.tmp'
shutil.copy('shared.bin', tmpfilename)
with open(tmpfilename, 'rb') as f:
    with open('shared.bin', 'wb') as g:
        g.write(f.read(0x40800))
        g.write(filedata)
        g.write(b'\x00' * (0x4D800 - g.tell()))
        f.seek(0x4D800)
        g.write(f.read())
os.remove(tmpfilename)
with open('EBOOT.BIN', 'rb') as f:
    filedata = bytearray(f.read())
filedata[0x15DCC8:0x15DCCC] = struct.pack('<I', size)
filedata[0x15DCD0:0x15DCD4] = struct.pack('<I', compressed_size)
with open('EBOOT.BIN', 'wb') as f:
    f.write(filedata)