• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright (C) 2020 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""apex_compression_tool is a tool that can compress/decompress APEX.
18
19Example:
20  apex_compression_tool compress --input /apex/to/compress --output output/path
21  apex_compression_tool decompress --input /apex/to/decompress --output dir/
22  apex_compression_tool verify-compressed --input /file/to/check
23"""
24from __future__ import print_function
25
26import argparse
27import os
28import shutil
29import subprocess
30import sys
31import tempfile
32from zipfile import ZipFile
33
34import apex_manifest_pb2
35
36tool_path_list = None
37
38
39def FindBinaryPath(binary):
40  for path in tool_path_list:
41    binary_path = os.path.join(path, binary)
42    if os.path.exists(binary_path):
43      return binary_path
44  raise Exception('Failed to find binary ' + binary + ' in path ' +
45                  ':'.join(tool_path_list))
46
47
48def RunCommand(cmd, verbose=False, env=None, expected_return_values=None):
49  expected_return_values = expected_return_values or {0}
50  env = env or {}
51  env.update(os.environ.copy())
52
53  cmd[0] = FindBinaryPath(cmd[0])
54
55  if verbose:
56    print('Running: ' + ' '.join(cmd))
57  p = subprocess.Popen(
58      cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env)
59  output, _ = p.communicate()
60
61  if verbose or p.returncode not in expected_return_values:
62    print(output.rstrip())
63
64  assert p.returncode in expected_return_values, 'Failed to execute: ' \
65                                                 + ' '.join(cmd)
66
67  return output, p.returncode
68
69
70def RunCompress(args, work_dir):
71  """RunCompress takes an uncompressed APEX and compresses into compressed APEX
72
73  Compressed apex will contain the following items:
74      - original_apex: The original uncompressed APEX
75      - Duplicates of various meta files inside the input APEX, e.g
76        AndroidManifest.xml, public_key
77
78  Args:
79      args.input: file path to uncompressed APEX
80      args.output: file path to where compressed APEX will be placed
81      work_dir: file path to a temporary folder
82  Returns:
83      True if compression was executed successfully, otherwise False
84  """
85  global tool_path_list
86  tool_path_list = args.apex_compression_tool_path
87
88  cmd = ['soong_zip']
89  cmd.extend(['-o', args.output])
90
91  # We want to put the input apex inside the compressed APEX with name
92  # "original_apex". Originally this was done by creating a hard link
93  # in order to put the renamed file inside the zip, but it causes some issue
94  # when running this tool with Bazel in a sandbox which restricts the function
95  # of creating cross-device links. So instead of creating hard links, we make a
96  # copy of the original_apex here.
97  original_apex = os.path.join(work_dir, 'original_apex')
98  shutil.copy2(args.input, original_apex)
99  cmd.extend(['-C', work_dir])
100  cmd.extend(['-f', original_apex])
101
102  # We also need to extract some files from inside of original_apex and zip
103  # together with compressed apex
104  with ZipFile(original_apex, 'r') as zip_obj:
105    extract_dir = os.path.join(work_dir, 'extract')
106    for meta_file in ['apex_manifest.json', 'apex_manifest.pb',
107                      'apex_pubkey', 'apex_build_info.pb',
108                      'AndroidManifest.xml']:
109      if meta_file in zip_obj.namelist():
110        zip_obj.extract(meta_file, path=extract_dir)
111        file_path = os.path.join(extract_dir, meta_file)
112        cmd.extend(['-C', extract_dir])
113        cmd.extend(['-f', file_path])
114        cmd.extend(['-s', meta_file])
115    # Extract the image for retrieving root digest
116    zip_obj.extract('apex_payload.img', path= work_dir)
117    image_path = os.path.join(work_dir, 'apex_payload.img')
118
119  # Set digest of original_apex to apex_manifest.pb
120  apex_manifest_path = os.path.join(extract_dir, 'apex_manifest.pb')
121  assert AddOriginalApexDigestToManifest(apex_manifest_path, image_path, args.verbose)
122
123  # Don't forget to compress
124  cmd.extend(['-L', '9'])
125
126  RunCommand(cmd, verbose=args.verbose)
127
128  return True
129
130
131def AddOriginalApexDigestToManifest(capex_manifest_path, apex_image_path, verbose=False):
132  # Retrieve the root digest of the image
133  avbtool_cmd = [
134        'avbtool',
135        'print_partition_digests', '--image',
136        apex_image_path]
137  # avbtool_cmd output has format "<name>: <value>"
138  root_digest = RunCommand(avbtool_cmd, verbose=verbose)[0].decode().split(': ')[1].strip()
139  # Update the manifest proto file
140  with open(capex_manifest_path, 'rb') as f:
141    pb = apex_manifest_pb2.ApexManifest()
142    pb.ParseFromString(f.read())
143  # Populate CompressedApexMetadata
144  capex_metadata = apex_manifest_pb2.ApexManifest().CompressedApexMetadata()
145  capex_metadata.originalApexDigest = root_digest
146  # Set updated value to protobuf
147  pb.capexMetadata.CopyFrom(capex_metadata)
148  with open(capex_manifest_path, 'wb') as f:
149    f.write(pb.SerializeToString())
150  return True
151
152
153def ParseArgs(argv):
154  parser = argparse.ArgumentParser()
155  subparsers = parser.add_subparsers(required=True, dest='cmd')
156
157  # Handle sub-command "compress"
158  parser_compress = subparsers.add_parser('compress',
159                                          help='compresses an APEX')
160  parser_compress.add_argument('-v', '--verbose', action='store_true',
161                               help='verbose execution')
162  parser_compress.add_argument('--input', type=str, required=True,
163                               help='path to input APEX file that will be '
164                                    'compressed')
165  parser_compress.add_argument('--output', type=str, required=True,
166                               help='output path to compressed APEX file')
167  apex_compression_tool_path_in_environ = \
168    'APEX_COMPRESSION_TOOL_PATH' in os.environ
169  parser_compress.add_argument(
170      '--apex_compression_tool_path',
171      required=not apex_compression_tool_path_in_environ,
172      default=os.environ['APEX_COMPRESSION_TOOL_PATH'].split(':')
173      if apex_compression_tool_path_in_environ else None,
174      type=lambda s: s.split(':'),
175      help="""A list of directories containing all the tools used by
176        apex_compression_tool (e.g. soong_zip etc.) separated by ':'. Can also
177        be set using the APEX_COMPRESSION_TOOL_PATH environment variable""")
178  parser_compress.set_defaults(func=RunCompress)
179
180  return parser.parse_args(argv)
181
182
183class TempDirectory(object):
184
185  def __enter__(self):
186    self.name = tempfile.mkdtemp()
187    return self.name
188
189  def __exit__(self, *unused):
190    shutil.rmtree(self.name)
191
192
193def main(argv):
194  args = ParseArgs(argv)
195
196  with TempDirectory() as work_dir:
197    success = args.func(args, work_dir)
198
199  if not success:
200    sys.exit(1)
201
202
203if __name__ == '__main__':
204  main(sys.argv[1:])
205