Tudor Brindus | 3e03eba | 2018-07-18 11:27:13 -0700 | [diff] [blame^] | 1 | # -*- coding: utf-8 -*- |
| 2 | # Copyright 2018 The Chromium OS Authors. All rights reserved. |
| 3 | # Use of this source code is governed by a BSD-style license that can be |
| 4 | # found in the LICENSE file. |
| 5 | |
| 6 | """Script to generate a Chromium OS update for use by the update engine. |
| 7 | |
| 8 | If a source .bin is specified, the update is assumed to be a delta update. |
| 9 | """ |
| 10 | |
| 11 | from __future__ import print_function |
| 12 | |
| 13 | import os |
| 14 | import shutil |
| 15 | import tempfile |
| 16 | |
| 17 | from chromite.lib import constants |
| 18 | from chromite.lib import commandline |
| 19 | from chromite.lib import cros_build_lib |
| 20 | from chromite.lib import cros_logging as logging |
| 21 | from chromite.lib import osutils |
| 22 | |
| 23 | |
| 24 | _DELTA_GENERATOR = 'delta_generator' |
| 25 | |
| 26 | |
| 27 | # TODO(tbrindus): move this to paygen/filelib.py. |
| 28 | def CopyFileSegment(in_file, in_mode, in_len, out_file, out_mode, in_seek=0): |
| 29 | """Simulates a `dd` operation with seeks.""" |
| 30 | with open(in_file, in_mode) as in_stream, \ |
| 31 | open(out_file, out_mode) as out_stream: |
| 32 | in_stream.seek(in_seek) |
| 33 | remaining = in_len |
| 34 | while remaining: |
| 35 | chunk = in_stream.read(min(8192 * 1024, remaining)) |
| 36 | remaining -= len(chunk) |
| 37 | out_stream.write(chunk) |
| 38 | |
| 39 | |
| 40 | # TODO(tbrindus): move this to paygen/filelib.py. |
| 41 | def ExtractPartitionToTempFile(filename, partition, temp_file=None): |
| 42 | """Extracts |partition| from |filename| into |temp_file|. |
| 43 | |
| 44 | If |temp_file| is not specified, an arbitrary file is used. |
| 45 | |
| 46 | Returns the location of the extracted partition. |
| 47 | """ |
| 48 | if temp_file is None: |
| 49 | temp_file = tempfile.mktemp(prefix='cros_generate_update_payload') |
| 50 | |
| 51 | parts = cros_build_lib.GetImageDiskPartitionInfo(filename, unit='B') |
| 52 | part_info = parts[partition] |
| 53 | |
| 54 | offset = int(part_info.start) |
| 55 | length = int(part_info.size) |
| 56 | |
| 57 | CopyFileSegment(filename, 'r', length, temp_file, 'w', in_seek=offset) |
| 58 | |
| 59 | return temp_file |
| 60 | |
| 61 | |
| 62 | def PatchKernel(image, kern_file): |
| 63 | """Patches kernel |kern_file| with vblock from |image| stateful partition.""" |
| 64 | state_out = ExtractPartitionToTempFile(image, constants.PART_STATE) |
| 65 | vblock = tempfile.mktemp(prefix='vmlinuz_hd.vblock') |
| 66 | cros_build_lib.RunCommand(['e2cp', '%s:/vmlinuz_hd.vblock' % state_out, |
| 67 | vblock]) |
| 68 | |
| 69 | CopyFileSegment(vblock, 'r', os.path.getsize(vblock), kern_file, 'r+') |
| 70 | |
| 71 | osutils.SafeUnlink(state_out) |
| 72 | osutils.SafeUnlink(vblock) |
| 73 | |
| 74 | |
| 75 | def ExtractKernel(bin_file, kern_out): |
| 76 | """Extracts the kernel from the given |bin_file|, into |kern_out|.""" |
| 77 | kern_out = ExtractPartitionToTempFile(bin_file, constants.PART_KERN_B, |
| 78 | kern_out) |
| 79 | with open(kern_out, 'r') as kern: |
| 80 | if not any(kern.read(65536)): |
| 81 | logging.warn('%s: Kernel B is empty, patching kernel A.', bin_file) |
| 82 | ExtractPartitionToTempFile(bin_file, constants.PART_KERN_A, kern_out) |
| 83 | PatchKernel(bin_file, kern_out) |
| 84 | |
| 85 | return kern_out |
| 86 | |
| 87 | |
| 88 | def Ext2FileSystemSize(rootfs): |
| 89 | """Return the size in bytes of the ext2 filesystem passed in |rootfs|.""" |
| 90 | # dumpe2fs is normally installed in /sbin but doesn't require root. |
| 91 | dump = cros_build_lib.RunCommand(['/sbin/dumpe2fs', '-h', rootfs], |
| 92 | print_cmd=False, capture_output=True).output |
| 93 | fs_blocks = 0 |
| 94 | fs_blocksize = 0 |
| 95 | for line in dump.split('\n'): |
| 96 | if not line: |
| 97 | continue |
| 98 | label, data = line.split(':')[:2] |
| 99 | if label == 'Block count': |
| 100 | fs_blocks = int(data) |
| 101 | elif label == 'Block size': |
| 102 | fs_blocksize = int(data) |
| 103 | |
| 104 | return fs_blocks * fs_blocksize |
| 105 | |
| 106 | |
| 107 | def ExtractRoot(bin_file, root_out, root_pretruncate=None): |
| 108 | """Extract the rootfs partition from the gpt image |bin_file|. |
| 109 | |
| 110 | Stores it in |root_out|. If |root_out| is empty, a new temp file will be used. |
| 111 | If |root_pretruncate| is non-empty, saves the pretruncated rootfs partition |
| 112 | there. |
| 113 | """ |
| 114 | root_out = ExtractPartitionToTempFile(bin_file, constants.PART_ROOT_A, |
| 115 | root_out) |
| 116 | |
| 117 | if root_pretruncate: |
| 118 | logging.info('Saving pre-truncated root as %s.', root_pretruncate) |
| 119 | shutil.copyfile(root_out, root_pretruncate) |
| 120 | |
| 121 | # We only update the filesystem part of the partition, which is stored in the |
| 122 | # gpt script. |
| 123 | root_out_size = Ext2FileSystemSize(root_out) |
| 124 | if root_out_size: |
| 125 | with open(root_out, 'a') as root: |
| 126 | logging.info('Root size currently %d bytes.', os.path.getsize(root_out)) |
| 127 | root.truncate(root_out_size) |
| 128 | logging.info('Truncated root to %d bytes.', root_out_size) |
| 129 | else: |
| 130 | raise IOError('Error truncating the rootfs to filesystem size.') |
| 131 | |
| 132 | return root_out |
| 133 | |
| 134 | |
| 135 | def GenerateUpdatePayload(opts): |
| 136 | """Generates the output files for the given commandline |opts|.""" |
| 137 | # TODO(tbrindus): we can support calling outside of chroot easily with |
| 138 | # RunCommand, so we can remove this restriction. |
| 139 | if opts.src_image and not opts.outside_chroot: |
| 140 | # We need to be in the chroot for generating delta images. By specifying |
| 141 | # --outside_chroot you can choose not to assert this, which will allow us to |
| 142 | # run this script outside chroot. Running this script outside chroot |
| 143 | # requires copying a delta_generator binary and some related shared |
| 144 | # libraries. |
| 145 | cros_build_lib.AssertInsideChroot() |
| 146 | try: |
| 147 | src_kernel_path = src_root_path = dst_kernel_path = dst_root_path = '' |
| 148 | |
| 149 | if opts.extract: |
| 150 | if opts.src_image: |
| 151 | src_kernel_path = opts.src_kern_path or 'old_kern.dat' |
| 152 | src_root_path = opts.src_root_path or 'old_root.dat' |
| 153 | ExtractKernel(opts.src_image, src_kernel_path) |
| 154 | ExtractRoot(opts.src_image, src_root_path) |
| 155 | if opts.image: |
| 156 | dst_kernel_path = opts.kern_path or 'new_kern.dat' |
| 157 | dst_root_path = opts.root_path or 'new_root.dat' |
| 158 | ExtractKernel(opts.image, dst_kernel_path) |
| 159 | ExtractRoot(opts.image, dst_root_path, opts.root_pretruncate_path) |
| 160 | logging.info('Done extracting kernel/root') |
| 161 | return |
| 162 | |
| 163 | delta, payload_type = (True, 'delta') if opts.src_image else (False, 'full') |
| 164 | |
| 165 | logging.info('Generating %s update', payload_type) |
| 166 | |
| 167 | if delta: |
| 168 | if not opts.full_kernel: |
| 169 | src_kernel_path = ExtractKernel(opts.src_image, opts.src_kern_path) |
| 170 | else: |
| 171 | logging.info('Generating full kernel update') |
| 172 | |
| 173 | src_root_path = ExtractRoot(opts.src_image, opts.src_root_path) |
| 174 | |
| 175 | dst_kernel_path = ExtractKernel(opts.image, opts.kern_path) |
| 176 | dst_root_path = ExtractRoot(opts.image, opts.root_path, |
| 177 | opts.root_pretruncate_path) |
| 178 | |
| 179 | # TODO(tbrindus): delta_generator should be called with partition lists |
| 180 | # to support major version 2 more easily. |
| 181 | generator_args = [ |
| 182 | # Common payload args: |
| 183 | '--major_version=1', |
| 184 | '--out_file=' + opts.output, |
| 185 | '--private_key=' + (opts.private_key or ''), |
| 186 | '--out_metadata_size_file=' + (opts.out_metadata_size_file or ''), |
| 187 | # Target image args: |
| 188 | '--new_image=' + dst_root_path, |
| 189 | '--new_kernel=' + dst_kernel_path, |
| 190 | '--new_channel=' + opts.channel, |
| 191 | '--new_board=' + opts.board, |
| 192 | '--new_version=' + opts.version, |
| 193 | '--new_key=' + opts.key, |
| 194 | '--new_build_channel=' + opts.build_channel, |
| 195 | '--new_build_version=' + opts.build_version, |
| 196 | ] |
| 197 | |
| 198 | if delta: |
| 199 | generator_args += [ |
| 200 | # Source image args: |
| 201 | '--old_image=' + src_root_path, |
| 202 | '--old_kernel=' + src_kernel_path, |
| 203 | '--old_channel=' + opts.src_channel, |
| 204 | '--old_board=' + opts.src_board, |
| 205 | '--old_version=' + opts.src_version, |
| 206 | '--old_key=' + opts.src_key, |
| 207 | '--old_build_channel=' + opts.src_build_channel, |
| 208 | '--old_build_version=' + opts.src_build_version, |
| 209 | ] |
| 210 | |
| 211 | # The passed chunk_size is only used for delta payload. Use |
| 212 | # delta_generator's default if no value is provided. |
| 213 | if opts.chunk_size: |
| 214 | logging.info('Forcing chunk size to %d', opts.chunk_size) |
| 215 | generator_args.append('--chunk_size=%d' % opts.chunk_size) |
| 216 | |
| 217 | # Add partition size. Only *required* for minor_version=1. |
| 218 | # TODO(tbrindus): deprecate this when we deprecate minor version 1. |
| 219 | dst_root_part_size = cros_build_lib.GetImageDiskPartitionInfo( |
| 220 | opts.image, unit='B')[constants.PART_ROOT_A].size |
| 221 | |
| 222 | if dst_root_part_size: |
| 223 | logging.info('Using rootfs partition size: %d', dst_root_part_size) |
| 224 | generator_args.append('--rootfs_partition_size=%d' % dst_root_part_size) |
| 225 | else: |
| 226 | logging.info('Using the default partition size') |
| 227 | |
| 228 | cros_build_lib.RunCommand([_DELTA_GENERATOR] + generator_args) |
| 229 | |
| 230 | if opts.out_payload_hash_file or opts.out_metadata_hash_file: |
| 231 | # The out_metadata_hash_file flag requires out_hash_file flag to be set |
| 232 | # in delta_generator, if caller doesn't provide it, we set it to |
| 233 | # /dev/null. |
| 234 | out_payload_hash_file = opts.out_payload_hash_file or '/dev/null' |
| 235 | |
| 236 | # The manifest - unfortunately - contains two fields called |
| 237 | # signature_offset and signature_size with data about how the manifest is |
| 238 | # signed. This means we have to pass the signature size used. The value |
| 239 | # 256 is the number of bytes the SHA-256 hash value of the manifest signed |
| 240 | # with a 2048-bit RSA key occupies. |
| 241 | generator_args = [ |
| 242 | '--in_file=' + opts.output, |
| 243 | '--signature_size=256', |
| 244 | '--out_hash_file=' + out_payload_hash_file, |
| 245 | '--out_metadata_hash_file=' + (opts.out_metadata_hash_file or ''), |
| 246 | ] |
| 247 | cros_build_lib.RunCommand([_DELTA_GENERATOR] + generator_args) |
| 248 | |
| 249 | logging.info('Done generating %s update', payload_type) |
| 250 | finally: |
| 251 | if not opts.src_kern_path: |
| 252 | osutils.SafeUnlink(src_kernel_path) |
| 253 | if not opts.src_root_path: |
| 254 | osutils.SafeUnlink(src_root_path) |
| 255 | if not opts.kern_path: |
| 256 | osutils.SafeUnlink(dst_kernel_path) |
| 257 | if not opts.root_path: |
| 258 | osutils.SafeUnlink(dst_root_path) |
| 259 | |
| 260 | |
| 261 | def ParseArguments(argv): |
| 262 | """Returns a namespace for the CLI arguments.""" |
| 263 | parser = commandline.ArgumentParser(description=__doc__) |
| 264 | parser.add_argument('--image', type='path', |
| 265 | help='The image that should be sent to clients.') |
| 266 | parser.add_argument('--src_image', type='path', |
| 267 | help='A source image. If specified, this makes a delta ' |
| 268 | 'update.') |
| 269 | parser.add_argument('--output', type='path', help='Output file.') |
| 270 | parser.add_argument('--outside_chroot', action='store_true', |
| 271 | help='Running outside of chroot.') |
| 272 | parser.add_argument('--private_key', type='path', |
| 273 | help='Path to private key in .pem format.') |
| 274 | parser.add_argument('--out_payload_hash_file', type='path', |
| 275 | help='Path to output payload hash file.') |
| 276 | parser.add_argument('--out_metadata_hash_file', type='path', |
| 277 | help='Path to output metadata hash file.') |
| 278 | parser.add_argument('--out_metadata_size_file', type='path', |
| 279 | help='Path to output metadata size file.') |
| 280 | parser.add_argument('--extract', action='store_true', |
| 281 | help='If set, extract old/new kernel/rootfs to ' |
| 282 | '[old|new]_[kern|root].dat. Useful for debugging.') |
| 283 | parser.add_argument('--full_kernel', action='store_true', |
| 284 | help='Generate a full kernel update even if generating a ' |
| 285 | 'delta update.') |
| 286 | parser.add_argument('--chunk_size', type=int, |
| 287 | help='Delta payload chunk size (-1 means whole files).') |
| 288 | |
| 289 | # TODO(tbrindus): Remove --build_version and --src_build_version flags. |
| 290 | src_group = parser.add_argument_group('Source options') |
| 291 | src_group.add_argument('--src_channel', type=str, default='', |
| 292 | help='Channel of the src image.') |
| 293 | src_group.add_argument('--src_board', type=str, default='', |
| 294 | help='Board of the src image.') |
| 295 | src_group.add_argument('--src_version', type=str, default='', |
| 296 | help='Version of the src image.') |
| 297 | src_group.add_argument('--src_key', type=str, default='', |
| 298 | help='Key of the src image.') |
| 299 | src_group.add_argument('--src_build_channel', type=str, default='', |
| 300 | help='Channel of the build of the src image.') |
| 301 | src_group.add_argument('--src_build_version', type=str, default='', |
| 302 | help='Channel of the build of the the src image.') |
| 303 | |
| 304 | dst_group = parser.add_argument_group('Target options') |
| 305 | dst_group.add_argument('--channel', type=str, default='', |
| 306 | help='Channel of the target image.') |
| 307 | dst_group.add_argument('--board', type=str, default='', |
| 308 | help='Board of the target image.') |
| 309 | dst_group.add_argument('--version', type=str, default='', |
| 310 | help='Version of the target image.') |
| 311 | dst_group.add_argument('--key', type=str, default='', |
| 312 | help='Key of the target image.') |
| 313 | dst_group.add_argument('--build_channel', type=str, default='', |
| 314 | help='Channel of the build of the target image.') |
| 315 | dst_group.add_argument('--build_version', type=str, default='', |
| 316 | help='Channel of the build of the the target image.') |
| 317 | |
| 318 | # Because we archive/call old versions of this script, we can't easily remove |
| 319 | # command line options, even if we ignore this one now. |
| 320 | parser.add_argument('--patch_kernel', action='store_true', |
| 321 | help='Ignored. Present for compatibility.') |
| 322 | |
| 323 | # Specifying any of the following will cause it to not be cleaned up on exit. |
| 324 | parser.add_argument('--kern_path', type='path', |
| 325 | help='File path for extracting the kernel partition.') |
| 326 | parser.add_argument('--root_path', type='path', |
| 327 | help='File path for extracting the rootfs partition.') |
| 328 | parser.add_argument('--root_pretruncate_path', type='path', |
| 329 | help='File path for extracting the rootfs partition, ' |
| 330 | 'pre-truncation.') |
| 331 | parser.add_argument('--src_kern_path', type='path', |
| 332 | help='File path for extracting the source kernel ' |
| 333 | 'partition.') |
| 334 | parser.add_argument('--src_root_path', type='path', |
| 335 | help='File path for extracting the source root ' |
| 336 | 'partition.') |
| 337 | |
| 338 | opts = parser.parse_args(argv) |
| 339 | opts.Freeze() |
| 340 | |
| 341 | if not opts.extract and not opts.output: |
| 342 | parser.error('You must specify an output filename with --output FILENAME') |
| 343 | |
| 344 | return opts |
| 345 | |
| 346 | |
| 347 | def main(argv): |
| 348 | opts = ParseArguments(argv) |
| 349 | |
| 350 | return GenerateUpdatePayload(opts) |