onekpaq.py

#!/usr/bin/env python3

import argparse, os, re, shutil, stat, subprocess, sys, tempfile

def main():
    scriptdir = os.path.dirname(os.path.realpath(__file__))
    # TODO: auto-parameterize this?
    stubbase = "onekpaq_stub_lnx32.asm"

    parser = argparse.ArgumentParser()

    parser.add_argument('--onekpaq', default=scriptdir+"/onekpaq_encode",
        help="The onekpaq_encode binary to use")
    parser.add_argument('--nasm', default=os.getenv('NASM') or shutil.which('nasm'),
        help="The nasm binary to use")
    parser.add_argument('--stub', default=scriptdir+"/"+stubbase,
        help="The assembly decompressor stub to use")

    parser.add_argument('mode', choices=[1,2,3,4], type=int,
        help="Selects the compression and decompression mode: "\
            +"'1' is 'single block, small decoder', "\
            +"'2' is 'multi-block, small decoder', "\
            +"'3' is 'single block, fast decoder', "\
            +"'4' is 'multi-block, fast decoder'.")
    parser.add_argument('complexity', choices=[1,2,3], type=int,
        help="Compression complexity to use. Higher values "+\
             "mean higher compression, but also longer compression times.")
    parser.add_argument('input', nargs='+', help="Input block(s)")
    parser.add_argument('output', type=str, help="Output executable")

    args = parser.parse_args()

    with tempfile.NamedTemporaryFile(prefix="onekpaq") as tf:
        okpargs = [args.onekpaq, str(args.mode), str(args.complexity),
                   *args.input, tf.name]
        #print(okpargs)
        proc = subprocess.Popen(okpargs, stdout=subprocess.PIPE,
                                stderr=sys.stderr, cwd=scriptdir)
        (sout, serr_) = proc.communicate()
        sys.stderr.buffer.write(sout)
        sys.stderr.buffer.flush()
        if proc.returncode != 0:
            print("oneKpaq compressor failed.")
            return proc.returncode

        out = sout.decode('utf-8')
        match = re.search('P offset=([0-9]+) shift=([0-9]+)', out)
        assert match is not None
        offset, shift = (int(match.group(1)), int(match.group(2)))
        #print(offset, shift)

        #print(tf.name)
        nasmargs = [args.nasm, "-fbin", "-o", args.output,
                    args.stub, "-DPAYLOAD_OFF=%d" % offset,
                    "-DONEKPAQ_DECOMPRESSOR_MODE=%d" % args.mode,
                    "-DONEKPAQ_DECOMPRESSOR_SHIFT=%d" % shift,
                    "-DPAYLOAD_BIN=\"%s\"" % tf.name]
        #print(nasmargs)
        subprocess.run(nasmargs, check=True)
        os.chmod(args.output, os.stat(args.output).st_mode|stat.S_IEXEC)

        return 0


if __name__ == '__main__':
    rv = main()
    if rv is None: pass
    else:
        try: sys.exit(int(rv))
        except Exception: sys.exit(1)