Include vendor partitions in target zip generation

This change will enable the target zip to include partition defined in
BRILLO_VENDOR_PARTITIONS. This is necessary for OTA to support
partitions other than boot and system in ab_partitions.txt.

BUG: 28623063
Change-Id: I6e0969c31c9ad2e8285c6f560825c160aa8c4d55
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index f98a281..48d5bfd 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -395,8 +395,9 @@
     banner("partition-table")
     AddPartitionTable(output_zip)
 
-  # For devices using A/B update, copy over images from RADIO/ to IMAGES/ and
-  # make sure we have all the needed images ready under IMAGES/.
+  # For devices using A/B update, copy over images from RADIO/ and/or
+  # VENDOR_IMAGES/ to IMAGES/ and make sure we have all the needed
+  # images ready under IMAGES/.
   ab_partitions = os.path.join(OPTIONS.input_tmp, "META", "ab_partitions.txt")
   if os.path.exists(ab_partitions):
     with open(ab_partitions, 'r') as f:
@@ -404,9 +405,14 @@
     for line in lines:
       img_name = line.strip() + ".img"
       img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
+      img_vendor_path = os.path.join(
+        OPTIONS.input_tmp, "VENDOR_IMAGES", img_name)
       if os.path.exists(img_radio_path):
         common.ZipWrite(output_zip, img_radio_path,
                         os.path.join("IMAGES", img_name))
+      elif os.path.exists(img_vendor_path):
+        common.ZipWrite(output_zip, img_vendor_path,
+                        os.path.join("IMAGES", img_name))
 
       # Zip spec says: All slashes MUST be forward slashes.
       img_path = 'IMAGES/' + img_name