diff --git a/sdat2img.py b/sdat2img.py index 71384f9..ef76133 100755 --- a/sdat2img.py +++ b/sdat2img.py @@ -1,143 +1,80 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- -#==================================================== -# FILE: sdat2img.py -# AUTHORS: xpirt - luxi78 - howellzhu -# DATE: 2018-10-27 10:33:21 CEST -#==================================================== - -from __future__ import print_function -import sys, os, errno - -def main(TRANSFER_LIST_FILE, NEW_DATA_FILE, OUTPUT_IMAGE_FILE): - __version__ = '1.2' - - if sys.hexversion < 0x02070000: - print >> sys.stderr, "Python 2.7 or newer is required." - try: - input = raw_input - except NameError: pass - input('Press ENTER to exit...') - sys.exit(1) - else: - print('sdat2img binary - version: {}\n'.format(__version__)) - - def rangeset(src): - src_set = src.split(',') - num_set = [int(item) for item in src_set] - if len(num_set) != num_set[0]+1: - print('Error on parsing following data to rangeset:\n{}'.format(src), file=sys.stderr) - sys.exit(1) - - return tuple ([ (num_set[i], num_set[i+1]) for i in range(1, len(num_set), 2) ]) - - def parse_transfer_list_file(path): - trans_list = open(TRANSFER_LIST_FILE, 'r') - - # First line in transfer list is the version number - version = int(trans_list.readline()) - - # Second line in transfer list is the total number of blocks we expect to write - new_blocks = int(trans_list.readline()) - - if version >= 2: - # Third line is how many stash entries are needed simultaneously - trans_list.readline() - # Fourth line is the maximum number of blocks that will be stashed simultaneously - trans_list.readline() - - # Subsequent lines are all individual transfer commands - commands = [] - for line in trans_list: - line = line.split(' ') - cmd = line[0] - if cmd in ['erase', 'new', 'zero']: - commands.append([cmd, rangeset(line[1])]) - else: - # Skip lines starting with numbers, they are not commands anyway - if not cmd[0].isdigit(): - print('Command "{}" is not valid.'.format(cmd), file=sys.stderr) - trans_list.close() - sys.exit(1) - - trans_list.close() - return version, new_blocks, commands - - BLOCK_SIZE = 4096 - - version, new_blocks, commands = parse_transfer_list_file(TRANSFER_LIST_FILE) - - if version == 1: - print('Android Lollipop 5.0 detected!\n') - elif version == 2: - print('Android Lollipop 5.1 detected!\n') - elif version == 3: - print('Android Marshmallow 6.x detected!\n') - elif version == 4: - print('Android Nougat 7.x / Oreo 8.x detected!\n') - else: - print('Unknown Android version!\n') - - # Don't clobber existing files to avoid accidental data loss - try: - output_img = open(OUTPUT_IMAGE_FILE, 'wb') - except IOError as e: - if e.errno == errno.EEXIST: - print('Error: the output file "{}" already exists'.format(e.filename), file=sys.stderr) - print('Remove it, rename it, or choose a different file name.', file=sys.stderr) - sys.exit(e.errno) +import argparse +import os + +BLOCK_SIZE = 4096 + + +def range_set(src): + src_set = src.split(',') + num_set = [int(item) for item in src_set] + if len(num_set) != num_set[0]+1: + raise ValueError( + 'Error on parsing following data to range_set:\n{}'.format(src)) + return tuple([(num_set[i], num_set[i+1]) + for i in range(1, len(num_set), 2)]) + + +def transfer_list_file_to_commands(trans_list): + version = int(trans_list.readline()) + trans_list.readline() # new blocks + if version >= 2: + trans_list.readline() # simultaneously stashed entries + trans_list.readline() # max num blocks simultaneously stashed + + commands = [] + for line in trans_list: + line = line.split(' ') + cmd = line[0] + if cmd in ['erase', 'new', 'zero']: + commands.append([cmd, range_set(line[1])]) else: - raise - - new_data_file = open(NEW_DATA_FILE, 'rb') - all_block_sets = [i for command in commands for i in command[1]] - max_file_size = max(pair[1] for pair in all_block_sets)*BLOCK_SIZE - - for command in commands: - if command[0] == 'new': - for block in command[1]: - begin = block[0] - end = block[1] - block_count = end - begin - print('Copying {} blocks into position {}...'.format(block_count, begin)) - - # Position output file - output_img.seek(begin*BLOCK_SIZE) - - # Copy one block at a time - while(block_count > 0): - output_img.write(new_data_file.read(BLOCK_SIZE)) - block_count -= 1 - else: - print('Skipping command {}...'.format(command[0])) + if not cmd[0].isdigit(): + raise ValueError('Command "{}" is not valid.'.format(cmd)) + + return commands + + +def main(transfer_list_file, new_dat_file, output_filename): + commands = transfer_list_file_to_commands(transfer_list_file) + + if os.path.exists(output_filename): + raise ValueError( + 'Error: the output file "{}" already exists' + 'Remove it, rename it, or choose a different file name.' + .format(output_filename)) + + with open(output_filename, 'wb') as output_img: + all_block_sets = [i for command in commands for i in command[1]] + max_file_size = max(pair[1] for pair in all_block_sets) * BLOCK_SIZE + + for command in commands: + if command[0] == 'new': + for block in command[1]: + begin = block[0] + end = block[1] + block_count = end - begin + print('Copying {} blocks to position {}...'.format( + block_count, begin)) + output_img.seek(begin * BLOCK_SIZE) + while block_count > 0: + output_img.write(new_dat_file.read(BLOCK_SIZE)) + block_count -= 1 + else: + print('Skipping command {}'.format(command[0])) - # Make file larger if necessary - if(output_img.tell() < max_file_size): - output_img.truncate(max_file_size) + # Make file larger if necessary + if(output_img.tell() < max_file_size): + output_img.truncate(max_file_size) - output_img.close() - new_data_file.close() - print('Done! Output image: {}'.format(os.path.realpath(output_img.name))) if __name__ == '__main__': - try: - TRANSFER_LIST_FILE = str(sys.argv[1]) - NEW_DATA_FILE = str(sys.argv[2]) - except IndexError: - print('\nUsage: sdat2img.py [system_img]\n') - print(' : transfer list file') - print(' : system new dat file') - print(' [system_img]: output system image\n\n') - print('Visit xda thread for more information.\n') - try: - input = raw_input - except NameError: pass - input('Press ENTER to exit...') - sys.exit() - - try: - OUTPUT_IMAGE_FILE = str(sys.argv[3]) - except IndexError: - OUTPUT_IMAGE_FILE = 'system.img' - - main(TRANSFER_LIST_FILE, NEW_DATA_FILE, OUTPUT_IMAGE_FILE) + parser = argparse.ArgumentParser() + parser.add_argument('transfer_list', help='transfer list file') + parser.add_argument('new_dat', help='system new dat file') + parser.add_argument('output', default='output.img', + help='output image') + args = parser.parse_args() + with open(args.transfer_list, 'r') as transfer_list_file: + with open(args.new_dat, 'rb') as new_dat_file: + main(transfer_list_file, new_dat_file, args.output)