pbzx2zip.py

#!/usr/bin/env python3
'''convert payload to zip'''
import struct
import sys
import zipfile
import datetime
import io
import lzma

def get_chunk(payload, chunknum):
    '''reads a chunk, decompressing if necessary'''
    chunkheader = payload.read(16)
    if len(chunkheader) < 16:
        return
    uncompsize, compsize = struct.unpack('>QQ', chunkheader)
    print('chunk #%03d, uncompressed size %d, compressed size %d' % (
        chunknum, uncompsize, compsize))
    chunkdata = payload.read(compsize)
    if uncompsize != compsize:
        chunkdata = lzma.decompress(chunkdata, lzma.FORMAT_XZ)
    return io.BytesIO(chunkdata)

def main():
    '''entry point'''
    if len(sys.argv) < 3:
        print('Usage: %s In Out (PayloadPath)' % sys.argv[0])
        sys.exit(1)

    instr = sys.argv[1]
    outstr = sys.argv[2]

    if instr[-4:] == '.zip':
        infile = zipfile.ZipFile(instr, 'r')
        payload = infile.open(sys.argv[3] if len(sys.argv) > 3 else 'AssetData/payloadv2/payload')
    else:
        infile = payload = open(instr, 'rb')

    if payload.read(12)[:4] != b'pbzx':
        print('not pbzx')
        infile.close()
        return

    main.chunknum = 1
    main.curchunk = get_chunk(payload, main.chunknum)

    def safe_read(reqlen):
        '''read data from chunks contiguously'''
        if not main.curchunk:
            return b''
        totalread = b''
        while True:
            thisread = main.curchunk.read(reqlen)
            reqlen -= len(thisread)
            totalread += thisread
            if reqlen <= 0:
                return totalread
            main.chunknum += 1
            main.curchunk = get_chunk(payload, main.chunknum)
            if not main.curchunk:
                return b''

    header = safe_read(30)
    outzip = zipfile.ZipFile(outstr, 'w', allowZip64=True)

    while header:
        kind, _a, filesize, _b, timestamp, _c, namelen, uid, gid, perms = (
            struct.unpack('>IHIIIIHHHH', header))
        filename = safe_read(namelen).decode('utf-8')
        if not filename:
            print('truncated')
            break
        extra1 = struct.pack('<HHBI', 0x5455, 5, 1, timestamp)
        extra2 = struct.pack('<HHHH', 0x7855, 4, uid, gid)
        zinfo = zipfile.ZipInfo(filename, datetime.datetime.utcfromtimestamp(timestamp).timetuple())
        zinfo.extra = extra1 + extra2
        zinfo.create_system = 3
        zinfo.external_attr = perms << 16
        if kind == 0x10020000:
            zinfo.filename += '/'
            zinfo.external_attr |= 0x10
        outzip.writestr(zinfo, safe_read(filesize))
        zinfo.extra = extra1 + struct.pack('<HH', 0x7855, 0)
        header = safe_read(30)

    outzip.close()
    infile.close()

if __name__ == "__main__":
    main()