#!/usr/bin/env python # # Copyright (C) 2021 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. """sign_virt_apex is a command line tool for sign the Virt APEX file. Typical usage: sign_virt_apex [-v] [--avbtool path_to_avbtool] path_to_key payload_contents_dir sign_virt_apex uses external tools which are assumed to be available via PATH. - avbtool (--avbtool can override the tool) - lpmake, lpunpack, simg2img, img2simg """ import argparse import os import re import shutil import subprocess import sys import tempfile def ParseArgs(argv): parser = argparse.ArgumentParser(description='Sign the Virt APEX') parser.add_argument( '-v', '--verbose', action='store_true', help='verbose execution') parser.add_argument( '--avbtool', default='avbtool', help='Optional flag that specifies the AVB tool to use. Defaults to `avbtool`.') parser.add_argument( 'key', help='path to the private key file.') parser.add_argument( 'input_dir', help='the directory having files to be packaged') return parser.parse_args(argv) def RunCommand(args, cmd, env=None, expected_return_values={0}): env = env or {} env.update(os.environ.copy()) # TODO(b/193504286): we need a way to find other tool (cmd[0]) in various contexts # e.g. sign_apex.py, sign_target_files_apk.py if cmd[0] == 'avbtool': cmd[0] = args.avbtool if args.verbose: print('Running: ' + ' '.join(cmd)) p = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env, universal_newlines=True) output, _ = p.communicate() if args.verbose or p.returncode not in expected_return_values: print(output.rstrip()) assert p.returncode in expected_return_values, ( '%d Failed to execute: ' + ' '.join(cmd)) % p.returncode return (output, p.returncode) def ReadBytesSize(value): return int(value.removesuffix(' bytes')) def AvbInfo(args, image_path, descriptor_name=None): """Parses avbtool --info image output Args: args: program arguments. image_path: The path to the image. descriptor_name: Descriptor name of interest. Returns: A pair of - a dict that contains VBMeta info. None if there's no VBMeta info. - a dict that contains target descriptor info. None if name is not specified or not found. """ if not os.path.exists(image_path): raise ValueError('Failed to find image: {}'.format(image_path)) output, ret_code = RunCommand( args, ['avbtool', 'info_image', '--image', image_path], expected_return_values={0, 1}) if ret_code == 1: return None, None info, descriptor = {}, None # Read `avbtool info_image` output as "key:value" lines matcher = re.compile(r'^(\s*)([^:]+):\s*(.*)$') def IterateLine(output): for line in output.split('\n'): line_info = matcher.match(line) if not line_info: continue yield line_info.group(1), line_info.group(2), line_info.group(3) gen = IterateLine(output) # Read VBMeta info for _, key, value in gen: if key == 'Descriptors': break info[key] = value if descriptor_name: for indent, key, _ in gen: # Read a target descriptor if key == descriptor_name: cur_indent = indent descriptor = {} for indent, key, value in gen: if indent == cur_indent: break descriptor[key] = value break return info, descriptor def AddHashFooter(args, key, image_path): info, descriptor = AvbInfo(args, image_path, 'Hash descriptor') if info: image_size = ReadBytesSize(info['Image size']) algorithm = info['Algorithm'] partition_name = descriptor['Partition Name'] partition_size = str(image_size) cmd = ['avbtool', 'add_hash_footer', '--key', key, '--algorithm', algorithm, '--partition_name', partition_name, '--partition_size', partition_size, '--image', image_path] RunCommand(args, cmd) def AddHashTreeFooter(args, key, image_path): info, descriptor = AvbInfo(args, image_path, 'Hashtree descriptor') if info: image_size = ReadBytesSize(info['Image size']) algorithm = info['Algorithm'] partition_name = descriptor['Partition Name'] partition_size = str(image_size) cmd = ['avbtool', 'add_hashtree_footer', '--key', key, '--algorithm', algorithm, '--partition_name', partition_name, '--partition_size', partition_size, '--do_not_generate_fec', '--image', image_path] RunCommand(args, cmd) def MakeVbmetaImage(args, key, vbmeta_img, images): info, _ = AvbInfo(args, vbmeta_img) if info: algorithm = info['Algorithm'] rollback_index = info['Rollback Index'] rollback_index_location = info['Rollback Index Location'] cmd = ['avbtool', 'make_vbmeta_image', '--key', key, '--algorithm', algorithm, '--rollback_index', rollback_index, '--rollback_index_location', rollback_index_location, '--output', vbmeta_img] for img in images: cmd.extend(['--include_descriptors_from_image', img]) RunCommand(args, cmd) # libavb expects to be able to read the maximum vbmeta size, so we must provide a partition # which matches this or the read will fail. RunCommand(args, ['truncate', '-s', '65536', vbmeta_img]) class TempDirectory(object): def __enter__(self): self.name = tempfile.mkdtemp() return self.name def __exit__(self, *unused): shutil.rmtree(self.name) def MakeSuperImage(args, partitions, output): with TempDirectory() as work_dir: cmd = ['lpmake', '--device-size=auto', '--metadata-slots=2', # A/B '--metadata-size=65536', '--sparse', '--output=' + output] for part, img in partitions.items(): tmp_img = os.path.join(work_dir, part) RunCommand(args, ['img2simg', img, tmp_img]) image_arg = '--image=%s=%s' % (part, img) partition_arg = '--partition=%s:readonly:%d:default' % ( part, os.path.getsize(img)) cmd.extend([image_arg, partition_arg]) RunCommand(args, cmd) def SignVirtApex(args): key = args.key input_dir = args.input_dir # target files in the Virt APEX bootloader = os.path.join(input_dir, 'etc', 'microdroid_bootloader') boot_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_boot-5.10.img') vendor_boot_img = os.path.join( input_dir, 'etc', 'fs', 'microdroid_vendor_boot-5.10.img') super_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_super.img') vbmeta_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_vbmeta.img') # re-sign bootloader, boot.img, vendor_boot.img AddHashFooter(args, key, bootloader) AddHashFooter(args, key, boot_img) AddHashFooter(args, key, vendor_boot_img) # re-sign super.img with TempDirectory() as work_dir: # unpack super.img tmp_super_img = os.path.join(work_dir, 'super.img') RunCommand(args, ['simg2img', super_img, tmp_super_img]) RunCommand(args, ['lpunpack', tmp_super_img, work_dir]) system_a_img = os.path.join(work_dir, 'system_a.img') vendor_a_img = os.path.join(work_dir, 'vendor_a.img') partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img} # re-sign partitions in super.img for img in partitions.values(): AddHashTreeFooter(args, key, img) # re-pack super.img MakeSuperImage(args, partitions, super_img) # re-generate vbmeta from re-signed {boot, vendor_boot, system_a, vendor_a}.img # Ideally, making VBmeta should be done out of TempDirectory block. But doing it here # to avoid unpacking re-signed super.img for system/vendor images which are available # in this block. MakeVbmetaImage(args, key, vbmeta_img, [ boot_img, vendor_boot_img, system_a_img, vendor_a_img]) def main(argv): try: args = ParseArgs(argv) SignVirtApex(args) except Exception as e: print(e) sys.exit(1) if __name__ == '__main__': main(sys.argv[1:])