Add a script to parse an block-based ota package

The script prints the meta info, new/patch data size. It also simulates
the transfer commands and calculates amount of I/O as well as cache
needed.

We can add the more options to parse system/vendor.map so that we can
analyze the location change of files.

Bug: 31514709
Test: parse a bullhead incremental/full ota.
Change-Id: I70ed4f4c15e595b1851109e9799d44ac4e815c16
diff --git a/tools/releasetools/ota_package_parser.py b/tools/releasetools/ota_package_parser.py
new file mode 100755
index 0000000..331122b
--- /dev/null
+++ b/tools/releasetools/ota_package_parser.py
@@ -0,0 +1,228 @@
+#!/usr/bin/env python
+# Copyright (C) 2017 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import sys
+import traceback
+import zipfile
+
+from rangelib import RangeSet
+
+class Stash(object):
+  """Build a map to track stashed blocks during update simulation."""
+
+  def __init__(self):
+    self.blocks_stashed = 0
+    self.overlap_blocks_stashed = 0
+    self.max_stash_needed = 0
+    self.current_stash_size = 0
+    self.stash_map = {}
+
+  def StashBlocks(self, SHA1, blocks):
+    if SHA1 in self.stash_map:
+      logging.info("already stashed {}: {}".format(SHA1, blocks))
+      return
+    self.blocks_stashed += blocks.size()
+    self.current_stash_size += blocks.size()
+    self.max_stash_needed = max(self.current_stash_size, self.max_stash_needed)
+    self.stash_map[SHA1] = blocks
+
+  def FreeBlocks(self, SHA1):
+    assert self.stash_map.has_key(SHA1), "stash {} not found".format(SHA1)
+    self.current_stash_size -= self.stash_map[SHA1].size()
+    del self.stash_map[SHA1]
+
+  def HandleOverlapBlocks(self, SHA1, blocks):
+    self.StashBlocks(SHA1, blocks)
+    self.overlap_blocks_stashed += blocks.size()
+    self.FreeBlocks(SHA1)
+
+
+class OtaPackageParser(object):
+  """Parse a block-based OTA package."""
+
+  def __init__(self, package):
+    self.package = package
+    self.new_data_size = 0
+    self.patch_data_size = 0
+    self.block_written = 0
+    self.block_stashed = 0
+
+  @staticmethod
+  def GetSizeString(size):
+    assert size >= 0
+    base = 1024.0
+    if size <= base:
+      return "{} bytes".format(size)
+    for units in ['K', 'M', 'G']:
+      if size <= base * 1024 or units == 'G':
+        return "{:.1f}{}".format(size / base, units)
+      base *= 1024
+
+  def ParseTransferList(self, name):
+    """Simulate the transfer commands and calculate the amout of I/O."""
+
+    logging.info("\nSimulating commands in '{}':".format(name))
+    lines = self.package.read(name).strip().splitlines()
+    assert len(lines) >= 4, "{} is too short; Transfer list expects at least" \
+        "4 lines, it has {}".format(name, len(lines))
+    assert int(lines[0]) >= 3
+    logging.info("(version: {})".format(lines[0]))
+
+    blocks_written = 0
+    my_stash = Stash()
+    for line in lines[4:]:
+      cmd_list = line.strip().split(" ")
+      cmd_name = cmd_list[0]
+      try:
+        if cmd_name == "new" or cmd_name == "zero":
+          assert len(cmd_list) == 2, "command format error: {}".format(line)
+          target_range = RangeSet.parse_raw(cmd_list[1])
+          blocks_written += target_range.size()
+        elif cmd_name == "move":
+          # Example:  move <onehash> <tgt_range> <src_blk_count> <src_range>
+          # [<loc_range> <stashed_blocks>]
+          assert len(cmd_list) >= 5, "command format error: {}".format(line)
+          target_range = RangeSet.parse_raw(cmd_list[2])
+          blocks_written += target_range.size()
+          if cmd_list[4] == '-':
+            continue
+          SHA1 = cmd_list[1]
+          source_range = RangeSet.parse_raw(cmd_list[4])
+          if target_range.overlaps(source_range):
+            my_stash.HandleOverlapBlocks(SHA1, source_range)
+        elif cmd_name == "bsdiff" or cmd_name == "imgdiff":
+          # Example:  bsdiff <offset> <len> <src_hash> <tgt_hash> <tgt_range>
+          # <src_blk_count> <src_range> [<loc_range> <stashed_blocks>]
+          assert len(cmd_list) >= 8, "command format error: {}".format(line)
+          target_range = RangeSet.parse_raw(cmd_list[5])
+          blocks_written += target_range.size()
+          if cmd_list[7] == '-':
+            continue
+          source_SHA1 = cmd_list[3]
+          source_range = RangeSet.parse_raw(cmd_list[7])
+          if target_range.overlaps(source_range):
+            my_stash.HandleOverlapBlocks(source_SHA1, source_range)
+        elif cmd_name == "stash":
+          assert len(cmd_list) == 3, "command format error: {}".format(line)
+          SHA1 = cmd_list[1]
+          source_range = RangeSet.parse_raw(cmd_list[2])
+          my_stash.StashBlocks(SHA1, source_range)
+        elif cmd_name == "free":
+          assert len(cmd_list) == 2, "command format error: {}".format(line)
+          SHA1 = cmd_list[1]
+          my_stash.FreeBlocks(SHA1)
+      except:
+        logging.error("failed to parse command in: " + line)
+        raise
+
+    self.block_written += blocks_written
+    self.block_stashed += my_stash.blocks_stashed
+
+    logging.info("blocks written: {}  (expected: {})".format(
+        blocks_written, lines[1]))
+    logging.info("max blocks stashed simultaneously: {}  (expected: {})".
+        format(my_stash.max_stash_needed, lines[3]))
+    logging.info("total blocks stashed: {}".format(my_stash.blocks_stashed))
+    logging.info("blocks stashed implicitly: {}".format(
+        my_stash.overlap_blocks_stashed))
+
+  def PrintDataInfo(self, partition):
+    logging.info("\nReading data info for {} partition:".format(partition))
+    new_data = self.package.getinfo(partition + ".new.dat")
+    patch_data = self.package.getinfo(partition + ".patch.dat")
+    logging.info("{:<40}{:<40}".format(new_data.filename, patch_data.filename))
+    logging.info("{:<40}{:<40}".format(
+          "compress_type: " + str(new_data.compress_type),
+          "compress_type: " + str(patch_data.compress_type)))
+    logging.info("{:<40}{:<40}".format(
+          "compressed_size: " + OtaPackageParser.GetSizeString(
+              new_data.compress_size),
+          "compressed_size: " + OtaPackageParser.GetSizeString(
+              patch_data.compress_size)))
+    logging.info("{:<40}{:<40}".format(
+        "file_size: " + OtaPackageParser.GetSizeString(new_data.file_size),
+        "file_size: " + OtaPackageParser.GetSizeString(patch_data.file_size)))
+
+    self.new_data_size += new_data.file_size
+    self.patch_data_size += patch_data.file_size
+
+  def AnalyzePartition(self, partition):
+    assert partition in ("system", "vendor")
+    assert partition + ".new.dat" in self.package.namelist()
+    assert partition + ".patch.dat" in self.package.namelist()
+    assert partition + ".transfer.list" in self.package.namelist()
+
+    self.PrintDataInfo(partition)
+    self.ParseTransferList(partition + ".transfer.list")
+
+  def PrintMetadata(self):
+    metadata_path = "META-INF/com/android/metadata"
+    logging.info("\nMetadata info:")
+    metadata_info = {}
+    for line in self.package.read(metadata_path).strip().splitlines():
+      index = line.find("=")
+      metadata_info[line[0 : index].strip()] = line[index + 1:].strip()
+    assert metadata_info.get("ota-type") == "BLOCK"
+    assert "pre-device" in metadata_info
+    logging.info("device: {}".format(metadata_info["pre-device"]))
+    if "pre-build" in metadata_info:
+      logging.info("pre-build: {}".format(metadata_info["pre-build"]))
+    assert "post-build" in metadata_info
+    logging.info("post-build: {}".format(metadata_info["post-build"]))
+
+  def Analyze(self):
+    logging.info("Analyzing ota package: " + self.package.filename)
+    self.PrintMetadata()
+    assert "system.new.dat" in self.package.namelist()
+    self.AnalyzePartition("system")
+    if "vendor.new.dat" in self.package.namelist():
+      self.AnalyzePartition("vendor")
+
+    #TODO Add analysis of other partitions(e.g. bootloader, boot, radio)
+
+    BLOCK_SIZE = 4096
+    logging.info("\nOTA package analyzed:")
+    logging.info("new data size (uncompressed): " +
+        OtaPackageParser.GetSizeString(self.new_data_size))
+    logging.info("patch data size (uncompressed): " +
+        OtaPackageParser.GetSizeString(self.patch_data_size))
+    logging.info("total data written: " +
+        OtaPackageParser.GetSizeString(self.block_written * BLOCK_SIZE))
+    logging.info("total data stashed: " +
+        OtaPackageParser.GetSizeString(self.block_stashed * BLOCK_SIZE))
+
+
+def main(argv):
+  parser = argparse.ArgumentParser(description='Analyze an OTA package.')
+  parser.add_argument("ota_package", help='Path of the OTA package.')
+  args = parser.parse_args(argv)
+
+  logging_format = '%(message)s'
+  logging.basicConfig(level=logging.INFO, format=logging_format)
+
+  try:
+    with zipfile.ZipFile(args.ota_package, 'r') as package:
+      package_parser = OtaPackageParser(package)
+      package_parser.Analyze()
+  except:
+    logging.error("Failed to read " + args.ota_package)
+    traceback.print_exc()
+    sys.exit(1)
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])