Remove logical partitions from ab_partitions.txt in retrofit package.

For retrofit updates:
ab_partitions -= dynamic_partition_list
ab_partitions += super_block_devices

For example,
if super_block_devices == ["system", "vendor", "odm"] and
   dynamic_partition_list == ["system", "vendor", "product"],
product partition needs to be removed from ab_partitions.txt
otherwise brillo_update_payload will generate a payload with
product partition in it.

Test: retrofit update that adds a partition
Fixes: 127425410

Change-Id: Id79a410cee3c611ac50d27f14282916aea34f938
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 4336cb3..02d07da 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -226,8 +226,9 @@
 METADATA_NAME = 'META-INF/com/android/metadata'
 POSTINSTALL_CONFIG = 'META/postinstall_config.txt'
 DYNAMIC_PARTITION_INFO = 'META/dynamic_partitions_info.txt'
+AB_PARTITIONS = 'META/ab_partitions.txt'
 UNZIP_PATTERN = ['IMAGES/*', 'META/*']
-SUPER_SPLIT_PATTERN = ['OTA/super_*.img']
+RETROFIT_DAP_UNZIP_PATTERN = ['OTA/super_*.img', AB_PARTITIONS]
 
 
 class BuildInfo(object):
@@ -1821,7 +1822,8 @@
 
 
 def GetTargetFilesZipForRetrofitDynamicPartitions(input_file,
-                                                  super_block_devices):
+                                                  super_block_devices,
+                                                  dynamic_partition_list):
   """Returns a target-files.zip for retrofitting dynamic partitions.
 
   This allows brillo_update_payload to generate an OTA based on the exact
@@ -1830,6 +1832,7 @@
   Args:
     input_file: The input target-files.zip filename.
     super_block_devices: The list of super block devices
+    dynamic_partition_list: The list of dynamic partitions
 
   Returns:
     The filename of target-files.zip with *.img replaced with super_*.img for
@@ -1846,8 +1849,34 @@
   with zipfile.ZipFile(input_file, 'r') as input_zip:
     namelist = input_zip.namelist()
 
+  input_tmp = common.UnzipTemp(input_file, RETROFIT_DAP_UNZIP_PATTERN)
+
+  # Remove partitions from META/ab_partitions.txt that is in
+  # dynamic_partition_list but not in super_block_devices so that
+  # brillo_update_payload won't generate update for those logical partitions.
+  ab_partitions_file = os.path.join(input_tmp, *AB_PARTITIONS.split('/'))
+  with open(ab_partitions_file) as f:
+    ab_partitions_lines = f.readlines()
+    ab_partitions = [line.strip() for line in ab_partitions_lines]
+  # Assert that all super_block_devices are in ab_partitions
+  super_device_not_updated = [partition for partition in super_block_devices
+                              if partition not in ab_partitions]
+  assert not super_device_not_updated, \
+      "{} is in super_block_devices but not in {}".format(
+          super_device_not_updated, AB_PARTITIONS)
+  # ab_partitions -= (dynamic_partition_list - super_block_devices)
+  new_ab_partitions = common.MakeTempFile(prefix="ab_partitions", suffix=".txt")
+  with open(new_ab_partitions, 'w') as f:
+    for partition in ab_partitions:
+      if (partition in dynamic_partition_list and
+          partition not in super_block_devices):
+          logger.info("Dropping %s from ab_partitions.txt", partition)
+          continue
+      f.write(partition + "\n")
+  to_delete = [AB_PARTITIONS]
+
   # Always skip postinstall for a retrofit update.
-  to_delete = [POSTINSTALL_CONFIG]
+  to_delete += [POSTINSTALL_CONFIG]
 
   # Delete dynamic_partitions_info.txt so that brillo_update_payload thinks this
   # is a regular update on devices without dynamic partitions support.
@@ -1859,7 +1888,6 @@
 
   common.ZipDelete(target_file, to_delete)
 
-  input_tmp = common.UnzipTemp(input_file, SUPER_SPLIT_PATTERN)
   target_zip = zipfile.ZipFile(target_file, 'a', allowZip64=True)
 
   # Write super_{foo}.img as {foo}.img.
@@ -1869,6 +1897,9 @@
     unzipped_file = os.path.join(input_tmp, *src.split('/'))
     common.ZipWrite(target_zip, unzipped_file, arcname=dst)
 
+  # Write new ab_partitions.txt file
+  common.ZipWrite(target_zip, new_ab_partitions, arcname=AB_PARTITIONS)
+
   common.ZipClose(target_zip)
 
   return target_file
@@ -1897,7 +1928,8 @@
 
   if OPTIONS.retrofit_dynamic_partitions:
     target_file = GetTargetFilesZipForRetrofitDynamicPartitions(
-        target_file, target_info.get("super_block_devices").strip().split())
+        target_file, target_info.get("super_block_devices").strip().split(),
+        target_info.get("dynamic_partition_list").strip().split())
   elif OPTIONS.skip_postinstall:
     target_file = GetTargetFilesZipWithoutPostinstallConfig(target_file)