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)
@@ -116,6 +116,8 @@ def main(filename, output_dir):
   payload.Init()
 
   for p in payload.manifest.partitions:
+    if (partition is not None and p.partition_name != partition):
+      continue
     name = p.partition_name + '.img'
     print("Extracting '%s'" % name)
     fname = os.path.join(output_dir, name)
@@ -131,15 +133,26 @@ 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_name]' % 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:
+      if not os.path.exists(sys.argv[2]):
+        partition = sys.argv[2]
+      else:
+        partition = None
+    except IndexError:
+      partition = None
+
   if not os.path.exists(output_dir):
     os.makedirs(output_dir)
 
-  main(filename, output_dir)
+  main(filename, output_dir, partition)