Support selective extraction of partitions

extract_android_ota_payload.py

@@ -104,7 +104,7 @@ def parse_payload(payload_f, partition, out_f):
     else:
       raise PayloadError('Unhandled operation type (%d)' % operation.type)
 
-def main(filename, output_dir):
+def main(filename, output_dir, partition):
   if filename.endswith('.zip'):
     print("Extracting 'payload.bin' from OTA file...")
     ota_zf = zipfile.ZipFile(filename)
@@ -117,6 +117,8 @@ def main(filename, output_dir):
 
   for p in payload.manifest.partitions:
     name = p.partition_name + '.img'
+    if (partition is not None and name != partition):
+      continue
     print("Extracting '%s'" % name)
     fname = os.path.join(output_dir, name)
     out_f = open(fname, 'w')
@@ -131,15 +133,29 @@ def main(filename, output_dir):
   try:
     filename = sys.argv[1]
   except:
-    print('Usage: %s payload.bin [output_dir]' % sys.argv[0])
+    print('Usage: %s payload.bin [output_dir] [partition.img]' % sys.argv[0])
     sys.exit()
 
   try:
     output_dir = sys.argv[2]
   except IndexError:
     output_dir = os.getcwd()
 
+  try:
+    partition = sys.argv[3]
+  except IndexError:
+    try:
+      partition = sys.argv[2]
+    except IndexError:
+      partition = None
+
+  if partition is not None and not partition.endswith('.img'):
+    partition = None
+
+  if output_dir.endswith('.img'):
+    output_dir = os.getcwd()
+
   if not os.path.exists(output_dir):
     os.makedirs(output_dir)
 
-  main(filename, output_dir)
+  main(filename, output_dir, partition)