From 88be24b0894b9580e214e256d02d2238189a15be Mon Sep 17 00:00:00 2001
From: Chirayu Desai <chirayudesai1@gmail.com>
Date: Wed, 25 May 2022 23:37:27 +0530
Subject: [PATCH] extract_firmware: Support A/B OTA

Change-Id: I845c62a63aaab223c0fb5f49a494d4aa1c165c92
---
 extract_utils.sh | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/extract_utils.sh b/extract_utils.sh
index 793567d..20e8cf1 100644
--- a/extract_utils.sh
+++ b/extract_utils.sh
@@ -1916,10 +1916,18 @@ function extract_firmware() {
         if [ ! -d "$OUTPUT_DIR" ]; then
             mkdir -p "$OUTPUT_DIR"
         fi
-        if [ -f "$SRC/$SRC_FILE" ]; then
-            cp "$SRC/$SRC_FILE" "$OUTPUT_DIR/$DST_FILE"
+        if [ -f "$SRC" ] && [ "${SRC##*.}" == "zip" ]; then
+            # Extract A/B OTA
+            if [ -a "$DUMPDIR"/payload.bin ]; then
+                python3 "$ANDROID_ROOT"/tools/extract-utils/extract_ota.py "$DUMPDIR"/payload.bin -o "$DUMPDIR" -p $(basename "${DST_FILE%.*}") 2>&1
+                cp "$DUMPDIR/$(basename $DST_FILE)" "$OUTPUT_DIR/$DST_FILE"
+            fi
         else
-            cp "$SRC/$DST_FILE" "$OUTPUT_DIR/$DST_FILE"
+            if [ -f "$SRC/$SRC_FILE" ]; then
+                cp "$SRC/$SRC_FILE" "$OUTPUT_DIR/$DST_FILE"
+            else
+                cp "$SRC/$DST_FILE" "$OUTPUT_DIR/$DST_FILE"
+            fi
         fi
         chmod 644 "$OUTPUT_DIR/$DST_FILE"
     done