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". So we create a hard link and put the renamed file inside 93 # the zip 94 original_apex = os.path.join(work_dir, 'original_apex') 95 os.link(args.input, original_apex) 96 cmd.extend(['-C', work_dir]) 97 cmd.extend(['-f', original_apex]) 98 99 # We also need to extract some files from inside of original_apex and zip 100 # together with compressed apex 101 with ZipFile(original_apex, 'r') as zip_obj: 102 extract_dir = os.path.join(work_dir, 'extract') 103 for meta_file in ['apex_manifest.json', 'apex_manifest.pb', 104 'apex_pubkey', 'apex_build_info.pb', 105 'AndroidManifest.xml']: 106 if meta_file in zip_obj.namelist(): 107 zip_obj.extract(meta_file, path=extract_dir) 108 file_path = os.path.join(extract_dir, meta_file) 109 cmd.extend(['-C', extract_dir]) 110 cmd.extend(['-f', file_path]) 111 cmd.extend(['-s', meta_file]) 112 # Extract the image for retrieving root digest 113 zip_obj.extract('apex_payload.img', path= work_dir) 114 image_path = os.path.join(work_dir, 'apex_payload.img') 115 116 # Set digest of original_apex to apex_manifest.pb 117 apex_manifest_path = os.path.join(extract_dir, 'apex_manifest.pb') 118 assert AddOriginalApexDigestToManifest(apex_manifest_path, image_path) 119 120 # Don't forget to compress 121 cmd.extend(['-L', '9']) 122 123 RunCommand(cmd, verbose=True) 124 125 return True 126 127 128def AddOriginalApexDigestToManifest(capex_manifest_path, apex_image_path): 129 # Retrieve the root digest of the image 130 avbtool_cmd = [ 131 'avbtool', 132 'print_partition_digests', '--image', 133 apex_image_path] 134 # avbtool_cmd output has format "<name>: <value>" 135 root_digest = RunCommand(avbtool_cmd, True)[0].decode().split(': ')[1].strip() 136 # Update the manifest proto file 137 with open(capex_manifest_path, 'rb') as f: 138 pb = apex_manifest_pb2.ApexManifest() 139 pb.ParseFromString(f.read()) 140 # Populate CompressedApexMetadata 141 capex_metadata = apex_manifest_pb2.ApexManifest().CompressedApexMetadata() 142 capex_metadata.originalApexDigest = root_digest 143 # Set updated value to protobuf 144 pb.capexMetadata.CopyFrom(capex_metadata) 145 with open(capex_manifest_path, 'wb') as f: 146 f.write(pb.SerializeToString()) 147 return True 148 149 150def ParseArgs(argv): 151 parser = argparse.ArgumentParser() 152 subparsers = parser.add_subparsers(required=True, dest='cmd') 153 154 # Handle sub-command "compress" 155 parser_compress = subparsers.add_parser('compress', 156 help='compresses an APEX') 157 parser_compress.add_argument('--input', type=str, required=True, 158 help='path to input APEX file that will be ' 159 'compressed') 160 parser_compress.add_argument('--output', type=str, required=True, 161 help='output path to compressed APEX file') 162 apex_compression_tool_path_in_environ = \ 163 'APEX_COMPRESSION_TOOL_PATH' in os.environ 164 parser_compress.add_argument( 165 '--apex_compression_tool_path', 166 required=not apex_compression_tool_path_in_environ, 167 default=os.environ['APEX_COMPRESSION_TOOL_PATH'].split(':') 168 if apex_compression_tool_path_in_environ else None, 169 type=lambda s: s.split(':'), 170 help="""A list of directories containing all the tools used by 171 apex_compression_tool (e.g. soong_zip etc.) separated by ':'. Can also 172 be set using the APEX_COMPRESSION_TOOL_PATH environment variable""") 173 parser_compress.set_defaults(func=RunCompress) 174 175 return parser.parse_args(argv) 176 177 178class TempDirectory(object): 179 180 def __enter__(self): 181 self.name = tempfile.mkdtemp() 182 return self.name 183 184 def __exit__(self, *unused): 185 shutil.rmtree(self.name) 186 187 188def main(argv): 189 args = ParseArgs(argv) 190 191 with TempDirectory() as work_dir: 192 success = args.func(args, work_dir) 193 194 if not success: 195 sys.exit(1) 196 197 198if __name__ == '__main__': 199 main(sys.argv[1:]) 200