Add total_data_length utility function am: b7162c1efd am: 1fa6171863

Original change: https://android-review.googlesource.com/c/platform/system/update_engine/+/2195163

Change-Id: Ib7f31f54fe1e9aa5fb2f5659ca39e4ea0883688d
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/scripts/update_payload/payload.py b/scripts/update_payload/payload.py
index e8f0f7e..d8bd83c 100644
--- a/scripts/update_payload/payload.py
+++ b/scripts/update_payload/payload.py
@@ -124,6 +124,8 @@
     """
     if zipfile.is_zipfile(payload_file):
       with zipfile.ZipFile(payload_file) as zfp:
+        if "payload.bin" not in zfp.namelist():
+          raise ValueError(f"payload.bin missing in archive {payload_file}")
         self.payload_file = zfp.open("payload.bin", "r")
     elif isinstance(payload_file, str):
       payload_fp = open(payload_file, "rb")
@@ -151,6 +153,20 @@
   def is_partial(self):
     return self.manifest.partial_update
 
+  @property
+  def total_data_length(self):
+    """Return the total data length of this payload, excluding payload
+    signature at the very end.
+    """
+    # Operations are sorted in ascending data_offset order, so iterating
+    # backwards and find the first one with non zero data_offset will tell
+    # us total data length
+    for partition in reversed(self.manifest.partitions):
+      for op in reversed(partition.operations):
+        if op.data_offset > 0:
+          return op.data_offset + op.data_length
+    return 0
+
   def _ReadHeader(self):
     """Reads and returns the payload header.