From 47851781e7cb9ee14ebaf613ddaaf839f7fb1ec6 Mon Sep 17 00:00:00 2001
From: Chirayu Desai <chirayudesai1@gmail.com>
Date: Fri, 26 Nov 2021 05:47:25 +0530
Subject: [PATCH] extract_utils: Support A/B OTAs

Change-Id: I52e66fa392f7a8fb5401db179b1242d5d216d137
---
 extract_utils.sh | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/extract_utils.sh b/extract_utils.sh
index c77a4d6..2010546 100644
--- a/extract_utils.sh
+++ b/extract_utils.sh
@@ -1519,10 +1519,9 @@ function extract() {
             unzip "$SRC" -d "$DUMPDIR"
             echo "$MD5" > "$DUMPDIR"/zipmd5.txt
 
-            # Stop if an A/B OTA zip is detected. We cannot extract these.
+            # Extract A/B OTA
             if [ -a "$DUMPDIR"/payload.bin ]; then
-                echo "A/B style OTA zip detected. This is not supported at this time. Stopping..."
-                exit 1
+                python3 "$ANDROID_ROOT"/tools/extract-utils/extract_ota.py "$DUMPDIR"/payload.bin -o "$DUMPDIR" -p "system" "odm" "product" "system_ext" "vendor" 2>&1
             fi
 
             for PARTITION in "system" "odm" "product" "system_ext" "vendor"
@@ -1541,6 +1540,9 @@ function extract() {
                     extract_img_data "$DUMPDIR"/"$PARTITION".img "$DUMPDIR"/"$PARTITION"/
                     rm "$DUMPDIR"/"$PARTITION".img
                 fi
+                if [ -a "$DUMPDIR"/"$PARTITION".img ]; then
+                    extract_img_data "$DUMPDIR"/"$PARTITION".img "$DUMPDIR"/"$PARTITION"/
+                fi
             done
         fi