From 254cbdc9825056fd6ac49aea5a4fc9bcefe0b3d0 Mon Sep 17 00:00:00 2001
From: Chirayu Desai <chirayudesai1@gmail.com>
Date: Tue, 17 Sep 2019 03:55:17 +0530
Subject: [PATCH] Support selective extraction of partitions

---
 README.md                      |  7 ++++---
 extract_android_ota_payload.py | 22 +++++++++++++++++++---
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index f7b2416..1b09d1f 100644
--- a/README.md
+++ b/README.md
@@ -10,9 +10,10 @@ Incremental firmware images are not supported (source_copy, source_bsdiff operat
 ## Usage
 
 ```
-$ extract_android_ota_payload.py <payload.bin> [target_dir]
-  <payload.bin> : file extracted from the OTA zip file or the OTA zip file
-  <target_dir>  : output directory for the extracted file
+$ extract_android_ota_payload.py <payload.bin> [target_dir] [partition.img]
+  <payload.bin>   : file extracted from the OTA zip file or the OTA zip file
+  <target_dir>    : output directory for the extracted file
+  <partition_img> : name of partition to be extracted
 ```
 
 ## Example
diff --git a/extract_android_ota_payload.py b/extract_android_ota_payload.py
index 7688ecd..1ef304e 100644
--- a/extract_android_ota_payload.py
+++ b/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,7 +133,7 @@ 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:
@@ -139,7 +141,21 @@ def main(filename, output_dir):
   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)