#!/usr/bin/env python3

# splituapp for Python2/3 by SuperR. @XDA
#
# For extracting img files from UPDATE.APP

# Based on the app_structure file in split_updata.pl by McSpoon

import os
import re
import sys
import string
import struct
import argparse
import subprocess

parser = argparse.ArgumentParser(description="Split UPDATE.APP file into img files", add_help=False)
required = parser.add_argument_group('Required')
required.add_argument("-f", "--filename", required=True, help="Path to update.app file")
required.add_argument("-o", "--output", required=True, help="Output directory")
optional = parser.add_argument_group('Optional')
optional.add_argument("--no-crc", action="store_true", help="Do not verify CRC")
optional.add_argument("-h", "--help", action="help", help="show this help message and exit")
optional.add_argument("-l", "--list", nargs="*", metavar=('img1', 'img2'), help="List of img files to extract")
args = parser.parse_args()

source = args.filename
outdir = args.output

flist = None
if args.list:
    flist = args.list


CRC_PROG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "crc")


bytenum = 4


with open(source, 'rb') as f:
    while True:
        i = f.read(bytenum)
        if not i:
            break
        if i != b'\x55\xAA\x5A\xA5':
            continue
        else:
            headersize = f.read(bytenum)
            headersize = list(struct.unpack('<L', headersize))[0]
            f.seek(16, 1)
            filesize = f.read(bytenum)
            filesize = list(struct.unpack('<L', filesize))[0]
            f.seek(32, 1)
            filename = f.read(16)
            print("Raw filename: %r" % filename)
            try:
                filename = str(filename.decode())
                filename = ''.join(f for f in filename if f in string.printable).lower()
            except:
                filename = ''

            assert "/" not in filename, "Potential directory traversal: %r" % filename

            f.seek(22, 1)
            crcdata = f.read(headersize - 98)
            print("FILE %r: filesize=%r  len(crcdata)=%r" % (filename, filesize, len(crcdata)))

            getimg = None
            if flist and filename in flist:
                getimg = 1
            elif not flist:
                getimg = 1

            if getimg:
                output_filename = os.path.join(outdir, "%s.img" % filename)
                while os.path.exists(output_filename):
                    output_filename += ".DUPLICATE"
                try:
                    # print('Extracting '+filename+'.img ...')
                    with open(output_filename, 'wb') as o:
                        # Copy data from UPDATE.app to image with 1MiB chunk size => Do not allocate gigabytes of memory
                        # for big images
                        pending_size = filesize
                        while pending_size > 0:
                            current_read_size = min(pending_size, 1024*1024)
                            chunk_data = f.read(current_read_size)
                            assert len(chunk_data) == current_read_size
                            o.write(chunk_data)
                            pending_size -= current_read_size
                except:
                    print('ERROR: Failed to create '+filename+'.img\n')
                    sys.exit()

                if not args.no_crc and filesize < 2**31:
                    # The crc binary is compiled as a 32-bit ELF file and is using standard open() instead of open64()
                    # so it is limited to 2GB file size
                    print("Calculating crc value for %r" % output_filename)
                    crcval = []
                    for i in crcdata:
                        crcval.append('%02X' % i)
                    crcval = ''.join(crcval)
                    cmd = [CRC_PROG, output_filename]
                    # print("CMD: %r" % cmd)
                    crcact = subprocess.check_output(cmd).decode().strip()
                    if crcval != crcact:
                        print('ERROR: crc value for %s does not match, %r <=> %r\n' % (output_filename, crcval, crcact))
                        sys.exit()
            else:
                f.seek(filesize, 1)

            xbytes = bytenum - f.tell() % bytenum
            if xbytes < bytenum:
                f.seek(xbytes, 1)

print('\nExtraction complete')

