sdat2img.py: fix Marshmallow image extraction (issue #6)

The max of the erase block can no longer be assumed to be the
file size. Fix by searching for the max of every type of block.

sdat2img.py

@@ -47,16 +47,12 @@ def parse_transfer_list_file(path):
         trans_list.readline()               # 3rd line = stash entries needed simultaneously
         trans_list.readline()               # 4th line = number of blocks that will be stashed
 
+    commands = []
     for line in trans_list:
         line = line.split(' ')              # 5th & next lines should be only commands
         cmd = line[0]
-        if cmd == 'erase':
-            erase_block_set = rangeset(line[1])
-        elif cmd == 'new':
-            new_block_set = rangeset(line[1])
-        elif cmd == 'zero':
-            # skip zero command, added on android-6.0.0_r1 but useless while decompressing
-            print ('Skipping %s' % cmd + ' command.')
+        if cmd in ['erase', 'new', 'zero']:
+            commands.append([cmd, rangeset(line[1])])
         else:
             # skip lines starting with numbers, they're not commands anyway.
             if not cmd[0].isdigit():
@@ -65,21 +61,24 @@ def parse_transfer_list_file(path):
                 sys.exit(1)
 
     trans_list.close()
-    return version, new_blocks, erase_block_set, new_block_set
+    return version, new_blocks, commands
 
-def init_output_file_size(output_file_obj, erase_block_set):
-    max_block_num = max(pair[1] for pair in erase_block_set)
+def init_output_file_size(output_file_obj, commands):
+    all_block_sets = [i for command in commands for i in command[1]]
+    max_block_num = max(pair[1] for pair in all_block_sets)
     output_file_obj.seek(max_block_num*BLOCK_SIZE - 1)
     output_file_obj.write('\0'.encode('utf-8'))
     output_file_obj.flush()
 
 def main(argv):
-    version, new_blocks, erase_block_set, new_block_set =  parse_transfer_list_file(TRANSFER_LIST_FILE)
+    version, new_blocks, commands = parse_transfer_list_file(TRANSFER_LIST_FILE)
     output_img = open(OUTPUT_IMAGE_FILE, 'wb')
-    init_output_file_size(output_img, erase_block_set)
+    init_output_file_size(output_img, commands)
     new_data_file = open(NEW_DATA_FILE, 'rb')
 
-    for block in new_block_set:
+    for command in commands:
+        if command[0] == 'new':
+            for block in command[1]:
                 begin = block[0]
                 end = block[1]
                 block_count = end - begin
@@ -87,6 +86,8 @@ def main(argv):
                 print('Copying {} blocks into position {}...'.format(block_count, begin))
                 output_img.seek(begin*BLOCK_SIZE)
                 output_img.write(data)
+        else:
+            print('Skipping command %s' % command[0])
 
     output_img.close()
     new_data_file.close()