[SMD] Streets of Rage 2 tileset (de)compressor (extractor, unpacker)

sor2_comp.py

#!/usr/bin/env python3

import io
import struct
from pathlib import Path

def decompress(inpt, artAddress):
    inpt.seek(artAddress)
    output = io.BytesIO()

    size = struct.unpack("<H", inpt.read(2))[0]

    prev = 0
    for _ in range(size):
        pos = inpt.tell() - artAddress
        if pos >= size:
            break
        combyte = inpt.read(1)[0]
        if (combyte & 0x80) == 0x80:
            d = (combyte & 0x60) // 0x20 + 4
            offset = (combyte & 0x1F) << 8
            offset = offset | inpt.read(1)[0]

            output.seek(-offset, io.SEEK_CUR)
            prev = output.tell() - 1
            for _ in range(d):
                prev += 1
                output.seek(prev, io.SEEK_SET)

                k = output.read(1)
                output.seek(0, io.SEEK_END)
                output.write(k)
        elif (combyte & 0x60) == 0x60:
            d = combyte & 0x1F
            for _ in range(d):
                prev += 1
                output.seek(prev, io.SEEK_SET)

                k = output.read(1)
                output.seek(0, io.SEEK_END)
                output.write(k)
        elif (combyte & 0x40) == 0x40:
            d = combyte & 0x1F # combyte ^ 0x40
            if (combyte & 0x10) == 0x10:
                d = (d ^ 0x10) << 8
                d =  d | inpt.read(1)[0]
            d += 4
            k = inpt.read(1)
            output.write(k * d)
        elif (combyte & 0x20) == 0x20:
            d = (combyte & 0x1F) << 8
            d = d | inpt.read(1)[0]
            output.write(inpt.read(d))
        else: #if (combyte & 0xE0) == 0x00:
            output.write(inpt.read(combyte))

    inpt.close()
    output.seek(0, io.SEEK_SET)
    return output

def compress(bi, ultra = False):
    def write_raw(b, bi, begin, end):
        """ ABCD...WXYZ """
        chunk = bi[begin:end]
        size = len(chunk)
        if size == 0:
            return
        if size <= 0x1F:
            b.append(size)
            b.extend(chunk)
        elif size <= 0x1FFF:
            b.append(0x20 | (size >> 8))
            b.append(size & 0xFF)
            b.extend(chunk)
        else:
            q = size // 0x1FFF
            offset = begin
            for _ in range(q):
                b.extend(b"\x3F\xFF")
                b.extend(bi[offset:offset+0x1FFF])
                offset += 0x1FFF
            r = size % 0x1FFF
            if r != 0:
                write_raw(b, bi, offset, offset + r)

    def write_simple_run(b, byte, size):
        """ BBBB...BBBB """
        if size < 4:
            print("<<< Something wrong >>>")
            return
        size -= 4
        if size <= 0xF:
            b.append(0x40 | size)
        elif size <= 0xFFF:
            b.append(0x50 | (size >> 8))
            b.append(size & 0xFF)
        else:
            size += 4
            q = size // (0xFFF + 4)
            b.extend((b"\x5F\xFF" + bytes([byte])) * q)
            r = size % (0xFFF + 4)
            if r != 0:
                write_simple_run(b, byte, r)
            return
        b.append(byte)
    
    def write_run(b, size, offset):
        """ ABAB...ABAB """
        if size < 4:
            print("<<< Something wrong >>>")
            return
        # 0x80
        b.append(0x80 | ((min(7, size) - 4) << 5) | (offset >> 8))
        b.append(offset & 0xFF)

        # 0x60
        if size > 7:
            size -= 7
            b.extend(b"\x7F" * (size // 0x1F))
            r = size % 0x1F
            if r != 0:
                b.append(0x60 | r)

    size_bi = len(bi)
    b1 = bytearray(b"??")
    b2 = bytearray(b"??")
    b = b1
    last_result_0 = 0
    save_point = [0, 0, 0, 0, 0, 0, 0]
    decision = 0
    skip = 0

    i = 0
    needw = 0
    while i <= size_bi - 4:
        start = max(i - 0x1FFF, 0)

        count_s = 0
        while i < size_bi - count_s \
            and bi[i+count_s] == bi[i]:
            count_s += 1

        count_r = 0
        max_i = -1
        while True:
            start = bi.find(bi[i:i+4], start, i + 2)
            if start != -1:
                count = 4
                while i < size_bi - count \
                    and bi[start+count] == bi[i+count]:
                    count += 1
                if count_r < count:
                    count_r = count
                    max_i = start
                start += 1
            else:
                break
        start = max_i

        if start != -1:
            if needw > 0:
                write_raw(b, bi, i - needw, i)
                needw = 0

            penalty = 0
            if count_s > 0xFFF + 4:
                penalty = 1
                penalty += ((count_s - (0xFFF + 4)) // (0xFFF + 4)) * 3
                r = (count_s - (0xFFF + 4)) % (0xFFF + 4)
                if r != 0:
                    if r < 4:
                        count_s -= r
                    elif r <= 0xF + 4:
                        penalty += 2
                    else:
                        penalty += 3
            elif count_s > 0xF + 4:
                penalty = 1
            penalty2 = 0
            if count_r > 7:
                penalty2 = (count_r - 7) // 0x1F
                if (count_r - 7) % 0x1F != 0:
                    penalty2 += 1

            if ultra: # TODO: Last check
                if count_s - penalty == count_r - penalty2 and not skip:
                    skip = 1
                    decision ^= 1
                    if decision == 1:
                        last_result_0 = len(b1)
                        save_point[1:] = start, count_s, penalty, count_r, penalty2, i
                        i = save_point[0]
                        b = b2
                        continue
                    else:
                        if last_result_0 < len(b2):
                            start, count_s, penalty, count_r, penalty2, i = save_point[1:]
                            b2 = b1[:]
                        else:
                            b1 = b2[:]
                        save_point[0] = i
                        b = b1
                if skip:
                    skip -= 1

            if (count_s - penalty + decision > count_r - penalty2):
                write_simple_run(b, bi[i], count_s)
                i += count_s - 1
            else:
                write_run(b, count_r, i - start)
                i += count_r - 1
        elif count_s >= 4:
            if needw > 0:
                write_raw(b, bi, i - needw, i)
                needw = 0

            write_simple_run(b, bi[i], count_s)
            i += count_s - 1
        else:
            needw += 1

        i += 1

    write_raw(b, bi, i - needw, size_bi)

    size = len(b)
    b[0] = size & 0xFF
    b[1] = size >> 8
    return bytes(b)

if __name__ == '__main__':
    import ast
    import argparse

    parser = argparse.ArgumentParser(description='Streets of Rage 2 [MD] tileset (de)compressor')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-d', '--decompress', nargs='+', metavar=('INPUT', 'OFFSET'))
    group.add_argument('-c', '--compress', metavar='INPUT')
    parser.add_argument('-o', '--output', required=True)
    parser.add_argument('-u', '--ultra', action='store_true', help='Better compression')
    parser.add_argument('-r', '--replace', nargs='+', metavar=('INPUT', 'OFFSET'), help='Compress and replace in file')
    parser.add_argument('-f', '--force', action='store_true', help='Force replace (no padding)')
    args = parser.parse_args()

    if args.decompress:
        offset = 0
        if len(args.decompress) > 1:
            offset = ast.literal_eval(args.decompress[1])
        fc = Path(args.decompress[0]).open("rb")
        fc.seek(offset, io.SEEK_SET)
        packed_size = struct.unpack("<H", fc.read(2))[0]
        fc.seek(0, io.SEEK_SET)
        with decompress(fc, offset) as fd:
            bd = fd.read()
            Path(args.output).write_bytes(bd)
        size = len(bd)
        print(f"Packed Size: {packed_size}")
        print(f"       Size: {size}")
    elif args.compress:
        bd = Path(args.compress).read_bytes()
        bc = compress(bd, args.ultra)
        size = len(bd)
        packed_size = len(bc)
        if args.replace:
            offset = 0
            b = bytearray(Path(args.replace[0]).read_bytes())
            if len(args.replace) > 1:
                offset = ast.literal_eval(args.replace[1])
            if not args.force:
                packed_size_old = struct.unpack("<H", b[offset:offset+2])[0]
                if packed_size > packed_size_old:
                    print("Too big!")
                else:
                    pad = b"\x00" * (packed_size_old - len(bc))
                    b[offset:offset+packed_size_old] = bc + pad
                    Path(args.output).write_bytes(b)
                print(f"Old Packed Size: {packed_size_old}")
            else:
                b[offset:offset+packed_size] = bc
                Path(args.output).write_bytes(b)
            print(f"New Packed Size: {packed_size}")
            print(f"           Size: {size}")
        else:
            Path(args.output).write_bytes(bc)
            print(f"Packed Size: {packed_size}")
            print(f"       Size: {size}")