# v2 pbzx stream handler
# My personal writeup on the differences here: https://gist.github.com/pudquick/29fcfe09c326a9b96cf5
#
# Pure python reimplementation of .cpio.xz content extraction from pbzx file payload originally here:
# http://www.tonymacx86.com/general-help/135458-pbzx-stream-parser.html
#
# Cleaned up C version (as the basis for my code) here, thanks to Pepijn Bruienne / @bruienne
# https://gist.github.com/bruienne/029494bbcfb358098b41

import struct
import sys


def seekread(f, offset = None, length = 0, relative = True):
    if offset is not None:
        # offset provided, let's seek
        f.seek(offset, [0, 1, 2][relative])
    if length:
        return f.read(length)
    return None


def parse_pbzx(pbzx_path):
    section = 0
    xar_out_path = '%s.part%02d.cpio.xz' % (pbzx_path, section)
    with open(pbzx_path, 'rb') as f:
        # pbzx = f.read()
        # f.close()
        magic = seekread(f, length = 4)
        if magic != 'pbzx':
            raise RuntimeError("Error: Not a pbzx file")
        # Read 8 bytes for initial flags
        flags = seekread(f, length = 8)
        # Interpret the flags as a 64-bit big-endian unsigned int
        flags = struct.unpack('>Q', flags)[0]
        while flags & (1 << 24):
            with open(xar_out_path, 'wb') as xar_f:
                xar_f.seek(0, os.SEEK_END)
                # Read in more flags
                flags = seekread(f, length = 8)
                flags = struct.unpack('>Q', flags)[0]
                # Read in length
                f_length = seekread(f, length = 8)
                f_length = struct.unpack('>Q', f_length)[0]
                xzmagic = seekread(f, length = 6)
                if xzmagic != '\xfd7zXZ\x00':
                    # This isn't xz content, this is actually _raw decompressed cpio_ chunk of 16MB in size...
                    # Let's back up ...
                    seekread(f, offset = -6, length = 0)
                    # ... and split it out ...
                    f_content = seekread(f, length = f_length)
                    section += 1
                    decomp_out = '%s.part%02d.cpio' % (pbzx_path, section)
                    with open(decomp_out, 'wb') as g:
                        g.write(f_content)
                    # Now to start the next section, which should hopefully be .xz (we'll just assume it is ...)
                    section += 1
                    xar_out_path = '%s.part%02d.cpio.xz' % (pbzx_path, section)
                else:
                    f_length -= 6
                    # This part needs buffering
                    f_content = seekread(f, length = f_length)
                    tail = seekread(f, offset = -2, length = 2)
                    xar_f.write(xzmagic)
                    xar_f.write(f_content)
                    if tail != 'YZ':
                        raise RuntimeError("Error: Footer is not xar file footer")


def main():
    parse_pbzx(sys.argv[1])
    print("Now xz decompress the .xz chunks, then 'cat' them all together in order into a single new.cpio file")


if __name__ == '__main__':
    main()