sdat2img

#!/usr/bin/env python3

import brotli, os, sys

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'releasetools'))
from releasetools import rangelib

def ParseTransferList(TRANSFER_LIST):
    lines = TRANSFER_LIST.read().strip().splitlines()
    commands = []
    for line in lines[4:]:
        cmd_list = line.strip().split(" ")
        cmd_name = cmd_list[0]

        if cmd_name == "new" or cmd_name == "zero" or cmd_name == "erase":
            commands.append([cmd_name, rangelib.RangeSet.parse_raw(cmd_list[1])])
        elif not cmd_name.isdigit():
            print("Error: failed to parse command in: " + line)
            sys.exit(1)

    return commands

def main(TRANSFER_LIST, NEW_DAT, OUT_FILE):
    __version__ = '3.0'
    print('sdat2img binary - version: %s' % __version__)

    # Decompress image if necessary
    if NEW_DAT.endswith('.br'):
        CHUNK_SIZE = 4096 * 4096
        NEW_DAT_BR = NEW_DAT
        NEW_DAT = os.path.splitext(NEW_DAT)[0]

        print('- Decompressing image...')
        with open(NEW_DAT_BR, 'rb') as input_image, open(NEW_DAT, 'wb') as output_image:
            decompressor = brotli.Decompressor()
            while chunk := input_image.read(CHUNK_SIZE):
                output_image.write(decompressor.process(chunk))
            os.remove(NEW_DAT_BR)

    # Set output path if none is specified
    if OUT_FILE is None:
        OUT_FILE = NEW_DAT.rsplit('/', 1)[-1].replace(".new.dat", ".img")

    # Parse transfer list
    with open(TRANSFER_LIST, 'r') as trans_list:
        commands = ParseTransferList(trans_list)

    # Convert image
    with open(NEW_DAT, 'rb') as input_image, open(OUT_FILE, 'wb') as output_image:
        BLOCK_SIZE = 4096
        block_ranges = (pair for command in commands for pair in command[1])
        max_file_size = max(end for start, end in block_ranges) * BLOCK_SIZE

        print('- Converting image...')
        for command in commands:
            if command[0] == 'new':
                for start, end in command[1]:
                    block_count = end - start
                    output_image.seek(start * BLOCK_SIZE)
                    data_chunk = input_image.read(block_count * BLOCK_SIZE)
                    output_image.write(data_chunk)

        if (output_image.tell() < max_file_size):
            output_image.truncate(max_file_size)

    print(f'\nDone! Output file: {OUT_FILE}')
    return

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='A tool to convert Android data images into raw images.')

    parser.add_argument('transfer_list', help='input transfer list')
    parser.add_argument('new_dat', help='input data image')
    parser.add_argument('-o', '--out', help='output file path (partition.img in the current directory by default)')

    args = parser.parse_args()

    main(args.transfer_list, args.new_dat, args.out)