1#
2# Copyright (C) 2020 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#
16
17"""Tools for running host side simulation of an OTA update."""
18
19
20from __future__ import print_function
21
22import argparse
23import filecmp
24import os
25import shutil
26import subprocess
27import sys
28import tempfile
29import zipfile
30
31import update_payload
32
33
34def extract_file(zip_file_path, entry_name, target_file_path):
35  """Extract a file from zip archive into |target_file_path|"""
36  with open(target_file_path, 'wb') as out_fp:
37    if isinstance(zip_file_path, zipfile.ZipFile):
38      with zip_file_path.open(entry_name) as fp:
39        shutil.copyfileobj(fp, out_fp)
40    elif os.path.isdir(zip_file_path):
41      with open(os.path.join(zip_file_path, entry_name), "rb") as fp:
42        shutil.copyfileobj(fp, out_fp)
43
44def is_sparse_image(filepath):
45  with open(filepath, 'rb') as fp:
46    # Magic for android sparse image format
47    # https://source.android.com/devices/bootloader/images
48    return fp.read(4) == b'\x3A\xFF\x26\xED'
49
50def extract_img(zip_archive, img_name, output_path):
51  entry_name = "IMAGES/" + img_name + ".img"
52  extract_file(zip_archive, entry_name, output_path)
53  if is_sparse_image(output_path):
54    raw_img_path = output_path + ".raw"
55    subprocess.check_output(["simg2img", output_path, raw_img_path])
56    os.rename(raw_img_path, output_path)
57
58def run_ota(source, target, payload_path, tempdir):
59  """Run an OTA on host side"""
60  payload = update_payload.Payload(payload_path)
61  payload.Init()
62  if zipfile.is_zipfile(source):
63    source = zipfile.ZipFile(source)
64  if zipfile.is_zipfile(target):
65    target = zipfile.ZipFile(target)
66
67  old_partitions = []
68  new_partitions = []
69  expected_new_partitions = []
70  for part in payload.manifest.partitions:
71    name = part.partition_name
72    old_image = os.path.join(tempdir, "source_" + name + ".img")
73    new_image = os.path.join(tempdir, "target_" + name + ".img")
74    print("Extracting source image for", name)
75    extract_img(source, name, old_image)
76    print("Extracting target image for", name)
77    extract_img(target, name, new_image)
78
79    old_partitions.append(old_image)
80    scratch_image_name = new_image + ".actual"
81    new_partitions.append(scratch_image_name)
82    with open(scratch_image_name, "wb") as fp:
83      fp.truncate(part.new_partition_info.size)
84    expected_new_partitions.append(new_image)
85
86  delta_generator_args = ["delta_generator", "--in_file=" + payload_path]
87  partition_names = [
88      part.partition_name for part in payload.manifest.partitions
89  ]
90  delta_generator_args.append("--partition_names=" + ":".join(partition_names))
91  delta_generator_args.append("--old_partitions=" + ":".join(old_partitions))
92  delta_generator_args.append("--new_partitions=" + ":".join(new_partitions))
93
94  subprocess.check_output(delta_generator_args)
95
96  valid = True
97  for (expected_part, actual_part, part_name) in \
98      zip(expected_new_partitions, new_partitions, partition_names):
99    if filecmp.cmp(expected_part, actual_part):
100      print("Partition `{}` is valid".format(part_name))
101    else:
102      valid = False
103      print(
104          "Partition `{}` is INVALID expected image: {} actual image: {}"
105          .format(part_name, expected_part, actual_part))
106
107  if not valid and sys.stdout.isatty():
108    input("Paused to investigate invalid partitions, press any key to exit.")
109
110
111def main():
112  parser = argparse.ArgumentParser(
113      description="Run host side simulation of OTA package")
114  parser.add_argument(
115      "--source",
116      help="Target file zip for the source build",
117      required=True, nargs=1)
118  parser.add_argument(
119      "--target",
120      help="Target file zip for the target build",
121      required=True, nargs=1)
122  parser.add_argument(
123      "payload",
124      help="payload.bin for the OTA package, or a zip of OTA package itself",
125      nargs=1)
126  args = parser.parse_args()
127  print(args)
128
129  assert os.path.exists(args.source[0]), \
130    "source target file must point to a valid zipfile or directory"
131  assert os.path.exists(args.target[0]), \
132    "target path must point to a valid zipfile or directory"
133
134  # pylint: disable=no-member
135  with tempfile.TemporaryDirectory() as tempdir:
136    payload_path = args.payload[0]
137    if zipfile.is_zipfile(payload_path):
138      with zipfile.ZipFile(payload_path, "r") as zfp:
139        payload_entry_name = 'payload.bin'
140        zfp.extract(payload_entry_name, tempdir)
141        payload_path = os.path.join(tempdir, payload_entry_name)
142
143    run_ota(args.source[0], args.target[0], payload_path, tempdir)
144
145
146if __name__ == '__main__':
147  main()
148