1#!/usr/bin/env python
2# Copyright (C) 2017 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import argparse
17import logging
18import sys
19import traceback
20import zipfile
21
22from rangelib import RangeSet
23
24class Stash(object):
25  """Build a map to track stashed blocks during update simulation."""
26
27  def __init__(self):
28    self.blocks_stashed = 0
29    self.overlap_blocks_stashed = 0
30    self.max_stash_needed = 0
31    self.current_stash_size = 0
32    self.stash_map = {}
33
34  def StashBlocks(self, SHA1, blocks):
35    if SHA1 in self.stash_map:
36      logging.info("already stashed {}: {}".format(SHA1, blocks))
37      return
38    self.blocks_stashed += blocks.size()
39    self.current_stash_size += blocks.size()
40    self.max_stash_needed = max(self.current_stash_size, self.max_stash_needed)
41    self.stash_map[SHA1] = blocks
42
43  def FreeBlocks(self, SHA1):
44    assert self.stash_map.has_key(SHA1), "stash {} not found".format(SHA1)
45    self.current_stash_size -= self.stash_map[SHA1].size()
46    del self.stash_map[SHA1]
47
48  def HandleOverlapBlocks(self, SHA1, blocks):
49    self.StashBlocks(SHA1, blocks)
50    self.overlap_blocks_stashed += blocks.size()
51    self.FreeBlocks(SHA1)
52
53
54class OtaPackageParser(object):
55  """Parse a block-based OTA package."""
56
57  def __init__(self, package):
58    self.package = package
59    self.new_data_size = 0
60    self.patch_data_size = 0
61    self.block_written = 0
62    self.block_stashed = 0
63
64  @staticmethod
65  def GetSizeString(size):
66    assert size >= 0
67    base = 1024.0
68    if size <= base:
69      return "{} bytes".format(size)
70    for units in ['K', 'M', 'G']:
71      if size <= base * 1024 or units == 'G':
72        return "{:.1f}{}".format(size / base, units)
73      base *= 1024
74
75  def ParseTransferList(self, name):
76    """Simulate the transfer commands and calculate the amout of I/O."""
77
78    logging.info("\nSimulating commands in '{}':".format(name))
79    lines = self.package.read(name).strip().splitlines()
80    assert len(lines) >= 4, "{} is too short; Transfer list expects at least" \
81        "4 lines, it has {}".format(name, len(lines))
82    assert int(lines[0]) >= 3
83    logging.info("(version: {})".format(lines[0]))
84
85    blocks_written = 0
86    my_stash = Stash()
87    for line in lines[4:]:
88      cmd_list = line.strip().split(" ")
89      cmd_name = cmd_list[0]
90      try:
91        if cmd_name == "new" or cmd_name == "zero":
92          assert len(cmd_list) == 2, "command format error: {}".format(line)
93          target_range = RangeSet.parse_raw(cmd_list[1])
94          blocks_written += target_range.size()
95        elif cmd_name == "move":
96          # Example:  move <onehash> <tgt_range> <src_blk_count> <src_range>
97          # [<loc_range> <stashed_blocks>]
98          assert len(cmd_list) >= 5, "command format error: {}".format(line)
99          target_range = RangeSet.parse_raw(cmd_list[2])
100          blocks_written += target_range.size()
101          if cmd_list[4] == '-':
102            continue
103          SHA1 = cmd_list[1]
104          source_range = RangeSet.parse_raw(cmd_list[4])
105          if target_range.overlaps(source_range):
106            my_stash.HandleOverlapBlocks(SHA1, source_range)
107        elif cmd_name == "bsdiff" or cmd_name == "imgdiff":
108          # Example:  bsdiff <offset> <len> <src_hash> <tgt_hash> <tgt_range>
109          # <src_blk_count> <src_range> [<loc_range> <stashed_blocks>]
110          assert len(cmd_list) >= 8, "command format error: {}".format(line)
111          target_range = RangeSet.parse_raw(cmd_list[5])
112          blocks_written += target_range.size()
113          if cmd_list[7] == '-':
114            continue
115          source_SHA1 = cmd_list[3]
116          source_range = RangeSet.parse_raw(cmd_list[7])
117          if target_range.overlaps(source_range):
118            my_stash.HandleOverlapBlocks(source_SHA1, source_range)
119        elif cmd_name == "stash":
120          assert len(cmd_list) == 3, "command format error: {}".format(line)
121          SHA1 = cmd_list[1]
122          source_range = RangeSet.parse_raw(cmd_list[2])
123          my_stash.StashBlocks(SHA1, source_range)
124        elif cmd_name == "free":
125          assert len(cmd_list) == 2, "command format error: {}".format(line)
126          SHA1 = cmd_list[1]
127          my_stash.FreeBlocks(SHA1)
128      except:
129        logging.error("failed to parse command in: " + line)
130        raise
131
132    self.block_written += blocks_written
133    self.block_stashed += my_stash.blocks_stashed
134
135    logging.info("blocks written: {}  (expected: {})".format(
136        blocks_written, lines[1]))
137    logging.info("max blocks stashed simultaneously: {}  (expected: {})".
138        format(my_stash.max_stash_needed, lines[3]))
139    logging.info("total blocks stashed: {}".format(my_stash.blocks_stashed))
140    logging.info("blocks stashed implicitly: {}".format(
141        my_stash.overlap_blocks_stashed))
142
143  def PrintDataInfo(self, partition):
144    logging.info("\nReading data info for {} partition:".format(partition))
145    new_data = self.package.getinfo(partition + ".new.dat")
146    patch_data = self.package.getinfo(partition + ".patch.dat")
147    logging.info("{:<40}{:<40}".format(new_data.filename, patch_data.filename))
148    logging.info("{:<40}{:<40}".format(
149          "compress_type: " + str(new_data.compress_type),
150          "compress_type: " + str(patch_data.compress_type)))
151    logging.info("{:<40}{:<40}".format(
152          "compressed_size: " + OtaPackageParser.GetSizeString(
153              new_data.compress_size),
154          "compressed_size: " + OtaPackageParser.GetSizeString(
155              patch_data.compress_size)))
156    logging.info("{:<40}{:<40}".format(
157        "file_size: " + OtaPackageParser.GetSizeString(new_data.file_size),
158        "file_size: " + OtaPackageParser.GetSizeString(patch_data.file_size)))
159
160    self.new_data_size += new_data.file_size
161    self.patch_data_size += patch_data.file_size
162
163  def AnalyzePartition(self, partition):
164    assert partition in ("system", "vendor")
165    assert partition + ".new.dat" in self.package.namelist()
166    assert partition + ".patch.dat" in self.package.namelist()
167    assert partition + ".transfer.list" in self.package.namelist()
168
169    self.PrintDataInfo(partition)
170    self.ParseTransferList(partition + ".transfer.list")
171
172  def PrintMetadata(self):
173    metadata_path = "META-INF/com/android/metadata"
174    logging.info("\nMetadata info:")
175    metadata_info = {}
176    for line in self.package.read(metadata_path).strip().splitlines():
177      index = line.find("=")
178      metadata_info[line[0 : index].strip()] = line[index + 1:].strip()
179    assert metadata_info.get("ota-type") == "BLOCK"
180    assert "pre-device" in metadata_info
181    logging.info("device: {}".format(metadata_info["pre-device"]))
182    if "pre-build" in metadata_info:
183      logging.info("pre-build: {}".format(metadata_info["pre-build"]))
184    assert "post-build" in metadata_info
185    logging.info("post-build: {}".format(metadata_info["post-build"]))
186
187  def Analyze(self):
188    logging.info("Analyzing ota package: " + self.package.filename)
189    self.PrintMetadata()
190    assert "system.new.dat" in self.package.namelist()
191    self.AnalyzePartition("system")
192    if "vendor.new.dat" in self.package.namelist():
193      self.AnalyzePartition("vendor")
194
195    #TODO Add analysis of other partitions(e.g. bootloader, boot, radio)
196
197    BLOCK_SIZE = 4096
198    logging.info("\nOTA package analyzed:")
199    logging.info("new data size (uncompressed): " +
200        OtaPackageParser.GetSizeString(self.new_data_size))
201    logging.info("patch data size (uncompressed): " +
202        OtaPackageParser.GetSizeString(self.patch_data_size))
203    logging.info("total data written: " +
204        OtaPackageParser.GetSizeString(self.block_written * BLOCK_SIZE))
205    logging.info("total data stashed: " +
206        OtaPackageParser.GetSizeString(self.block_stashed * BLOCK_SIZE))
207
208
209def main(argv):
210  parser = argparse.ArgumentParser(description='Analyze an OTA package.')
211  parser.add_argument("ota_package", help='Path of the OTA package.')
212  args = parser.parse_args(argv)
213
214  logging_format = '%(message)s'
215  logging.basicConfig(level=logging.INFO, format=logging_format)
216
217  try:
218    with zipfile.ZipFile(args.ota_package, 'r') as package:
219      package_parser = OtaPackageParser(package)
220      package_parser.Analyze()
221  except:
222    logging.error("Failed to read " + args.ota_package)
223    traceback.print_exc()
224    sys.exit(1)
225
226
227if __name__ == '__main__':
228  main(sys.argv[1:])
229